projector.ts 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400
  1. export * as SessionProjector from "./projector"
  2. import { and, desc, eq, sql } from "drizzle-orm"
  3. import { DateTime, Effect, Layer, Schema } from "effect"
  4. import { Database } from "../database/database"
  5. import { EventV2 } from "../event"
  6. import { LayerNode } from "../effect/layer-node"
  7. import { SessionEvent } from "./event"
  8. import { SessionV1 } from "../v1/session"
  9. import { WorkspaceTable } from "../control-plane/workspace.sql"
  10. import { SessionMessage } from "./message"
  11. import { SessionMessageUpdater } from "./message-updater"
  12. import { SessionInput } from "./input"
  13. import { WorkspaceV2 } from "../workspace"
  14. import { SessionContextEpoch } from "./context-epoch"
  15. import { MessageTable, PartTable, SessionMessageTable, SessionTable } from "./sql"
  16. import type { DeepMutable } from "../schema"
  17. type DatabaseService = Database.Interface["db"]
  18. const decodeMessage = Schema.decodeUnknownSync(SessionMessage.Message)
  19. const encodeMessage = Schema.encodeSync(SessionMessage.Message)
  20. export class SessionAlreadyProjected extends Error {}
  21. type Usage = {
  22. cost: number
  23. tokens: {
  24. input: number
  25. output: number
  26. reasoning: number
  27. cache: { read: number; write: number }
  28. }
  29. }
  30. function usage(part: (typeof SessionV1.Event.PartUpdated.Type)["data"]["part"] | unknown): Usage | undefined {
  31. if (typeof part !== "object" || part === null) return undefined
  32. const value = part as Record<string, unknown>
  33. if (value.type !== "step-finish") return undefined
  34. if (!("cost" in value) || !("tokens" in value)) return undefined
  35. return { cost: value.cost as Usage["cost"], tokens: value.tokens as Usage["tokens"] }
  36. }
  37. function sessionRow(info: SessionV1.SessionInfo): typeof SessionTable.$inferInsert {
  38. return {
  39. id: info.id,
  40. project_id: info.projectID,
  41. workspace_id: info.workspaceID ?? null,
  42. parent_id: info.parentID,
  43. slug: info.slug,
  44. directory: info.directory,
  45. path: info.path,
  46. title: info.title,
  47. agent: info.agent,
  48. model: info.model,
  49. version: info.version,
  50. share_url: info.share?.url,
  51. summary_additions: info.summary?.additions,
  52. summary_deletions: info.summary?.deletions,
  53. summary_files: info.summary?.files,
  54. summary_diffs: info.summary?.diffs ? [...info.summary.diffs] : undefined,
  55. metadata: info.metadata,
  56. cost: info.cost ?? 0,
  57. tokens_input: (info.tokens ?? { input: 0 }).input,
  58. tokens_output: (info.tokens ?? { output: 0 }).output,
  59. tokens_reasoning: (info.tokens ?? { reasoning: 0 }).reasoning,
  60. tokens_cache_read: (info.tokens ?? { cache: { read: 0 } }).cache.read,
  61. tokens_cache_write: (info.tokens ?? { cache: { write: 0 } }).cache.write,
  62. revert: info.revert ?? null,
  63. permission: info.permission ? [...info.permission] : undefined,
  64. time_created: info.time.created,
  65. time_updated: info.time.updated,
  66. time_compacting: info.time.compacting,
  67. time_archived: info.time.archived,
  68. }
  69. }
  70. function messageData(
  71. info: (typeof SessionV1.Event.MessageUpdated.Type)["data"]["info"],
  72. ): typeof MessageTable.$inferInsert.data {
  73. const { id: _, sessionID: __, ...rest } = info
  74. return rest as DeepMutable<typeof rest>
  75. }
  76. function partData(part: (typeof SessionV1.Event.PartUpdated.Type)["data"]["part"]): typeof PartTable.$inferInsert.data {
  77. const { id: _, messageID: __, sessionID: ___, ...rest } = part
  78. return rest as DeepMutable<typeof rest>
  79. }
  80. function applyUsage(
  81. db: DatabaseService,
  82. sessionID: (typeof SessionV1.Event.MessageUpdated.Type)["data"]["sessionID"],
  83. value: Usage,
  84. sign = 1,
  85. ) {
  86. return db
  87. .update(SessionTable)
  88. .set({
  89. cost: sql`${SessionTable.cost} + ${value.cost * sign}`,
  90. tokens_input: sql`${SessionTable.tokens_input} + ${value.tokens.input * sign}`,
  91. tokens_output: sql`${SessionTable.tokens_output} + ${value.tokens.output * sign}`,
  92. tokens_reasoning: sql`${SessionTable.tokens_reasoning} + ${value.tokens.reasoning * sign}`,
  93. tokens_cache_read: sql`${SessionTable.tokens_cache_read} + ${value.tokens.cache.read * sign}`,
  94. tokens_cache_write: sql`${SessionTable.tokens_cache_write} + ${value.tokens.cache.write * sign}`,
  95. time_updated: sql`${SessionTable.time_updated}`,
  96. })
  97. .where(eq(SessionTable.id, sessionID))
  98. .run()
  99. .pipe(Effect.orDie)
  100. }
  101. function run(db: DatabaseService, event: SessionEvent.Event) {
  102. return Effect.gen(function* () {
  103. const decodeRow = (row: typeof SessionMessageTable.$inferSelect) =>
  104. decodeMessage({ ...row.data, id: row.id, type: row.type })
  105. const updateMessage = (message: SessionMessage.Message) => {
  106. if (event.durable === undefined) return Effect.die("Durable Session event is missing aggregate sequence")
  107. const encoded = encodeMessage(message)
  108. const { id, type, ...data } = encoded
  109. return db
  110. .update(SessionMessageTable)
  111. .set({ type, time_created: DateTime.toEpochMillis(message.time.created), data })
  112. .where(
  113. and(
  114. eq(SessionMessageTable.id, SessionMessage.ID.make(id)),
  115. eq(SessionMessageTable.session_id, event.data.sessionID),
  116. ),
  117. )
  118. .run()
  119. .pipe(Effect.orDie)
  120. }
  121. const appendMessage = (message: SessionMessage.Message) => insertMessage(db, event, message)
  122. const adapter: SessionMessageUpdater.Adapter = {
  123. getCurrentAssistant() {
  124. return Effect.gen(function* () {
  125. // A newer turn supersedes stale incomplete rows; never resume an older assistant projection.
  126. const row = yield* db
  127. .select()
  128. .from(SessionMessageTable)
  129. .where(
  130. and(eq(SessionMessageTable.session_id, event.data.sessionID), eq(SessionMessageTable.type, "assistant")),
  131. )
  132. .orderBy(desc(SessionMessageTable.seq))
  133. .limit(1)
  134. .get()
  135. .pipe(Effect.orDie)
  136. if (!row) return
  137. const message = decodeRow(row)
  138. return message.type === "assistant" && !message.time.completed ? message : undefined
  139. })
  140. },
  141. getAssistant(messageID) {
  142. return Effect.gen(function* () {
  143. const row = yield* db
  144. .select()
  145. .from(SessionMessageTable)
  146. .where(
  147. and(
  148. eq(SessionMessageTable.id, messageID),
  149. eq(SessionMessageTable.session_id, event.data.sessionID),
  150. eq(SessionMessageTable.type, "assistant"),
  151. ),
  152. )
  153. .get()
  154. .pipe(Effect.orDie)
  155. if (!row) return
  156. const message = decodeRow(row)
  157. return message.type === "assistant" ? message : undefined
  158. })
  159. },
  160. getCurrentShell(callID) {
  161. return Effect.gen(function* () {
  162. const rows = yield* db
  163. .select()
  164. .from(SessionMessageTable)
  165. .where(and(eq(SessionMessageTable.session_id, event.data.sessionID), eq(SessionMessageTable.type, "shell")))
  166. .orderBy(desc(SessionMessageTable.seq))
  167. .all()
  168. .pipe(Effect.orDie)
  169. return rows
  170. .map(decodeRow)
  171. .find((message): message is SessionMessage.Shell => message.type === "shell" && message.callID === callID)
  172. })
  173. },
  174. updateAssistant: updateMessage,
  175. updateShell: updateMessage,
  176. appendMessage,
  177. }
  178. yield* SessionMessageUpdater.update(adapter, event)
  179. })
  180. }
  181. function insertMessage(db: DatabaseService, event: SessionEvent.Event, message: SessionMessage.Message) {
  182. if (event.durable === undefined) return Effect.die("Durable Session event is missing aggregate sequence")
  183. const encoded = encodeMessage(message)
  184. const { id, type, ...data } = encoded
  185. return db
  186. .insert(SessionMessageTable)
  187. .values({
  188. id: SessionMessage.ID.make(id),
  189. session_id: event.data.sessionID,
  190. type,
  191. seq: event.durable.seq,
  192. time_created: DateTime.toEpochMillis(message.time.created),
  193. data,
  194. })
  195. .run()
  196. .pipe(Effect.orDie)
  197. }
  198. export const layer = Layer.effectDiscard(
  199. Effect.gen(function* () {
  200. const events = yield* EventV2.Service
  201. const { db } = yield* Database.Service
  202. yield* events.project(SessionV1.Event.Created, (event) =>
  203. Effect.gen(function* () {
  204. const stored = yield* db
  205. .insert(SessionTable)
  206. .values(sessionRow(event.data.info))
  207. .onConflictDoNothing()
  208. .returning({ sessionID: SessionTable.id })
  209. .get()
  210. .pipe(Effect.orDie)
  211. if (!stored) return yield* Effect.die(new SessionAlreadyProjected())
  212. if (event.data.info.workspaceID) {
  213. yield* db
  214. .update(WorkspaceTable)
  215. .set({ time_used: Date.now() })
  216. .where(eq(WorkspaceTable.id, event.data.info.workspaceID))
  217. .run()
  218. .pipe(Effect.orDie)
  219. }
  220. }),
  221. )
  222. yield* events.project(SessionV1.Event.Updated, (event) =>
  223. db
  224. .update(SessionTable)
  225. .set(sessionRow(event.data.info))
  226. .where(eq(SessionTable.id, event.data.sessionID))
  227. .run()
  228. .pipe(Effect.orDie),
  229. )
  230. yield* events.project(SessionEvent.Moved, (event) =>
  231. Effect.gen(function* () {
  232. yield* db
  233. .update(SessionTable)
  234. .set({
  235. directory: event.data.location.directory,
  236. path: event.data.subdirectory,
  237. workspace_id: event.data.location.workspaceID ? WorkspaceV2.ID.make(event.data.location.workspaceID) : null,
  238. time_updated: DateTime.toEpochMillis(event.data.timestamp),
  239. })
  240. .where(eq(SessionTable.id, event.data.sessionID))
  241. .run()
  242. .pipe(Effect.orDie)
  243. yield* SessionContextEpoch.reset(db, event.data.sessionID)
  244. }),
  245. )
  246. yield* events.project(SessionV1.Event.Deleted, (event) =>
  247. db.delete(SessionTable).where(eq(SessionTable.id, event.data.sessionID)).run().pipe(Effect.orDie),
  248. )
  249. yield* events.project(SessionV1.Event.MessageUpdated, (event) =>
  250. Effect.gen(function* () {
  251. const time_created = event.data.info.time.created
  252. const id = event.data.info.id
  253. const sessionID = event.data.info.sessionID
  254. const data = messageData(event.data.info)
  255. yield* db
  256. .insert(MessageTable)
  257. .values({ id, session_id: sessionID, time_created, data })
  258. .onConflictDoUpdate({ target: MessageTable.id, set: { data } })
  259. .run()
  260. .pipe(Effect.orDie)
  261. }),
  262. )
  263. yield* events.project(SessionV1.Event.MessageRemoved, (event) =>
  264. Effect.gen(function* () {
  265. const rows = yield* db
  266. .select()
  267. .from(PartTable)
  268. .where(and(eq(PartTable.message_id, event.data.messageID), eq(PartTable.session_id, event.data.sessionID)))
  269. .all()
  270. .pipe(Effect.orDie)
  271. for (const row of rows) {
  272. const previous = usage(row.data)
  273. if (previous) yield* applyUsage(db, event.data.sessionID, previous, -1)
  274. }
  275. yield* db
  276. .delete(MessageTable)
  277. .where(and(eq(MessageTable.id, event.data.messageID), eq(MessageTable.session_id, event.data.sessionID)))
  278. .run()
  279. .pipe(Effect.orDie)
  280. }),
  281. )
  282. yield* events.project(SessionV1.Event.PartRemoved, (event) =>
  283. Effect.gen(function* () {
  284. const row = yield* db
  285. .select()
  286. .from(PartTable)
  287. .where(and(eq(PartTable.id, event.data.partID), eq(PartTable.session_id, event.data.sessionID)))
  288. .get()
  289. .pipe(Effect.orDie)
  290. const previous = row && usage(row.data)
  291. if (previous) yield* applyUsage(db, event.data.sessionID, previous, -1)
  292. yield* db
  293. .delete(PartTable)
  294. .where(and(eq(PartTable.id, event.data.partID), eq(PartTable.session_id, event.data.sessionID)))
  295. .run()
  296. .pipe(Effect.orDie)
  297. }),
  298. )
  299. yield* events.project(SessionV1.Event.PartUpdated, (event) =>
  300. Effect.gen(function* () {
  301. const id = event.data.part.id
  302. const messageID = event.data.part.messageID
  303. const sessionID = event.data.part.sessionID
  304. const data = partData(event.data.part)
  305. const row = yield* db.select().from(PartTable).where(eq(PartTable.id, id)).get().pipe(Effect.orDie)
  306. yield* db
  307. .insert(PartTable)
  308. .values({ id, message_id: messageID, session_id: sessionID, time_created: event.data.time, data })
  309. .onConflictDoUpdate({ target: PartTable.id, set: { data } })
  310. .run()
  311. .pipe(Effect.orDie)
  312. const previous = row && usage(row.data)
  313. const next = usage(event.data.part)
  314. if (previous) yield* applyUsage(db, row.session_id, previous, -1)
  315. if (next) yield* applyUsage(db, sessionID, next)
  316. }),
  317. )
  318. yield* events.project(SessionEvent.AgentSwitched, (event) =>
  319. db
  320. .update(SessionTable)
  321. .set({ agent: event.data.agent, time_updated: DateTime.toEpochMillis(event.data.timestamp) })
  322. .where(eq(SessionTable.id, event.data.sessionID))
  323. .run()
  324. .pipe(Effect.orDie, Effect.andThen(run(db, event))),
  325. )
  326. yield* events.project(SessionEvent.ModelSwitched, (event) =>
  327. Effect.gen(function* () {
  328. yield* db
  329. .update(SessionTable)
  330. .set({ model: event.data.model, time_updated: DateTime.toEpochMillis(event.data.timestamp) })
  331. .where(eq(SessionTable.id, event.data.sessionID))
  332. .run()
  333. .pipe(Effect.orDie)
  334. yield* run(db, event)
  335. }),
  336. )
  337. yield* events.project(SessionEvent.Prompted, (event) =>
  338. Effect.gen(function* () {
  339. if (event.durable === undefined) return yield* Effect.die("Durable Session event is missing aggregate sequence")
  340. yield* SessionInput.projectPrompted(db, {
  341. id: event.data.messageID,
  342. sessionID: event.data.sessionID,
  343. prompt: event.data.prompt,
  344. delivery: event.data.delivery,
  345. timeCreated: event.data.timestamp,
  346. promotedSeq: event.durable.seq,
  347. })
  348. yield* run(db, event)
  349. }),
  350. )
  351. yield* events.project(SessionEvent.PromptAdmitted, (event) =>
  352. Effect.gen(function* () {
  353. if (event.durable === undefined) return yield* Effect.die("Durable Session event is missing aggregate sequence")
  354. yield* SessionInput.projectAdmitted(db, {
  355. admittedSeq: event.durable.seq,
  356. id: event.data.messageID,
  357. sessionID: event.data.sessionID,
  358. prompt: event.data.prompt,
  359. delivery: event.data.delivery,
  360. timeCreated: event.data.timestamp,
  361. })
  362. }),
  363. )
  364. yield* events.project(SessionEvent.ContextUpdated, (event) => run(db, event))
  365. yield* events.project(SessionEvent.Synthetic, (event) => run(db, event))
  366. yield* events.project(SessionEvent.Shell.Started, (event) => run(db, event))
  367. yield* events.project(SessionEvent.Shell.Ended, (event) => run(db, event))
  368. yield* events.project(SessionEvent.Step.Started, (event) => run(db, event))
  369. yield* events.project(SessionEvent.Step.Ended, (event) => run(db, event))
  370. yield* events.project(SessionEvent.Step.Failed, (event) => run(db, event))
  371. yield* events.project(SessionEvent.Text.Started, (event) => run(db, event))
  372. yield* events.project(SessionEvent.Text.Ended, (event) => run(db, event))
  373. yield* events.project(SessionEvent.Tool.Input.Started, (event) => run(db, event))
  374. yield* events.project(SessionEvent.Tool.Input.Ended, (event) => run(db, event))
  375. yield* events.project(SessionEvent.Tool.Called, (event) => run(db, event))
  376. yield* events.project(SessionEvent.Tool.Progress, (event) => run(db, event))
  377. yield* events.project(SessionEvent.Tool.Success, (event) => run(db, event))
  378. yield* events.project(SessionEvent.Tool.Failed, (event) => run(db, event))
  379. yield* events.project(SessionEvent.Reasoning.Started, (event) => run(db, event))
  380. yield* events.project(SessionEvent.Reasoning.Ended, (event) => run(db, event))
  381. // yield* events.project(SessionEvent.Retried, (event) => run(db, event))
  382. yield* events.project(SessionEvent.Compaction.Ended, (event) => run(db, event))
  383. }),
  384. )
  385. export const defaultLayer = layer.pipe(Layer.provide(EventV2.defaultLayer), Layer.provide(Database.defaultLayer))
  386. export const node = LayerNode.make(layer, [EventV2.node, Database.node])