From 8992320444bada968f80f2670b814dcd8b64b69f Mon Sep 17 00:00:00 2001 From: Cristina Poncela Cubeiro <140309543+cristinaponcela@users.noreply.github.com> Date: Mon, 17 Aug 2026 10:38:59 +0200 Subject: [PATCH] feat(sqlite): co-ordinate fork with open session --- .../sqlite-node/src/sqlite/repo.ts | 152 ++++++++++++------ .../sqlite-node/src/sqlite/session/entries.ts | 5 + .../src/sqlite/session/registers.ts | 8 + .../sqlite-node/src/sqlite/storage.ts | 44 ++++- .../sqlite-node/test/repo.test.ts | 35 ++++ 5 files changed, 189 insertions(+), 55 deletions(-) diff --git a/packages/session-backends/sqlite-node/src/sqlite/repo.ts b/packages/session-backends/sqlite-node/src/sqlite/repo.ts index a64fd2899..75604d38a 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/repo.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/repo.ts @@ -10,7 +10,7 @@ import type { import { StorageBackedSession } from "@earendil-works/pi-agent-core"; import { uuidv7 } from "@earendil-works/pi-ai"; import { applyInitialSchema } from "./migrations.ts"; -import { appendEntryToBranchIndex } from "./session/branch-entries.ts"; +import { appendEntryToBranchIndex, scanBranchEntries } from "./session/branch-entries.ts"; import { decodeEntryRow, type EntryRow, EntryRowWriter } from "./session/entries.ts"; import { insertInitialMainLaneRegisters } from "./session/registers.ts"; import { @@ -23,7 +23,7 @@ import { import { claimWriterLease, releaseWriterLease, renewWriterLease, type WriterLeaseRow } from "./session/writer-lease.ts"; import { SqliteOpenSession } from "./session.ts"; import { sql } from "./sql.ts"; -import { SqliteStorage } from "./storage.ts"; +import { SqliteStorage, type SqliteStorageSnapshot } from "./storage.ts"; import type { SqliteDatabase, SqliteDatabaseFactory } from "./types.ts"; export const SQLITE_STORAGE_VERSION = 1; @@ -100,7 +100,7 @@ function hasRegister(registers: readonly Register[], namespace: RegisterNamespac return registers.some((register) => register.namespace === namespace && register.key === key); } -function validateForkSourceSnapshot( +function validateTreeForkSourceSnapshot( entriesById: ReadonlyMap, registers: readonly Register[], sourceLeaves: readonly Register<"lane.leaf">[], @@ -130,6 +130,16 @@ function validateForkSourceSnapshot( } } +function validateBranchForkSourceSnapshot( + registers: readonly Register[], + sourceLeaves: readonly Register<"lane.leaf">[], +): void { + if (!sourceLeaves.some((register) => register.key === "main")) + throw new Error("Source session is missing main lane"); + if (!hasRegister(registers, "lane.state", "main")) + throw new Error('Source session lane "main" is missing lane.state'); +} + function selectForkContents( entriesById: ReadonlyMap, sourceLeaves: readonly Register<"lane.leaf">[], @@ -164,45 +174,64 @@ function selectForkContents( return { entryIds, laneToLeafId }; } +function buildForkSnapshot(source: SqliteStorageSnapshot, options: ForkOptions): ForkSnapshot { + const sourceEntries = source.entries; + const sourceRegisters = source.registers; + const sourceEntriesById = new Map(sourceEntries.map((entry) => [entry.id, entry])); + const sourceLeaves = sourceRegisters.filter((register) => isRegisterNamespace(register, "lane.leaf")); + if (options.scope === "tree") validateTreeForkSourceSnapshot(sourceEntriesById, sourceRegisters, sourceLeaves); + else validateBranchForkSourceSnapshot(sourceRegisters, sourceLeaves); + const { entryIds, laneToLeafId } = selectForkContents(sourceEntriesById, sourceLeaves, options); + const entries = sourceEntries.filter((entry) => entryIds.has(entry.id)); + const registers: Register[] = []; + let nextSeq = Math.max(0, ...entries.map((entry) => entry.seq)) + 1; + const setRegister = (namespace: RegisterNamespace, key: string, value: Register["value"]): void => { + registers.push({ namespace, key, value, seq: nextSeq++ } as Register); + }; + for (const [lane, leaf] of laneToLeafId) { + const configuration = sourceRegisters.find( + (register) => register.namespace === "lane.config" && register.key === lane, + ); + if (configuration !== undefined) setRegister("lane.config", lane, configuration.value); + setRegister("lane.leaf", lane, leaf); + setRegister("lane.state", lane, { currentOperationId: null, pendingNextRun: [] }); + } + for (const register of sourceRegisters) { + if ( + register.namespace === "fact.name" || + register.namespace === "fact.custom" || + (register.namespace === "fact.label" && entryIds.has(register.key)) + ) { + setRegister(register.namespace, register.key, register.value); + } + } + return { + entries, + registers, + messageCount: entries.filter((entry) => entry.type === "message").length, + nextSeq, + }; +} + +function readForkSourceEntries(db: SqliteDatabase, registers: readonly Register[], options: ForkOptions): Entry[] { + if (options.scope === "tree") return readSourceEntries(db); + const mainLeaf = registers.find((register) => register.namespace === "lane.leaf" && register.key === "main") as + | Register<"lane.leaf"> + | undefined; + if (mainLeaf === undefined) throw new Error("Source session is missing main lane"); + const requested = options.entryId ?? mainLeaf.value; + return requested === null ? [] : scanBranchEntries(db, { start: requested, order: "oldestFirst" }); +} + function createSqliteForkSnapshot(sourceDb: SqliteDatabase, options: ForkOptions): ForkSnapshot { sourceDb.exec("BEGIN"); let committed = false; try { - const sourceEntries = readSourceEntries(sourceDb); - const sourceRegisters = readSourceRegisters(sourceDb); - const sourceEntriesById = new Map(sourceEntries.map((entry) => [entry.id, entry])); - const sourceLeaves = sourceRegisters.filter((register) => isRegisterNamespace(register, "lane.leaf")); - validateForkSourceSnapshot(sourceEntriesById, sourceRegisters, sourceLeaves); - const { entryIds, laneToLeafId } = selectForkContents(sourceEntriesById, sourceLeaves, options); - const entries = sourceEntries.filter((entry) => entryIds.has(entry.id)); - const registers: Register[] = []; - let nextSeq = Math.max(0, ...entries.map((entry) => entry.seq)) + 1; - const setRegister = (namespace: RegisterNamespace, key: string, value: Register["value"]): void => { - registers.push({ namespace, key, value, seq: nextSeq++ } as Register); - }; - for (const [lane, leaf] of laneToLeafId) { - const configuration = sourceRegisters.find( - (register) => register.namespace === "lane.config" && register.key === lane, - ); - if (configuration !== undefined) setRegister("lane.config", lane, configuration.value); - setRegister("lane.leaf", lane, leaf); - setRegister("lane.state", lane, { currentOperationId: null, pendingNextRun: [] }); - } - for (const register of sourceRegisters) { - if ( - register.namespace === "fact.name" || - register.namespace === "fact.custom" || - (register.namespace === "fact.label" && entryIds.has(register.key)) - ) { - setRegister(register.namespace, register.key, register.value); - } - } - const snapshot = { - entries, - registers, - messageCount: entries.filter((entry) => entry.type === "message").length, - nextSeq, - }; + const registers = readSourceRegisters(sourceDb); + const snapshot = buildForkSnapshot( + { entries: readForkSourceEntries(sourceDb, registers, options), registers }, + options, + ); sourceDb.exec("COMMIT"); committed = true; return snapshot; @@ -226,6 +255,7 @@ export class SqliteSessionRepo { private readonly databaseFactory: SqliteDatabaseFactory; private readonly now: () => number; private readonly pendingIds = new Set(); + private readonly openStorages = new Map(); constructor(options: SqliteSessionRepoOptions) { this.directory = options.directory; @@ -352,12 +382,13 @@ export class SqliteSessionRepo { } async fork(source: SqliteSessionMetadata, options: ForkOptions): Promise { - await mkdir(this.directory, { recursive: true }); const createdAt = this.now(); const id = options.id ?? uuidv7(createdAt); this.reserveId(id); + const sourceStorage = this.openStorages.get(source.id); + const activeSourceSnapshot = sourceStorage?.snapshot(); + await mkdir(this.directory, { recursive: true }); const path = sessionPath(this.directory, id); - let sourceDb: SqliteDatabase | undefined; let db: SqliteDatabase | undefined; let lease: WriterLeaseRow | undefined; let reservedFile = false; @@ -368,18 +399,10 @@ export class SqliteSessionRepo { await file.close(); reservedFile = true; - sourceDb = await this.databaseFactory.open(source.path); - configureConnection(sourceDb); - const storedSource = metadataFromSessionRow( - source.path, - sessionIdFromPath(source.path), - readSingleSessionRow(sourceDb), - SQLITE_STORAGE_VERSION, - ); - if (storedSource.id !== source.id) { - throw new Error(`SQLite session path ${source.path} contains ${storedSource.id}, not ${source.id}`); - } - const snapshot = createSqliteForkSnapshot(sourceDb, options); + const snapshot = + activeSourceSnapshot === undefined + ? await this.createForkSnapshotFromClosedSource(source, options) + : buildForkSnapshot(await activeSourceSnapshot, options); const activeDb = await this.databaseFactory.open(path); db = activeDb; @@ -410,7 +433,6 @@ export class SqliteSessionRepo { if (reservedFile && !initialized) await removeSessionFiles(path, { force: true }); throw error; } finally { - sourceDb?.close(); if (session === undefined) { const failedDb = db; const failedLease = lease; @@ -423,6 +445,28 @@ export class SqliteSessionRepo { } } + private async createForkSnapshotFromClosedSource( + source: SqliteSessionMetadata, + options: ForkOptions, + ): Promise { + const sourceDb = await this.databaseFactory.open(source.path); + try { + configureConnection(sourceDb); + const storedSource = metadataFromSessionRow( + source.path, + sessionIdFromPath(source.path), + readSingleSessionRow(sourceDb), + SQLITE_STORAGE_VERSION, + ); + if (storedSource.id !== source.id) { + throw new Error(`SQLite session path ${source.path} contains ${storedSource.id}, not ${source.id}`); + } + return createSqliteForkSnapshot(sourceDb, options); + } finally { + sourceDb.close(); + } + } + private openStorageBackedSession( metadata: SqliteSessionMetadata, db: SqliteDatabase, @@ -435,6 +479,7 @@ export class SqliteSessionRepo { ); }; const storage = new SqliteStorage(db, { now: this.now, beforeCommit: renew }); + this.openStorages.set(metadata.id, storage); const session = new StorageBackedSession(metadata, storage); return new SqliteOpenSession(session, { renewWriterLease: renew, @@ -442,6 +487,7 @@ export class SqliteSessionRepo { renewIntervalMs: WRITER_LEASE_RENEW_INTERVAL_MS, onClose: () => { db.close(); + this.openStorages.delete(metadata.id); this.pendingIds.delete(metadata.id); }, }); diff --git a/packages/session-backends/sqlite-node/src/sqlite/session/entries.ts b/packages/session-backends/sqlite-node/src/sqlite/session/entries.ts index 1e1480445..22cbc6c81 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/session/entries.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/session/entries.ts @@ -139,6 +139,11 @@ export function readEntryRows(db: SqliteDatabase, ids: readonly string[]): Entry WHERE id IN (${placeholders})`.all(db); } +export function readAllEntryRows(db: SqliteDatabase): EntryRow[] { + return sql`SELECT id, parent_id, seq, type, custom_type, timestamp, payload + FROM entries ORDER BY seq ASC`.all(db); +} + export function scanEntryRows(db: SqliteDatabase, query: EntryScan): EntryRow[] { const filters: SqlQuery[] = []; if (query.type !== undefined) filters.push(sql`type = ${query.type}`); diff --git a/packages/session-backends/sqlite-node/src/sqlite/session/registers.ts b/packages/session-backends/sqlite-node/src/sqlite/session/registers.ts index 76e3721f4..8ec73a3e5 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/session/registers.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/session/registers.ts @@ -54,6 +54,14 @@ export function readRegisterRow( return row === undefined ? undefined : decodeRegisterRow(namespace, row); } +export function readAllRegisterRows(db: SqliteDatabase): Register[] { + return sql`SELECT namespace, key, seq, value FROM registers ORDER BY seq ASC` + .all(db) + .map( + (row) => ({ namespace: row.namespace, key: row.key, seq: row.seq, value: JSON.parse(row.value) }) as Register, + ); +} + function nextPrefixBoundary(prefix: string): string | undefined { if (prefix === "") return undefined; const codePoints = Array.from(prefix); diff --git a/packages/session-backends/sqlite-node/src/sqlite/storage.ts b/packages/session-backends/sqlite-node/src/sqlite/storage.ts index 949014560..d0d7af49b 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/storage.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/storage.ts @@ -3,6 +3,7 @@ import type { Entry, EntryScan, EntryStructure, + ForkOptions, Register, RegisterNamespace, SessionStats, @@ -14,8 +15,14 @@ import type { } from "@earendil-works/pi-agent-core"; import { prepareStorageCommit } from "@earendil-works/pi-agent-core"; import { appendEntryToBranchIndex, scanBranchEntries, scanBranchEntryStructures } from "./session/branch-entries.ts"; -import { decodeEntryRow, EntryRowWriter, readEntryRows, scanEntryRows } from "./session/entries.ts"; -import { deleteRegisterRow, listRegisterRows, readRegisterRow, setRegisterRow } from "./session/registers.ts"; +import { decodeEntryRow, EntryRowWriter, readAllEntryRows, readEntryRows, scanEntryRows } from "./session/entries.ts"; +import { + deleteRegisterRow, + listRegisterRows, + readAllRegisterRows, + readRegisterRow, + setRegisterRow, +} from "./session/registers.ts"; import { advanceNextSeq, readNextSeq } from "./session/session-sequences.ts"; import { addUsageToSessionStats, incrementMessageCount, readSessionStats } from "./session/session-stats.ts"; import { decodeUsageLedgerRow, scanUsageLedgerRows, UsageLedgerRowWriter } from "./session/usage-ledger.ts"; @@ -26,6 +33,11 @@ export interface SqliteStorageOptions { beforeCommit?: () => void; } +export interface SqliteStorageSnapshot { + entries: Entry[]; + registers: Register[]; +} + export class SqliteStorage implements Storage { private readonly db: SqliteDatabase; private readonly now: () => number; @@ -106,6 +118,34 @@ export class SqliteStorage implements Storage { return Promise.resolve(readSessionStats(this.db)); } + snapshot(options: ForkOptions = {}): Promise { + if (this.state !== "open") return Promise.reject(new Error("SqliteStorage is closed")); + const result = this.commitQueue.then(() => this.readSnapshot(options)); + this.commitQueue = result.then( + () => undefined, + () => undefined, + ); + return result; + } + + private readSnapshot(options: ForkOptions): SqliteStorageSnapshot { + const registers = readAllRegisterRows(this.db); + return { + entries: this.readSnapshotEntries(options, registers), + registers, + }; + } + + private readSnapshotEntries(options: ForkOptions, registers: readonly Register[]): Entry[] { + if (options.scope === "tree") return readAllEntryRows(this.db).map(decodeEntryRow); + const mainLeaf = registers.find((register) => register.namespace === "lane.leaf" && register.key === "main") as + | Register<"lane.leaf"> + | undefined; + if (mainLeaf === undefined) throw new Error("Source session is missing main lane"); + const requested = options.entryId ?? mainLeaf.value; + return requested === null ? [] : scanBranchEntries(this.db, { start: requested, order: "oldestFirst" }); + } + private applyCommit(transaction: Transaction): CommitResult { this.beforeCommit(); return this.db.transaction(() => { diff --git a/packages/session-backends/sqlite-node/test/repo.test.ts b/packages/session-backends/sqlite-node/test/repo.test.ts index 6fb33e495..82477e9be 100644 --- a/packages/session-backends/sqlite-node/test/repo.test.ts +++ b/packages/session-backends/sqlite-node/test/repo.test.ts @@ -148,6 +148,41 @@ describe("SqliteSessionRepo", () => { }); }); + it("forks at a serialized boundary for an open source session", async () => { + await withTempDir(async (directory) => { + const repo = new SqliteSessionRepo({ + directory, + databaseFactory: createNodeSqliteFactory(), + now: () => 1_700_000_000_000, + }); + const source = await repo.create({ id: "source" }); + const firstCommit = source.mutate("main", (mutator) => + mutator.commit({ + writes: [ + { kind: "entry", entry: { id: "root", parentId: null, type: "custom", customType: "root" } }, + { kind: "register", op: "set", namespace: "lane.leaf", key: "main", value: "root" }, + ], + }), + ); + const forkPromise = repo.fork(source.metadata, { id: "fork" }); + const secondCommit = source.mutate("main", (mutator) => + mutator.commit({ + writes: [ + { kind: "entry", entry: { id: "child", parentId: "root", type: "custom", customType: "child" } }, + { kind: "register", op: "set", namespace: "lane.leaf", key: "main", value: "child" }, + ], + }), + ); + + const [, fork] = await Promise.all([firstCommit, forkPromise]); + await secondCommit; + + expect(await fork.getLeafId()).toBe("root"); + expect((await fork.findEntries({ order: "asc" })).map((entry) => entry.id)).toEqual(["root"]); + await Promise.all([source.close(), fork.close()]); + }); + }); + it("forks one branch with scoped facts and a zero ledger", async () => { await withTempDir(async (directory) => { const repo = new SqliteSessionRepo({