feat(core): persist v2 session context epochs (#30789)

packages/core/src/session/context.ts

@@ -1,47 +1,92 @@
-import { and, asc, desc, eq, gt, gte, or } from "drizzle-orm"
+import { and, asc, desc, eq, gt, gte, ne, or } from "drizzle-orm"
 import { Effect, Schema } from "effect"
 import { Database } from "../database/database"
 import { MessageDecodeError } from "./error"
 import { SessionMessage } from "./message"
 import { SessionSchema } from "./schema"
-import { SessionMessageTable } from "./sql"
+import { SessionContextEpochTable, SessionMessageTable } from "./sql"
 
 type DatabaseService = Database.Interface["db"]
 
 const decode = Schema.decodeUnknownEffect(SessionMessage.Message)
 
-export const load = Effect.fn("SessionContext.load")(function* (db: DatabaseService, sessionID: SessionSchema.ID) {
-  const compaction = yield* db
-    .select()
+const latestCompaction = Effect.fnUntraced(function* (db: DatabaseService, sessionID: SessionSchema.ID) {
+  return yield* db
+    .select({ seq: SessionMessageTable.seq })
     .from(SessionMessageTable)
     .where(and(eq(SessionMessageTable.session_id, sessionID), eq(SessionMessageTable.type, "compaction")))
     .orderBy(desc(SessionMessageTable.seq))
     .limit(1)
     .get()
     .pipe(Effect.orDie)
-  const rows = yield* db
+})
+
+const messageRows = Effect.fnUntraced(function* (
+  db: DatabaseService,
+  sessionID: SessionSchema.ID,
+  compaction: { readonly seq: number } | undefined,
+  baselineSeq?: number,
+) {
+  return yield* db
     .select()
     .from(SessionMessageTable)
     .where(
       and(
         eq(SessionMessageTable.session_id, sessionID),
-        compaction ? or(gte(SessionMessageTable.seq, compaction.seq)) : undefined,
+        compaction
+          ? or(
+              gte(SessionMessageTable.seq, compaction.seq),
+              baselineSeq === undefined
+                ? undefined
+                : and(eq(SessionMessageTable.type, "system"), gt(SessionMessageTable.seq, baselineSeq)),
+            )
+          : undefined,
+        baselineSeq === undefined
+          ? undefined
+          : or(ne(SessionMessageTable.type, "system"), gt(SessionMessageTable.seq, baselineSeq)),
       ),
     )
     .orderBy(asc(SessionMessageTable.seq))
     .all()
     .pipe(Effect.orDie)
-  return yield* Effect.forEach(rows, (row) =>
-    decode({ ...row.data, id: row.id, type: row.type }).pipe(
-      Effect.mapError(
-        () =>
-          new MessageDecodeError({
-            sessionID: SessionSchema.ID.make(row.session_id),
-            messageID: SessionMessage.ID.make(row.id),
-          }),
-      ),
+})
+
+const decodeMessageRow = (row: typeof SessionMessageTable.$inferSelect) =>
+  decode({ ...row.data, id: row.id, type: row.type }).pipe(
+    Effect.mapError(
+      () =>
+        new MessageDecodeError({
+          sessionID: SessionSchema.ID.make(row.session_id),
+          messageID: SessionMessage.ID.make(row.id),
+        }),
     ),
   )
+
+export const load = Effect.fn("SessionContext.load")(function* (db: DatabaseService, sessionID: SessionSchema.ID) {
+  const [epoch, compaction] = yield* Effect.all(
+    [
+      db
+        .select({ baselineSeq: SessionContextEpochTable.baseline_seq })
+        .from(SessionContextEpochTable)
+        .where(eq(SessionContextEpochTable.session_id, sessionID))
+        .get()
+        .pipe(Effect.orDie),
+      latestCompaction(db, sessionID),
+    ],
+    { concurrency: "unbounded" },
+  )
+  return yield* Effect.forEach(yield* messageRows(db, sessionID, compaction, epoch?.baselineSeq), decodeMessageRow)
+})
+
+export const loadForRunner = Effect.fn("SessionContext.loadForRunner")(function* (
+  db: DatabaseService,
+  sessionID: SessionSchema.ID,
+  baselineSeq: number,
+) {
+  return yield* Effect.forEach(
+    yield* messageRows(db, sessionID, yield* latestCompaction(db, sessionID), baselineSeq),
+    decodeMessageRow,
+  )
 })
 
 export * as SessionContext from "./context"