diff --git a/packages/session-backends/sqlite-node/src/sqlite/session/branch-entries.ts b/packages/session-backends/sqlite-node/src/sqlite/session/branch-entries.ts index d0b1ac063..ea89192e0 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/session/branch-entries.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/session/branch-entries.ts @@ -35,6 +35,19 @@ interface EntryStructureRow { timestamp: number; } +interface BranchTipRow { + branch_id: string; +} + +interface CompactionBoundary { + branchId: string; + seq: number; +} + +interface CompactionBoundaryRow { + entry_seq: number | null; +} + function readBranchMembership(db: SqliteDatabase, entryId: string): BranchMembershipRow { const row = sql`SELECT b.branch_id, b.entry_seq FROM branch_entries b @@ -56,6 +69,31 @@ function readBranchMeta(db: SqliteDatabase, branchId: string): BranchMetaRow { return row; } +function insertBranchEntry(db: SqliteDatabase, branchId: string, entry: Entry): void { + sql`INSERT INTO branch_entries (branch_id, entry_id, entry_seq, entry_type) + VALUES (${branchId}, ${entry.id}, ${entry.seq}, ${entry.type})`.run(db); +} + +function readBranchTipForParent(db: SqliteDatabase, parentId: string): BranchTipRow | undefined { + return sql`SELECT branch_id + FROM branch_meta + WHERE tip_entry_id = ${parentId}`.get(db); +} + +function createRootBranchForEntry(db: SqliteDatabase, entry: Entry): void { + sql`INSERT INTO branch_meta (branch_id, tip_entry_id, tip_seq, base_branch_id, base_seq) + VALUES (${entry.id}, ${entry.id}, ${entry.seq}, ${null}, ${null})`.run(db); + insertBranchEntry(db, entry.id, entry); +} + +function appendEntryToExistingBranch(db: SqliteDatabase, branchId: string, entry: Entry): void { + insertBranchEntry(db, branchId, entry); + const result = sql`UPDATE branch_meta + SET tip_entry_id = ${entry.id}, tip_seq = ${entry.seq} + WHERE branch_id = ${branchId}`.run(db); + if (result.changes !== 1) throw new Error(`Expected to update branch ${branchId}, updated ${result.changes}`); +} + function readBranchSegmentsNewestFirst(db: SqliteDatabase, start: string): BranchSegment[] { let { branch_id: branchId, entry_seq: upperSeq } = readBranchMembership(db, start); const segments: BranchSegment[] = []; @@ -71,6 +109,69 @@ function readBranchSegmentsNewestFirst(db: SqliteDatabase, start: string): Branc return segments; } +function readNewestCompactionBoundary( + db: SqliteDatabase, + segmentsNewestFirst: readonly BranchSegment[], +): CompactionBoundary | undefined { + for (const segment of segmentsNewestFirst) { + const row = sql`SELECT MAX(entry_seq) AS entry_seq + FROM branch_entries + WHERE branch_id = ${segment.branchId} + AND entry_seq > ${segment.lowerSeq} + AND entry_seq <= ${segment.upperSeq} + AND entry_type = ${"compaction"}`.get(db); + if (row?.entry_seq !== null && row?.entry_seq !== undefined) + return { branchId: segment.branchId, seq: row.entry_seq }; + } + return undefined; +} + +function copyBranchEntriesAfterSeqThroughParent( + db: SqliteDatabase, + targetBranchId: string, + segmentsNewestFirst: readonly BranchSegment[], + afterSeq: number, +): void { + for (const segment of [...segmentsNewestFirst].reverse()) { + const lowerSeq = Math.max(segment.lowerSeq, afterSeq); + if (segment.upperSeq <= lowerSeq) continue; + sql`INSERT INTO branch_entries (branch_id, entry_id, entry_seq, entry_type) + SELECT ${targetBranchId}, entry_id, entry_seq, entry_type + FROM branch_entries + WHERE branch_id = ${segment.branchId} + AND entry_seq > ${lowerSeq} + AND entry_seq <= ${segment.upperSeq}`.run(db); + } +} + +function createDivergentBranchForEntry(db: SqliteDatabase, entry: Entry): void { + if (entry.parentId === null) throw new Error("Root entries do not create divergent branches"); + const segmentsNewestFirst = readBranchSegmentsNewestFirst(db, entry.parentId); + const compaction = readNewestCompactionBoundary(db, segmentsNewestFirst); + const branchId = entry.id; + // A null base means this segment stores its own root-through-parent prefix. + sql`INSERT INTO branch_meta (branch_id, tip_entry_id, tip_seq, base_branch_id, base_seq) + VALUES (${branchId}, ${entry.id}, ${entry.seq}, ${compaction?.branchId ?? null}, ${compaction?.seq ?? null})`.run( + db, + ); + copyBranchEntriesAfterSeqThroughParent(db, branchId, segmentsNewestFirst, compaction?.seq ?? 0); + insertBranchEntry(db, branchId, entry); +} + +export function appendEntryToBranchIndex(db: SqliteDatabase, entry: Entry): void { + if (entry.parentId === null) { + createRootBranchForEntry(db, entry); + return; + } + + const branch = readBranchTipForParent(db, entry.parentId); + if (branch === undefined) { + createDivergentBranchForEntry(db, entry); + return; + } + appendEntryToExistingBranch(db, branch.branch_id, entry); +} + function stopPredicates(query: StorageBranchScan): SqlQuery[] { const predicates: SqlQuery[] = []; if (query.stopAtType !== undefined) predicates.push(sql`b.entry_type = ${query.stopAtType}`); diff --git a/packages/session-backends/sqlite-node/src/sqlite/storage.ts b/packages/session-backends/sqlite-node/src/sqlite/storage.ts index f7a3c1e33..b53438454 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/storage.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/storage.ts @@ -13,7 +13,7 @@ import type { UsageScan, } from "@earendil-works/pi-agent-core"; import { prepareStorageCommit } from "@earendil-works/pi-agent-core"; -import { scanBranchEntries, scanBranchEntryStructures } from "./session/branch-entries.ts"; +import { appendEntryToBranchIndex, scanBranchEntries, scanBranchEntryStructures } from "./session/branch-entries.ts"; import { decodeEntryRow, insertEntryRow, readEntryRows, scanEntryRows } from "./session/entries.ts"; import { deleteRegisterRow, listRegisterRows, readRegisterRow, setRegisterRow } from "./session/registers.ts"; import { advanceNextSeq, readNextSeq } from "./session/session-sequences.ts"; @@ -109,7 +109,7 @@ export class SqliteStorage implements Storage { case "entry": { const { kind: _kind, ...entry } = write; insertEntryRow(this.db, entry); - // TODO: Maintain branch_entries/branch_meta for every inserted entry. + appendEntryToBranchIndex(this.db, entry); if (entry.type === "message") incrementMessageCount(this.db); break; } diff --git a/packages/session-backends/sqlite-node/test/storage.test.ts b/packages/session-backends/sqlite-node/test/storage.test.ts index 08e24b769..92814e124 100644 --- a/packages/session-backends/sqlite-node/test/storage.test.ts +++ b/packages/session-backends/sqlite-node/test/storage.test.ts @@ -29,7 +29,218 @@ function expectBranchPlan(plan: string[]): void { expect(plan.some((detail) => detail.includes("SCAN e"))).toBe(false); } +function insertCommitSessionRow(db: SqliteDatabase, nextSeq = 1): void { + sql`INSERT INTO session + (created_at, parent_session_id, storage_version, metadata, message_count, usage_payload, next_seq) + VALUES ( + ${1}, + ${null}, + ${1}, + ${null}, + ${0}, + ${JSON.stringify({ + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + })}, + ${nextSeq} + )`.run(db); +} + describe("SqliteStorage", () => { + it("commits root entries and append-to-tip entries into the branch index", async () => { + await withStorage(async (storage, db) => { + insertCommitSessionRow(db); + + expect( + await storage.commit({ + writes: [ + { + kind: "entry", + entry: { + id: "root", + parentId: null, + type: "message", + message: { role: "user", content: "root", timestamp: 10 }, + }, + }, + ], + }), + ).toEqual({ firstSeq: 1, seqs: [1], timestamp: 1_700_000_000_000 }); + expect( + await storage.commit({ + writes: [ + { + kind: "entry", + entry: { + id: "child", + parentId: "root", + type: "message", + message: { role: "user", content: "child", timestamp: 11 }, + }, + }, + ], + }), + ).toEqual({ firstSeq: 2, seqs: [2], timestamp: 1_700_000_000_000 }); + + expect(sql`SELECT branch_id, tip_entry_id, tip_seq FROM branch_meta`.all(db)).toEqual([ + { branch_id: "root", tip_entry_id: "child", tip_seq: 2 }, + ]); + expect( + sql`SELECT branch_id, entry_id, entry_seq, entry_type FROM branch_entries ORDER BY entry_seq`.all(db), + ).toEqual([ + { branch_id: "root", entry_id: "root", entry_seq: 1, entry_type: "message" }, + { branch_id: "root", entry_id: "child", entry_seq: 2, entry_type: "message" }, + ]); + expect((await storage.scanBranch({ start: "child" })).map((entry) => entry.id)).toEqual(["child", "root"]); + }); + }); + + it("commits divergent branch entries by materializing a new segment", async () => { + await withStorage(async (storage, db) => { + insertCommitSessionRow(db); + await storage.commit({ + writes: [ + { + kind: "entry", + entry: { + id: "root", + parentId: null, + type: "message", + message: { role: "user", content: "root", timestamp: 10 }, + }, + }, + { + kind: "entry", + entry: { + id: "left", + parentId: "root", + type: "message", + message: { role: "user", content: "left", timestamp: 11 }, + }, + }, + { + kind: "entry", + entry: { + id: "right", + parentId: "root", + type: "message", + message: { role: "user", content: "right", timestamp: 12 }, + }, + }, + ], + }); + + expect( + sql`SELECT branch_id, tip_entry_id, tip_seq, base_branch_id, base_seq FROM branch_meta ORDER BY branch_id`.all( + db, + ), + ).toEqual([ + { branch_id: "right", tip_entry_id: "right", tip_seq: 3, base_branch_id: null, base_seq: null }, + { branch_id: "root", tip_entry_id: "left", tip_seq: 2, base_branch_id: null, base_seq: null }, + ]); + expect( + sql`SELECT branch_id, entry_id, entry_seq, entry_type FROM branch_entries ORDER BY branch_id, entry_seq`.all( + db, + ), + ).toEqual([ + { branch_id: "right", entry_id: "root", entry_seq: 1, entry_type: "message" }, + { branch_id: "right", entry_id: "right", entry_seq: 3, entry_type: "message" }, + { branch_id: "root", entry_id: "root", entry_seq: 1, entry_type: "message" }, + { branch_id: "root", entry_id: "left", entry_seq: 2, entry_type: "message" }, + ]); + expect((await storage.scanBranch({ start: "right" })).map((entry) => entry.id)).toEqual(["right", "root"]); + expect((await storage.scanBranch({ start: "left" })).map((entry) => entry.id)).toEqual(["left", "root"]); + }); + }); + + it("bases divergent branch segments at the newest compaction", async () => { + await withStorage(async (storage, db) => { + insertCommitSessionRow(db); + await storage.commit({ + writes: [ + { + kind: "entry", + entry: { + id: "root", + parentId: null, + type: "message", + message: { role: "user", content: "root", timestamp: 10 }, + }, + }, + { + kind: "entry", + entry: { + id: "compact", + parentId: "root", + type: "compaction", + summary: "summary", + retainedTail: [], + tokensBefore: 1, + fromHook: false, + }, + }, + { + kind: "entry", + entry: { + id: "left", + parentId: "compact", + type: "message", + message: { role: "user", content: "left", timestamp: 11 }, + }, + }, + { + kind: "entry", + entry: { + id: "leaf", + parentId: "left", + type: "message", + message: { role: "user", content: "leaf", timestamp: 12 }, + }, + }, + { + kind: "entry", + entry: { + id: "right", + parentId: "left", + type: "message", + message: { role: "user", content: "right", timestamp: 13 }, + }, + }, + ], + }); + + expect( + sql`SELECT branch_id, tip_entry_id, tip_seq, base_branch_id, base_seq FROM branch_meta WHERE branch_id = ${"right"}`.get( + db, + ), + ).toEqual({ + branch_id: "right", + tip_entry_id: "right", + tip_seq: 5, + base_branch_id: "root", + base_seq: 2, + }); + expect( + sql`SELECT branch_id, entry_id, entry_seq, entry_type FROM branch_entries WHERE branch_id = ${"right"} ORDER BY entry_seq`.all( + db, + ), + ).toEqual([ + { branch_id: "right", entry_id: "left", entry_seq: 3, entry_type: "message" }, + { branch_id: "right", entry_id: "right", entry_seq: 5, entry_type: "message" }, + ]); + expect((await storage.scanBranch({ start: "right" })).map((entry) => entry.id)).toEqual([ + "right", + "left", + "compact", + "root", + ]); + }); + }); + it("gets entries by requested id order", async () => { await withStorage(async (storage, db) => { sql`INSERT INTO entries (id, parent_id, seq, type, custom_type, timestamp, payload)