diff --git a/packages/agent/docs/harness-v2.md b/packages/agent/docs/harness-v2.md index cafa8cf4c..350602bde 100644 --- a/packages/agent/docs/harness-v2.md +++ b/packages/agent/docs/harness-v2.md @@ -1649,7 +1649,7 @@ lane_moves (session_id, seq, lane, leaf_id) -- history; getLog parity facts (session_id, seq, kind, key, value) -- name, labels; latest by seq branch_entries (session_id, branch_id, entry_id, entry_seq, entry_type, custom_type) branch_tips (session_id, branch_id, tip_id) -- PRIMARY KEY (session_id, tip_id) -leases (session_id, owner, heartbeat) -- writer claim +leases (session_id, owner_id, fence, expires_at_ms) -- writer claim -- indexes records: (session_id, lane, type, seq), (session_id, lane, type, op_kind, seq) @@ -1657,6 +1657,8 @@ branch_entries: (session_id, branch_id, entry_type, entry_seq) (session_id, entry_id) -- reverse lookup: entry → branches ``` +`leases` enforces one writer per session with expiring, fenced claims. Storage renews the claim inside every write transaction and while idle. Repository-owned cleanup releases only its matching owner and fence. + `branch_entries` and `branch_tips` are a private read cache. No interface exposes them; no other backend has them; rebuilding them from parent pointers is an explicit repair operation, never a runtime fallback. Two invariants carry the whole design: diff --git a/packages/agent/src/harness/experimental/session/session.ts b/packages/agent/src/harness/experimental/session/session.ts index 9d8269feb..aa19cd7ae 100644 --- a/packages/agent/src/harness/experimental/session/session.ts +++ b/packages/agent/src/harness/experimental/session/session.ts @@ -37,7 +37,7 @@ function assertValidCursor(afterSeq: number | undefined): void { } } -function assertJsonSerializable(value: unknown): void { +export function assertJsonSerializable(value: unknown): void { const active = new WeakSet(); const stack: JsonValidationFrame[] = [{ value }]; while (stack.length > 0) { diff --git a/packages/agent/test/harness/branch-query.test.ts b/packages/agent/test/harness/branch-query.test.ts index fb2f4b062..d67161c63 100644 --- a/packages/agent/test/harness/branch-query.test.ts +++ b/packages/agent/test/harness/branch-query.test.ts @@ -171,7 +171,7 @@ describe("bounded session branch queries", () => { const db = await sqlite.open(databasePath); try { await db - .prepare("UPDATE session_entries SET payload = ? WHERE session_id = ? AND id = ?") + .prepare("UPDATE entries SET payload = ? WHERE session_id = ? AND id = ?") .run("not json", "bounded-sqlite", middleId); const branch = await db .prepare("SELECT branch_id FROM branch_entries WHERE session_id = ? AND entry_id = ?") @@ -187,21 +187,6 @@ describe("bounded session branch queries", () => { expect((await session.findEntriesOnBranch({ start: tailId, stopAtId: tailId })).map((entry) => entry.id)).toEqual( [tailId], ); - const inspection = await sqlite.open(databasePath); - try { - const repaired = await inspection - .prepare( - `SELECT entry_id FROM branch_entries - WHERE session_id = ? AND branch_id = ( - SELECT branch_id FROM branch_entries WHERE session_id = ? AND entry_id = ? LIMIT 1 - ) - ORDER BY entry_seq`, - ) - .all<{ entry_id: string }>("bounded-sqlite", "bounded-sqlite", tailId); - expect(repaired.map((row) => row.entry_id)).toEqual([rootId, tailId]); - } finally { - await inspection.close(); - } expect( ( await session.findEntriesOnBranch({ @@ -214,7 +199,7 @@ describe("bounded session branch queries", () => { ).toEqual([rootId]); await expect(session.findEntriesOnBranch({ start: tailId, limit: 2 })).rejects.toMatchObject({ code: "invalid_entry", - message: expect.stringContaining(`failed to decode entry ${middleId}`), + message: expect.stringContaining(`Entry ${middleId} not found`), }); }); @@ -236,7 +221,7 @@ describe("bounded session branch queries", () => { const db = await sqlite.open(databasePath); try { await db - .prepare("UPDATE session_entries SET payload = ? WHERE session_id = ? AND id = ?") + .prepare("UPDATE entries SET payload = ? WHERE session_id = ? AND id = ?") .run("{}", "invalid-filtered-sqlite", customId); } finally { await db.close(); @@ -249,7 +234,7 @@ describe("bounded session branch queries", () => { const invalidJsonDb = await sqlite.open(databasePath); try { await invalidJsonDb - .prepare("UPDATE session_entries SET payload = ? WHERE session_id = ? AND id = ?") + .prepare("UPDATE entries SET payload = ? WHERE session_id = ? AND id = ?") .run("not json", "invalid-filtered-sqlite", customId); } finally { await invalidJsonDb.close(); @@ -277,7 +262,7 @@ describe("bounded session branch queries", () => { const db = await sqlite.open(databasePath); try { await db - .prepare("UPDATE session_entries SET parent_id = ? WHERE session_id = ? AND id = ?") + .prepare("UPDATE entries SET parent_id = ? WHERE session_id = ? AND id = ?") .run("missing-parent", "bounded-corrupt-sqlite", childId); } finally { await db.close(); @@ -289,17 +274,17 @@ describe("bounded session branch queries", () => { (await session.findEntriesOnBranch({ start: childId, stopAtType: "message" })).map((entry) => entry.id), ).toEqual([childId]); await expect(session.findEntriesOnBranch({ start: childId })).rejects.toMatchObject({ - code: "invalid_session", + code: "invalid_entry", message: expect.stringContaining("Entry missing-parent not found"), }); const cycleDb = await sqlite.open(databasePath); try { await cycleDb - .prepare("UPDATE session_entries SET parent_id = ? WHERE session_id = ? AND id = ?") + .prepare("UPDATE entries SET parent_id = ? WHERE session_id = ? AND id = ?") .run(rootId, "bounded-corrupt-sqlite", childId); await cycleDb - .prepare("UPDATE session_entries SET parent_id = ? WHERE session_id = ? AND id = ?") + .prepare("UPDATE entries SET parent_id = ? WHERE session_id = ? AND id = ?") .run(childId, "bounded-corrupt-sqlite", rootId); } finally { await cycleDb.close(); @@ -311,24 +296,8 @@ describe("bounded session branch queries", () => { (await session.findEntriesOnBranch({ start: childId, stopAtType: "message" })).map((entry) => entry.id), ).toEqual([childId]); await expect(session.findEntriesOnBranch({ start: childId })).rejects.toMatchObject({ - code: "invalid_session", - message: expect.stringContaining(`cycle in parent chain at entry ${childId}`), + code: "invalid_entry", + message: expect.stringContaining(`Entry ${childId} not found`), }); }); - - it("provides identical SQLite query semantics", async () => { - const root = createTempDir(); - const repo = new SqliteSessionRepository({ - env: new NodeExecutionEnv({ cwd: root }), - sqlite: createNodeSqliteFactory(), - databasePath: join(root, "sessions.sqlite"), - }); - ownedRepositories.push(repo); - const session = await repo.create({ id: "sqlite", cwd: root }); - const expected = await verifyBranchQueries(session); - const reopened = await repo.open(await session.getMetadata()); - expect( - (await reopened.findEntriesOnBranch({ start: expected.tail, order: "oldestFirst" })).map((entry) => entry.id), - ).toEqual(expected.fullPath); - }); }); diff --git a/packages/agent/test/harness/experimental/session/sqlite.test.ts b/packages/agent/test/harness/experimental/session/sqlite.test.ts new file mode 100644 index 000000000..b5ce72922 --- /dev/null +++ b/packages/agent/test/harness/experimental/session/sqlite.test.ts @@ -0,0 +1,58 @@ +import { join } from "node:path"; +import type { SessionMetadata, SessionRepo } from "@earendil-works/pi-agent-core/experimental"; +import { describe, it } from "vitest"; +import { + createNodeSqliteFactory, + type SqliteSessionMetadata, + SqliteSessionRepository, +} from "../../../../../storage/sqlite-node/src/index.ts"; +import { NodeExecutionEnv } from "../../../../src/harness/env/nodejs.ts"; +import { + createSessionBackendConformance, + type SessionBackendFixture, +} from "../../../../src/harness/experimental/session/testing/index.ts"; +import { createTempDir } from "../../session-test-utils.ts"; + +function requireSqliteMetadata(metadata: SessionMetadata): SqliteSessionMetadata { + const cwd = "cwd" in metadata ? metadata.cwd : undefined; + if (typeof cwd !== "string") { + throw new Error(`Expected SQLite metadata for session ${metadata.id}`); + } + const path = "path" in metadata ? metadata.path : undefined; + if (typeof path !== "string") { + throw new Error(`Expected SQLite metadata for session ${metadata.id}`); + } + return { ...metadata, cwd, path }; +} + +const conformance = createSessionBackendConformance(async () => { + const root = createTempDir(); + const sqliteRepository = new SqliteSessionRepository({ + env: new NodeExecutionEnv({ cwd: root }), + sqlite: createNodeSqliteFactory(), + databasePath: join(root, "sessions.sqlite"), + }); + const repository: SessionRepo = { + create: (options = {}) => sqliteRepository.create({ ...options, cwd: root }), + open: (metadata) => sqliteRepository.open(requireSqliteMetadata(metadata)), + list: () => sqliteRepository.list(), + delete: (metadata) => sqliteRepository.delete(requireSqliteMetadata(metadata)), + fork: (source, options = {}) => sqliteRepository.fork(requireSqliteMetadata(source), { ...options, cwd: root }), + }; + return { + repository, + async [Symbol.asyncDispose]() { + await sqliteRepository.close(); + }, + } satisfies SessionBackendFixture; +}); + +describe("SqliteSessionRepository conformance", () => { + for (const group of new Set(conformance.map((testCase) => testCase.group))) { + describe(group, () => { + for (const testCase of conformance.filter((candidate) => candidate.group === group)) { + it(testCase.name, () => testCase.run()); + } + }); + } +}); diff --git a/packages/agent/test/harness/session-test-utils.ts b/packages/agent/test/harness/session-test-utils.ts index 621ce84c5..81f98491d 100644 --- a/packages/agent/test/harness/session-test-utils.ts +++ b/packages/agent/test/harness/session-test-utils.ts @@ -2,7 +2,18 @@ import { existsSync, mkdirSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import type { + BranchSummaryEntry, + CompactionEntry, + Session as CoreSession, + Entry, + MessageEntry, + ModelChangeEntry, + ThinkingLevelChangeEntry, +} from "@earendil-works/pi-agent-core/experimental"; +import type { Usage } from "@earendil-works/pi-ai"; import { afterEach } from "vitest"; +import type { SqliteSessionMetadata } from "../../../storage/sqlite-node/src/index.ts"; import { InMemorySessionRepository } from "../../src/harness/session/memory-repo.ts"; import type { Session } from "../../src/harness/session/session.ts"; @@ -38,6 +49,123 @@ export function createAssistantMessage(text: string): AgentMessage { }; } +export type SqliteTestSession = CoreSession; +export type SqliteTestMessage = MessageEntry["message"]; + +export async function appendSqliteCompaction( + session: SqliteTestSession, + summary: string, + _firstKeptEntryId: string | undefined, + tokensBefore: number, + details?: unknown, + _fromHook?: boolean, + usage?: Usage, + retainedTail: SqliteTestMessage[] = [], +): Promise { + const provisioned = { + type: "compaction", + id: session.idGenerator.next(), + summary, + retainedTail, + tokensBefore, + ...(details === undefined ? {} : { details }), + ...(usage === undefined ? {} : { usage }), + } satisfies Omit; + const entry = await session.appendEntry(provisioned, "main"); + return entry.id; +} + +export async function moveSqliteMainLane( + session: SqliteTestSession, + entryId: string | null, + summary?: { summary: string; details?: unknown; usage?: Usage; fromHook?: boolean }, +): Promise { + await session.moveLane("main", entryId); + if (!summary) return undefined; + const provisioned = { + type: "branch_summary", + id: session.idGenerator.next(), + fromId: entryId ?? "root", + summary: summary.summary, + ...(summary.details === undefined ? {} : { details: summary.details }), + ...(summary.usage === undefined ? {} : { usage: summary.usage }), + } satisfies Omit; + const entry = await session.appendEntry(provisioned, "main"); + return entry.id; +} + +export async function getSqliteBranch(session: SqliteTestSession, fromId?: string | null): Promise { + const start = fromId === undefined ? await session.getLeafId() : fromId; + if (start === null) return []; + const newestWindow = await session.findEntriesOnBranch({ start, stopAtType: "compaction" }); + return newestWindow.reverse(); +} + +export async function getSqliteEntries( + session: SqliteTestSession, + options?: { afterEntrySeq?: number; limit?: number }, +): Promise { + return session.findEntries({ + order: "oldestFirst", + limit: options?.limit, + cursor: options?.afterEntrySeq === undefined ? undefined : { afterSeq: options.afterEntrySeq }, + }); +} + +export async function appendSqliteSessionName(session: SqliteTestSession, name: string): Promise { + await session.setName(name.replace(/[\r\n]+/g, " ").trim()); +} + +export async function appendSqliteLabel( + session: SqliteTestSession, + targetId: string, + label: string | undefined, +): Promise { + await session.setLabel(targetId, label); +} + +export async function buildSqliteContext(session: SqliteTestSession): Promise<{ messages: SqliteTestMessage[] }> { + const entries = await getSqliteBranch(session); + const messages = entries.flatMap((entry): SqliteTestMessage[] => { + if (entry.type === "message") return [entry.message]; + if (entry.type === "compaction") return entry.retainedTail; + return []; + }); + return { messages }; +} + +export async function appendSqliteThinkingLevelChange( + session: SqliteTestSession, + thinkingLevel: string, +): Promise { + const entry = await session.appendEntry( + { + type: "thinking_level_change", + id: session.idGenerator.next(), + thinkingLevel, + } satisfies Omit, + "main", + ); + return entry.id; +} + +export async function appendSqliteModelChange( + session: SqliteTestSession, + provider: string, + modelId: string, +): Promise { + const entry = await session.appendEntry( + { + type: "model_change", + id: session.idGenerator.next(), + provider, + modelId, + } satisfies Omit, + "main", + ); + return entry.id; +} + const tempDirs: string[] = []; export function createTempDir(): string { diff --git a/packages/agent/test/harness/sqlite-adapter.test.ts b/packages/agent/test/harness/sqlite-adapter.test.ts new file mode 100644 index 000000000..34daaf286 --- /dev/null +++ b/packages/agent/test/harness/sqlite-adapter.test.ts @@ -0,0 +1,36 @@ +import { describe, expect, it } from "vitest"; +import { createNodeSqliteFactory } from "../../../storage/sqlite-node/src/index.ts"; + +describe("node:sqlite adapter", () => { + it("runs transaction callbacks synchronously", async () => { + const db = await createNodeSqliteFactory().open(":memory:"); + try { + db.exec("CREATE TABLE values_table (value INTEGER NOT NULL)"); + const result = db.transaction(() => { + db.prepare("INSERT INTO values_table (value) VALUES (?)").run(42); + return "committed"; + }); + + expect(result).toBe("committed"); + expect(db.prepare("SELECT value FROM values_table").get()).toEqual({ value: 42 }); + } finally { + db.close(); + } + }); + + it("rejects asynchronous transaction callbacks", async () => { + const db = await createNodeSqliteFactory().open(":memory:"); + try { + db.exec("CREATE TABLE values_table (value INTEGER NOT NULL)"); + const asynchronous = async () => { + db.prepare("INSERT INTO values_table (value) VALUES (?)").run(42); + await Promise.resolve(); + }; + expect(() => db.transaction(asynchronous)).toThrow("SQLite transaction callbacks must be synchronous"); + await Promise.resolve(); + expect(db.prepare("SELECT value FROM values_table").all()).toEqual([]); + } finally { + db.close(); + } + }); +}); diff --git a/packages/agent/test/harness/sqlite-branch-cache.test.ts b/packages/agent/test/harness/sqlite-branch-cache.test.ts index 28bdfaa8c..30f3dbb7a 100644 --- a/packages/agent/test/harness/sqlite-branch-cache.test.ts +++ b/packages/agent/test/harness/sqlite-branch-cache.test.ts @@ -4,7 +4,13 @@ import { join } from "node:path"; import { describe, expect, it } from "vitest"; import { createNodeSqliteFactory, SqliteSessionRepository } from "../../../storage/sqlite-node/src/index.ts"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; -import { createAssistantMessage, createUserMessage } from "./session-test-utils.ts"; +import { + appendSqliteCompaction, + createAssistantMessage, + createUserMessage, + getSqliteBranch, + moveSqliteMainLane, +} from "./session-test-utils.ts"; function createTempDir(): string { return mkdtempSync(join(tmpdir(), "pi-agent-sqlite-branch-cache-")); @@ -20,9 +26,9 @@ describe("SQLite branch cache", () => { const session = await repo.create({ cwd: root, id: "session-1" }); const rootId = await session.appendMessage(createUserMessage("root")); const keptId = await session.appendMessage(createUserMessage("kept")); - const compactionId = await session.appendCompaction("summary", keptId, 100); + const compactionId = await appendSqliteCompaction(session, "summary", keptId, 100); await session.appendMessage(createAssistantMessage("first child")); - await session.moveTo(compactionId); + await moveSqliteMainLane(session, compactionId); const branchedId = await session.appendMessage(createAssistantMessage("branched child")); const db = await sqlite.open(databasePath); @@ -49,19 +55,19 @@ describe("SQLite branch cache", () => { const session = await repo.create({ cwd: root, id: "session-1" }); const oldId = await session.appendMessage(createUserMessage("old")); const keptId = await session.appendMessage(createUserMessage("kept")); - const compactionId = await session.appendCompaction("summary", keptId, 100); + const compactionId = await appendSqliteCompaction(session, "summary", keptId, 100); const leafId = await session.appendMessage(createAssistantMessage("new")); const db = await sqlite.open(databasePath); try { await db - .prepare("UPDATE session_entries SET payload = ? WHERE session_id = ? AND id = ?") + .prepare("UPDATE entries SET payload = ? WHERE session_id = ? AND id = ?") .run("not json", "session-1", oldId); } finally { await db.close(); } - expect((await session.getBranch()).map((entry) => entry.id)).toEqual([keptId, compactionId, leafId]); + expect((await getSqliteBranch(session)).map((entry) => entry.id)).toEqual([compactionId, leafId]); }); it("preserves nested compaction boundaries when reading the cache", async () => { @@ -71,7 +77,8 @@ describe("SQLite branch cache", () => { const repo = new SqliteSessionRepository({ env, sqlite: createNodeSqliteFactory(), databasePath }); const session = await repo.create({ cwd: root, id: "session-1" }); const rootId = await session.appendMessage(createUserMessage("root")); - const firstCompactionId = await session.appendCompaction( + const firstCompactionId = await appendSqliteCompaction( + session, "first summary", rootId, 100, @@ -81,18 +88,15 @@ describe("SQLite branch cache", () => { [], ); const middleId = await session.appendMessage(createUserMessage("middle")); - const secondCompactionId = await session.appendCompaction("second summary", rootId, 200); + const secondCompactionId = await appendSqliteCompaction(session, "second summary", rootId, 200); const leafId = await session.appendMessage(createAssistantMessage("new")); - expect((await session.getBranch()).map((entry) => entry.id)).toEqual([ - firstCompactionId, - middleId, - secondCompactionId, - leafId, - ]); + expect(firstCompactionId).not.toBe(secondCompactionId); + expect(middleId).not.toBe(leafId); + expect((await getSqliteBranch(session)).map((entry) => entry.id)).toEqual([secondCompactionId, leafId]); }); - it("repairs a missing branch cache from canonical parent links", async () => { + it("fails loudly when the private branch cache is missing", async () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); @@ -110,25 +114,11 @@ describe("SQLite branch cache", () => { await db.close(); } - expect((await session.getBranch()).map((entry) => entry.id)).toEqual([rootId, childId]); - - const inspection = await sqlite.open(databasePath); - try { - const rows = await inspection - .prepare("SELECT branch_id, entry_id FROM branch_entries WHERE session_id = ? ORDER BY entry_seq") - .all<{ branch_id: string; entry_id: string }>("session-1"); - expect(rows.map((row) => row.entry_id)).toEqual([rootId, childId]); - const tip = await inspection - .prepare("SELECT branch_id, tip_id FROM branch_tips WHERE session_id = ?") - .get<{ branch_id: string; tip_id: string }>("session-1"); - expect(rows).toHaveLength(2); - expect(tip).toEqual({ branch_id: rows[0]?.branch_id, tip_id: childId }); - } finally { - await inspection.close(); - } + expect(rootId).not.toBe(childId); + await expect(getSqliteBranch(session)).rejects.toMatchObject({ code: "invalid_entry" }); }); - it("rolls back an interrupted branch cache repair", async () => { + it("does not repair the private branch cache during normal reads", async () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); @@ -141,35 +131,49 @@ describe("SQLite branch cache", () => { try { await db.prepare("DELETE FROM branch_tips WHERE session_id = ?").run("session-1"); await db.prepare("DELETE FROM branch_entries WHERE session_id = ?").run("session-1"); - await db.exec(` - CREATE TRIGGER fail_branch_cache_repair - BEFORE INSERT ON branch_tips - BEGIN - SELECT RAISE(ABORT, 'repair failed'); - END; - `); } finally { await db.close(); } - await expect(session.getBranch()).rejects.toMatchObject({ - code: "storage", - message: expect.stringContaining("Failed to rebuild SQLite branch cache"), - }); + await expect(getSqliteBranch(session)).rejects.toMatchObject({ code: "invalid_entry" }); const inspection = await sqlite.open(databasePath); try { expect( await inspection.prepare("SELECT entry_id FROM branch_entries WHERE session_id = ?").all("session-1"), ).toEqual([]); - await inspection.exec("DROP TRIGGER fail_branch_cache_repair"); } finally { await inspection.close(); } - expect(await session.getBranch()).toHaveLength(1); }); - it("repairs a missing source branch cache while forking transactionally", async () => { + it("repairs the private branch cache explicitly", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const repo = new SqliteSessionRepository({ env, sqlite, databasePath }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const rootId = await session.appendMessage(createUserMessage("root")); + const childId = await session.appendMessage(createAssistantMessage("child")); + const metadata = await session.getMetadata(); + + const db = await sqlite.open(databasePath); + try { + await db.prepare("DELETE FROM branch_tips WHERE session_id = ?").run("session-1"); + await db.prepare("DELETE FROM branch_entries WHERE session_id = ?").run("session-1"); + } finally { + await db.close(); + } + + await expect(getSqliteBranch(session)).rejects.toMatchObject({ code: "invalid_entry" }); + + await repo.repairBranchCache(metadata); + + expect((await getSqliteBranch(session)).map((entry) => entry.id)).toEqual([rootId, childId]); + }); + + it("fails when forking from a source with a missing branch cache", async () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); @@ -187,16 +191,18 @@ describe("SQLite branch cache", () => { await db.close(); } - const fork = await repo.fork(await source.getMetadata(), { - cwd: root, - id: "fork", - entryId: childId, - position: "at", - }); - expect((await fork.getEntries()).map((entry) => entry.id)).toEqual([rootId, childId]); + expect(rootId).not.toBe(childId); + await expect( + repo.fork(await source.getMetadata(), { + cwd: root, + id: "fork", + entryId: childId, + position: "at", + }), + ).rejects.toMatchObject({ code: "invalid_fork_target" }); }); - it("repairs a stale branch cache from canonical parent links", async () => { + it("fails when the private branch cache is stale", async () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); @@ -210,31 +216,16 @@ describe("SQLite branch cache", () => { const db = await sqlite.open(databasePath); try { await db - .prepare("UPDATE session_entries SET parent_id = ? WHERE session_id = ? AND id = ?") + .prepare("UPDATE entries SET parent_id = ? WHERE session_id = ? AND id = ?") .run(rootId, "session-1", leafId); } finally { await db.close(); } - expect( - (await session.findEntriesOnBranch({ start: leafId, order: "oldestFirst" })).map((entry) => entry.id), - ).toEqual([rootId, leafId]); - const inspection = await sqlite.open(databasePath); - try { - const rows = await inspection - .prepare( - `SELECT entry_id FROM branch_entries - WHERE session_id = ? AND branch_id = ( - SELECT branch_id FROM branch_entries WHERE session_id = ? AND entry_id = ? LIMIT 1 - ) - ORDER BY entry_seq`, - ) - .all<{ entry_id: string }>("session-1", "session-1", leafId); - expect(rows.map((row) => row.entry_id)).toEqual([rootId, leafId]); - expect(rows.map((row) => row.entry_id)).not.toContain(staleId); - } finally { - await inspection.close(); - } + expect(staleId).not.toBe(leafId); + await expect(session.findEntriesOnBranch({ start: leafId, order: "oldestFirst" })).rejects.toMatchObject({ + code: "invalid_entry", + }); }); it("deletes branch entries and tips with the session", async () => { diff --git a/packages/agent/test/harness/sqlite-leases.test.ts b/packages/agent/test/harness/sqlite-leases.test.ts new file mode 100644 index 000000000..9ea4f9c6a --- /dev/null +++ b/packages/agent/test/harness/sqlite-leases.test.ts @@ -0,0 +1,136 @@ +import { join } from "node:path"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { + createNodeSqliteFactory, + type SqliteSessionMetadata, + SqliteSessionRepository, +} from "../../../storage/sqlite-node/src/index.ts"; +import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; +import { createTempDir, createUserMessage } from "./session-test-utils.ts"; + +const repositories: SqliteSessionRepository[] = []; + +function createRepository(root: string, databasePath: string, lease?: { ttlMs: number; heartbeatIntervalMs: number }) { + const repository = new SqliteSessionRepository({ + env: new NodeExecutionEnv({ cwd: root }), + sqlite: createNodeSqliteFactory(), + databasePath, + writerLease: lease, + }); + repositories.push(repository); + return repository; +} + +afterEach(async () => { + vi.useRealTimers(); + for (const repository of repositories.splice(0)) await repository.close(); +}); + +describe("SQLite session writer leases", () => { + it("rejects a second writer until the first session releases its claim", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const firstRepository = createRepository(root, databasePath); + const secondRepository = createRepository(root, databasePath); + const first = await firstRepository.create({ cwd: root, id: "session-1" }); + const metadata = await first.getMetadata(); + + await expect(secondRepository.open(metadata)).rejects.toMatchObject({ + code: "storage", + message: expect.stringContaining("already has an active writer"), + }); + + await firstRepository.close(); + const second = await secondRepository.open(metadata); + await expect(second.appendMessage(createUserMessage("new owner"))).resolves.toBeTypeOf("string"); + }); + + it("fences a stale owner after an expired lease is acquired by another writer", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const lease = { ttlMs: 120_000, heartbeatIntervalMs: 60_000 }; + const firstRepository = createRepository(root, databasePath, lease); + const secondRepository = createRepository(root, databasePath, lease); + const first = await firstRepository.create({ cwd: root, id: "session-1" }); + const metadata = await first.getMetadata(); + const sqlite = createNodeSqliteFactory(); + const db = await sqlite.open(databasePath); + try { + await db.prepare("UPDATE leases SET expires_at_ms = 0 WHERE session_id = ?").run(metadata.id); + } finally { + await db.close(); + } + + const second = await secondRepository.open(metadata); + await expect(first.appendMessage(createUserMessage("stale owner"))).rejects.toMatchObject({ + code: "storage", + message: expect.stringContaining("writer lease was lost"), + }); + expect(await second.findEntries()).toEqual([]); + + const inspection = await sqlite.open(databasePath); + let currentLease: { owner_id: string; fence: number } | undefined; + try { + currentLease = await inspection + .prepare("SELECT owner_id, fence FROM leases WHERE session_id = ?") + .get<{ owner_id: string; fence: number }>(metadata.id); + expect(currentLease?.fence).toBe(2); + } finally { + await inspection.close(); + } + + await firstRepository.close(); + const afterStaleClose = await sqlite.open(databasePath); + try { + expect( + await afterStaleClose.prepare("SELECT owner_id, fence FROM leases WHERE session_id = ?").get(metadata.id), + ).toEqual(currentLease); + } finally { + await afterStaleClose.close(); + } + await expect(second.appendMessage(createUserMessage("current owner"))).resolves.toBeTypeOf("string"); + }); + + it("serializes lease-checked writes for sessions sharing one database connection", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const repository = createRepository(root, databasePath); + const first = await repository.create({ cwd: root, id: "session-1" }); + const second = await repository.create({ cwd: root, id: "session-2" }); + + await expect( + Promise.all([ + first.appendMessage(createUserMessage("first")), + second.appendMessage(createUserMessage("second")), + ]), + ).resolves.toHaveLength(2); + }); + + it("renews an idle writer lease with a heartbeat", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-01-01T00:00:00.000Z")); + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const repository = createRepository(root, databasePath, { ttlMs: 30_000, heartbeatIntervalMs: 10_000 }); + const session = await repository.create({ cwd: root, id: "session-1" }); + const metadata = (await session.getMetadata()) as SqliteSessionMetadata; + const sqlite = createNodeSqliteFactory(); + + const readExpiry = async (): Promise => { + const db = await sqlite.open(databasePath); + try { + return ( + await db + .prepare("SELECT expires_at_ms FROM leases WHERE session_id = ?") + .get<{ expires_at_ms: number }>(metadata.id) + )?.expires_at_ms; + } finally { + await db.close(); + } + }; + + const initialExpiry = await readExpiry(); + await vi.advanceTimersByTimeAsync(10_000); + expect(await readExpiry()).toBe((initialExpiry ?? 0) + 10_000); + }); +}); diff --git a/packages/agent/test/harness/sqlite-migrations.test.ts b/packages/agent/test/harness/sqlite-migrations.test.ts index 9056754bd..397febe1c 100644 --- a/packages/agent/test/harness/sqlite-migrations.test.ts +++ b/packages/agent/test/harness/sqlite-migrations.test.ts @@ -3,9 +3,7 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { describe, expect, it } from "vitest"; import { - applyMigrations, createNodeSqliteFactory, - loadMigrations, type SqliteDatabase, type SqliteDatabaseFactory, type SqliteRunResult, @@ -13,30 +11,38 @@ import { SqliteSessionRepository, type SqliteStatement, } from "../../../storage/sqlite-node/src/index.ts"; -import { SqliteSessionConnection } from "../../../storage/sqlite-node/src/sqlite/storage/index.ts"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; -import { createAssistantMessage, createUserMessage } from "./session-test-utils.ts"; +import { + appendSqliteCompaction, + appendSqliteLabel, + appendSqliteSessionName, + buildSqliteContext, + createAssistantMessage, + createUserMessage, + getSqliteEntries, + moveSqliteMainLane, +} from "./session-test-utils.ts"; function createTempDir(): string { return mkdtempSync(join(tmpdir(), "pi-agent-sqlite-")); } class ThrowingStatement implements SqliteStatement { - private readonly onRun: () => Promise; + private readonly onRun: () => SqliteRunResult; - constructor(onRun: () => Promise) { + constructor(onRun: () => SqliteRunResult) { this.onRun = onRun; } - async run(..._params: unknown[]): Promise { + run(..._params: unknown[]): SqliteRunResult { return this.onRun(); } - async get(..._params: unknown[]): Promise { + get(..._params: unknown[]): TRow | undefined { return undefined; } - async all(..._params: unknown[]): Promise { + all(..._params: unknown[]): TRow[] { return []; } } @@ -49,17 +55,17 @@ class CountingDatabase implements SqliteDatabase { this.statementFactory = statementFactory; } - async exec(_sql: string): Promise {} + exec(_sql: string): void {} prepare(sql: string): SqliteStatement { return this.statementFactory(sql); } - async transaction(fn: () => Promise): Promise { + transaction(fn: () => T): T { return fn(); } - async close(): Promise { + close(): void { this.closeCount += 1; } } @@ -80,9 +86,9 @@ function createCloseCountingSqliteFactory(): { exec: (sql) => db.exec(sql), prepare: (sql) => db.prepare(sql), transaction: (fn) => db.transaction(fn), - async close() { + close() { counts.closes += 1; - await db.close(); + db.close(); }, }; }, @@ -102,7 +108,7 @@ describe("SQLite migrations", () => { const db = await sqlite.open(databasePath); try { const rows = await db.prepare("SELECT id FROM migrations ORDER BY id").all<{ id: string }>(); - expect(rows.map((row) => row.id)).toEqual(["001_initial.sql", "002_branch_tips.sql"]); + expect(rows.map((row) => row.id)).toEqual(["001_initial.sql"]); const tables = await db .prepare("SELECT name, sql FROM sqlite_master WHERE type = 'table' ORDER BY name") .all<{ name: string; sql: string | null }>(); @@ -110,16 +116,17 @@ describe("SQLite migrations", () => { expect.arrayContaining([ "migrations", "sessions", - "session_entries", + "entries", "session_sequences", "branch_entries", "branch_tips", - "session_materialized", - "entry_materialized", ]), ); const sessionColumns = await db.prepare("PRAGMA table_info(sessions)").all<{ name: string }>(); - expect(sessionColumns.map((column) => column.name)).toContain("active_leaf_id"); + expect(sessionColumns.map((column) => column.name)).not.toContain("leaf_id"); + expect(tables.map((row) => row.name)).toEqual( + expect.arrayContaining(["lanes", "records", "lane_moves", "facts", "leases", "session_stats"]), + ); const branchIndexes = await db .prepare("SELECT name FROM sqlite_master WHERE type = 'index' AND tbl_name = 'branch_entries'") .all<{ name: string }>(); @@ -128,10 +135,13 @@ describe("SQLite migrations", () => { for (const tableName of [ "sessions", "session_sequences", + "session_stats", "branch_entries", "branch_tips", - "session_materialized", - "entry_materialized", + "lanes", + "records", + "lane_moves", + "facts", ]) { const table = tables.find((row) => row.name === tableName); expect(table?.sql).toContain("WITHOUT ROWID"); @@ -141,37 +151,6 @@ describe("SQLite migrations", () => { } }); - it("clears legacy branch projections when adding explicit tips", async () => { - const root = createTempDir(); - const databasePath = join(root, "sessions.sqlite"); - const sqlite = createNodeSqliteFactory(); - const db = await sqlite.open(databasePath); - try { - const initial = (await loadMigrations()).find((migration) => migration.id === "001_initial.sql"); - if (!initial) throw new Error("Missing initial SQLite migration"); - await db.exec(initial.sql); - await db.exec("CREATE TABLE migrations (id TEXT PRIMARY KEY, applied_at TEXT NOT NULL)"); - await db - .prepare("INSERT INTO migrations (id, applied_at) VALUES (?, ?)") - .run(initial.id, "2026-01-01T00:00:00.000Z"); - await db - .prepare("INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) VALUES (?, ?, ?, ?)") - .run("session-1", "legacy-branch", "entry-1", 1); - - await applyMigrations(db); - - expect(await db.prepare("SELECT entry_id FROM branch_entries").all<{ entry_id: string }>()).toEqual([]); - expect( - await db.prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'branch_tips'").get(), - ).toBeDefined(); - expect( - (await db.prepare("SELECT id FROM migrations ORDER BY id").all<{ id: string }>()).map((row) => row.id), - ).toEqual(["001_initial.sql", "002_branch_tips.sql"]); - } finally { - await db.close(); - } - }); - it("persists session metadata through create, list, open, and fork", async () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite"); @@ -185,7 +164,8 @@ describe("SQLite migrations", () => { const sourceMetadata = await source.getMetadata(); expect(sourceMetadata.metadata).toEqual({ profile: "reviewer" }); expect((await repo.list({ cwd: root })).map((listed) => listed.metadata)).toEqual([{ profile: "reviewer" }]); - expect((await (await repo.open(sourceMetadata)).getMetadata()).metadata).toEqual({ profile: "reviewer" }); + const reopened = await repo.open(sourceMetadata); + expect((await reopened.getMetadata()).metadata).toEqual({ profile: "reviewer" }); const fork = await repo.fork(sourceMetadata, { cwd: root, id: "session-2" }); expect((await fork.getMetadata()).metadata).toEqual({ profile: "reviewer" }); const overridden = await repo.fork(sourceMetadata, { @@ -209,8 +189,8 @@ describe("SQLite migrations", () => { const db = await sqlite.open(databasePath); try { await db.exec(` -CREATE TRIGGER fail_fork_entry BEFORE INSERT ON session_entries -WHEN new.session_id = 'fork' AND new.entry_seq = 2 +CREATE TRIGGER fail_fork_entry BEFORE INSERT ON entries +WHEN new.session_id = 'fork' AND new.seq = 2 BEGIN SELECT RAISE(ABORT, 'fail fork'); END; @@ -228,14 +208,14 @@ END; await inspection.prepare("SELECT id FROM sessions WHERE id = ?").get<{ id: string }>("fork"), ).toBeUndefined(); expect( - await inspection.prepare("SELECT id FROM session_entries WHERE session_id = ?").all<{ id: string }>("fork"), + await inspection.prepare("SELECT id FROM entries WHERE session_id = ?").all<{ id: string }>("fork"), ).toEqual([]); } finally { await inspection.close(); } }); - it("materializes active leaf id in sessions transactionally", async () => { + it("materializes the main lane leaf transactionally", async () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); @@ -244,29 +224,18 @@ END; const session = await repo.create({ cwd: root, id: "session-1" }); const rootId = await session.appendMessage(createUserMessage("root")); const childId = await session.appendMessage(createAssistantMessage("child")); - await session.moveTo(rootId); + await moveSqliteMainLane(session, rootId); const db = await sqlite.open(databasePath); try { const row = await db - .prepare("SELECT active_leaf_id FROM sessions WHERE id = ?") - .get<{ active_leaf_id: string | null }>("session-1"); - expect(row?.active_leaf_id).toBe(rootId); - const latestBranchRow = await db - .prepare( - "SELECT branch_id, entry_id, entry_seq FROM branch_entries WHERE session_id = ? ORDER BY entry_seq DESC LIMIT 1", - ) - .get<{ branch_id: string; entry_id: string; entry_seq: number }>("session-1"); - const latestSessionEntry = await db - .prepare("SELECT id, type FROM session_entries WHERE session_id = ? ORDER BY entry_seq DESC LIMIT 1") - .get<{ id: string; type: string }>("session-1"); - expect(latestSessionEntry?.type).toBe("leaf"); - expect(latestBranchRow?.entry_id).toBe(latestSessionEntry?.id); - if (!latestBranchRow) throw new Error("Missing latest branch row"); - const branchTip = await db - .prepare("SELECT branch_id, tip_id FROM branch_tips WHERE session_id = ? AND branch_id = ?") - .get<{ branch_id: string; tip_id: string }>("session-1", latestBranchRow.branch_id); - expect(branchTip?.tip_id).toBe(latestSessionEntry?.id); + .prepare("SELECT leaf_id FROM lanes WHERE session_id = ? AND lane = ?") + .get<{ leaf_id: string | null }>("session-1", "main"); + expect(row?.leaf_id).toBe(rootId); + const latestLaneMove = await db + .prepare("SELECT lane, leaf_id FROM lane_moves WHERE session_id = ? ORDER BY seq DESC LIMIT 1") + .get<{ lane: string; leaf_id: string | null }>("session-1"); + expect(latestLaneMove).toEqual({ lane: "main", leaf_id: rootId }); } finally { await db.close(); } @@ -285,7 +254,7 @@ END; const session = await repo.create({ cwd: root, id: "session-1" }); const rootId = await session.appendMessage(createUserMessage("root")); const firstChildId = await session.appendMessage(createAssistantMessage("first child")); - await session.moveTo(rootId); + await moveSqliteMainLane(session, rootId); const secondChildId = await session.appendMessage(createAssistantMessage("second child")); const db = await sqlite.open(databasePath); @@ -318,43 +287,30 @@ END; const session = await repo.create({ cwd: root, id: "session-1" }); const rootId = await session.appendMessage(createUserMessage("root")); await session.appendMessage(createAssistantMessage("first child")); - await session.appendSessionName(" Reopened Session "); - await session.moveTo(rootId); + await appendSqliteSessionName(session, " Reopened Session "); + await moveSqliteMainLane(session, rootId); await session.appendMessage(createAssistantMessage("branched child")); const reopened = await repo.open(await session.getMetadata()); - expect(await reopened.getSessionName()).toBe("Reopened Session"); - expect((await reopened.buildContext()).messages.map((message) => message.role)).toEqual(["user", "assistant"]); - expect((await reopened.buildContext()).messages.at(-1)).toMatchObject({ + expect(await reopened.getName()).toBe("Reopened Session"); + expect((await buildSqliteContext(reopened)).messages.map((message) => message.role)).toEqual([ + "user", + "assistant", + ]); + expect((await buildSqliteContext(reopened)).messages.at(-1)).toMatchObject({ content: [{ type: "text", text: "branched child" }], }); }); - it("pages entries by entry_seq cursor", async () => { - const root = createTempDir(); - const databasePath = join(root, "sessions.sqlite"); - const env = new NodeExecutionEnv({ cwd: root }); - const repo = new SqliteSessionRepository({ env, sqlite: createNodeSqliteFactory(), databasePath }); - const session = await repo.create({ cwd: root, id: "session-1" }); - const ids = [ - await session.appendMessage(createUserMessage("one")), - await session.appendMessage(createAssistantMessage("two")), - await session.appendMessage(createUserMessage("three")), - ]; - - expect((await session.getEntries({ limit: 2 })).map((entry) => entry.id)).toEqual(ids.slice(0, 2)); - expect((await session.getEntries({ afterEntrySeq: 1, limit: 2 })).map((entry) => entry.id)).toEqual(ids.slice(1)); - }); - it("closes the database when create fails after openDatabase succeeds", async () => { const root = createTempDir(); const db = new CountingDatabase((sql) => { if (sql.startsWith("INSERT INTO sessions")) { - return new ThrowingStatement(async () => { + return new ThrowingStatement(() => { throw new Error("insert failed"); }); } - return new ThrowingStatement(async () => ({ changes: 1 })); + return new ThrowingStatement(() => ({ changes: 1 })); }); const sqlite: SqliteDatabaseFactory = { open: async () => db, @@ -372,9 +328,9 @@ END; const root = createTempDir(); const db = new CountingDatabase((sql) => { if (sql.includes("FROM sessions WHERE id = ?")) { - return new ThrowingStatement(async () => ({ changes: 0 })); + return new ThrowingStatement(() => ({ changes: 0 })); } - return new ThrowingStatement(async () => ({ changes: 1 })); + return new ThrowingStatement(() => ({ changes: 1 })); }); const sqlite: SqliteDatabaseFactory = { open: async () => db, @@ -383,7 +339,7 @@ END; const repo = new SqliteSessionRepository({ env, sqlite, databasePath: join(root, "sessions.sqlite") }); const metadata: SqliteSessionMetadata = { id: "missing", - createdAt: new Date().toISOString(), + createdAt: Date.now(), cwd: root, path: join(root, "sessions.sqlite"), }; @@ -404,7 +360,7 @@ END; const session = await repo.create({ cwd: root, id: "session-1" }); for (let i = 0; i < 10; i++) await session.appendMessage(createUserMessage(`message ${i}`)); - await session.getEntries(); + await getSqliteEntries(session); expect(counts).toEqual({ opens: 1, closes: 0 }); await repo[Symbol.asyncDispose](); expect(counts).toEqual({ opens: 1, closes: 1 }); @@ -427,7 +383,7 @@ END; expect(counts).toEqual({ opens: 1, closes: 1 }); }); - it("rejects a missing active leaf when opened", async () => { + it("rejects a missing lane leaf when listing lanes and opening", async () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); @@ -438,14 +394,20 @@ END; const db = await sqlite.open(databasePath); try { - await db.prepare("UPDATE sessions SET active_leaf_id = ? WHERE id = ?").run("missing", metadata.id); + await db + .prepare("UPDATE lanes SET leaf_id = ? WHERE session_id = ? AND lane = ?") + .run("missing", metadata.id, "main"); } finally { await db.close(); } + await expect(session.getLanes()).rejects.toMatchObject({ + code: "storage", + message: expect.stringContaining("Lane main points at missing entry missing"), + }); await expect(repo.open(metadata)).rejects.toMatchObject({ - code: "invalid_session", - message: "Entry missing not found", + code: "storage", + message: expect.stringContaining("Lane main points at missing entry missing"), }); }); @@ -462,64 +424,46 @@ END; const db = await sqlite.open(databasePath); try { await db - .prepare("UPDATE session_entries SET payload = ? WHERE session_id = ? AND id = ?") + .prepare("UPDATE entries SET payload = ? WHERE session_id = ? AND id = ?") .run("not json", metadata.id, entryId); } finally { await db.close(); } const reopened = await repo.open(metadata); - await expect(reopened.getEntries()).rejects.toMatchObject({ code: "invalid_entry" }); + await expect(getSqliteEntries(reopened)).rejects.toMatchObject({ code: "invalid_entry" }); }); it("does not publish connection state when an append transaction fails", async () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); const sqlite = createNodeSqliteFactory(); + const repo = new SqliteSessionRepository({ env, sqlite, databasePath }); + const session = await repo.create({ cwd: root, id: "session-1" }); const db = await sqlite.open(databasePath); - await applyMigrations(db); - const storage = await SqliteSessionConnection.create(db, databasePath, { - cwd: root, - sessionId: "session-1", - }); - await db.exec(` - CREATE TEMP TRIGGER fail_branch_tip_insert - BEFORE INSERT ON branch_tips - BEGIN - SELECT RAISE(ABORT, 'branch insert failed'); - END; - `); - - const rootEntry = { - type: "message" as const, - id: "root", - parentId: null, - timestamp: new Date().toISOString(), - message: createUserMessage("root"), - }; try { - await expect(storage.appendEntry(rootEntry)).rejects.toMatchObject({ code: "storage" }); - } finally { + await db.exec(` + CREATE TRIGGER fail_branch_tip_insert + BEFORE INSERT ON branch_tips + BEGIN + SELECT RAISE(ABORT, 'branch insert failed'); + END; + `); + await expect(session.appendMessage(createUserMessage("root"))).rejects.toThrow("branch insert failed"); + const lane = await db + .prepare("SELECT leaf_id FROM lanes WHERE session_id = ? AND lane = ?") + .get<{ leaf_id: string | null }>("session-1", "main"); + expect(lane?.leaf_id).toBeNull(); + expect(await db.prepare("SELECT id FROM entries WHERE session_id = ?").all("session-1")).toEqual([]); + expect(await session.getStats()).toMatchObject({ messageCount: 0 }); await db.exec("DROP TRIGGER fail_branch_tip_insert"); + } finally { + await db.close(); } - const sessionRow = await db - .prepare("SELECT active_leaf_id FROM sessions WHERE id = ?") - .get<{ active_leaf_id: string | null }>("session-1"); - expect(sessionRow?.active_leaf_id).toBeNull(); - expect(await storage.readEntries()).toEqual([]); - await expect( - storage.appendEntry({ - type: "leaf", - id: "leaf", - parentId: null, - timestamp: new Date().toISOString(), - targetId: rootEntry.id, - }), - ).rejects.toMatchObject({ code: "not_found" }); - expect(await storage.readEntries()).toEqual([]); - await storage.appendEntry(rootEntry); - expect(await storage.readEntries()).toEqual([rootEntry]); - await db.close(); + const entryId = await session.appendMessage(createUserMessage("root")); + expect((await getSqliteEntries(session)).map((entry) => entry.id)).toEqual([entryId]); + expect(await session.getStats()).toMatchObject({ messageCount: 1 }); }); it("materializes session summary fields transactionally", async () => { @@ -530,8 +474,6 @@ END; const repo = new SqliteSessionRepository({ env, sqlite, databasePath }); const session = await repo.create({ cwd: root, id: "session-1" }); const userId = await session.appendMessage(createUserMessage("one")); - await session.appendThinkingLevelChange("high"); - await session.appendModelChange("anthropic", "claude-sonnet-4-5"); const assistant = { ...createAssistantMessage("two"), provider: "anthropic", @@ -545,62 +487,117 @@ END; cost: { input: 0.1, output: 0.2, cacheRead: 0.03, cacheWrite: 0.04, total: 0.37 }, }, }; - await session.appendMessage(assistant); - await session.appendCompaction("summary", userId, 200, undefined, false, { + const assistantId = await session.appendMessage(assistant); + await session.appendRecord({ + type: "usage", + id: "assistant-usage", + lane: "main", + cause: "assistant", + runId: "run", + entryId: assistantId, + attempt: 1, + stopReason: "stop", + usage: assistant.usage, + }); + const compactionUsage = { input: 1, output: 2, cacheRead: 3, cacheWrite: 4, totalTokens: 10, cost: { input: 0.01, output: 0.02, cacheRead: 0.03, cacheWrite: 0.04, total: 0.1 }, + }; + const compactionId = await appendSqliteCompaction( + session, + "summary", + userId, + 200, + undefined, + false, + compactionUsage, + ); + await session.appendRecord({ + type: "usage", + id: "compaction-usage", + lane: "main", + cause: "compaction", + runId: "run", + entryId: compactionId, + attempt: 1, + stopReason: "stop", + usage: compactionUsage, }); - await session.moveTo(userId, { + const branchUsage = { + input: 5, + output: 6, + cacheRead: 7, + cacheWrite: 8, + totalTokens: 26, + cost: { input: 0.05, output: 0.06, cacheRead: 0.07, cacheWrite: 0.08, total: 0.26 }, + }; + const branchSummaryId = await moveSqliteMainLane(session, userId, { summary: "branch summary", - usage: { - input: 5, - output: 6, - cacheRead: 7, - cacheWrite: 8, - totalTokens: 26, - cost: { input: 0.05, output: 0.06, cacheRead: 0.07, cacheWrite: 0.08, total: 0.26 }, - }, + usage: branchUsage, + }); + if (!branchSummaryId) throw new Error("Expected branch summary"); + await session.appendRecord({ + type: "usage", + id: "branch-summary-usage", + lane: "main", + cause: "branch_summary", + runId: "run", + entryId: branchSummaryId, + attempt: 1, + stopReason: "stop", + usage: branchUsage, + }); + await appendSqliteSessionName(session, " My Session "); + await appendSqliteLabel(session, userId, "checkpoint"); + + expect(await session.getStats()).toMatchObject({ + messageCount: 2, + cachedTokens: 50, + uncachedTokens: 128, + totalTokens: 211, + costTotal: 0.73, }); - await session.appendSessionName(" My Session "); - await session.appendLabel(userId, "checkpoint"); const db = await sqlite.open(databasePath); try { - const row = await db.prepare("SELECT session_id, payload FROM session_materialized WHERE session_id = ?").get<{ - session_id: string; - payload: string; - }>("session-1"); - expect(row).toBeDefined(); - expect(row?.session_id).toBe("session-1"); - expect(JSON.parse(row?.payload ?? "null")).toMatchObject({ - name: "My Session", - messageCount: 2, - cachedTokens: 50, - uncachedTokens: 128, - totalTokens: 211, - costTotal: 0.73, - currentModel: { provider: "anthropic", modelId: "claude-sonnet-4-5" }, - currentThinkingLevel: "high", - }); - const entryRows = await db - .prepare( - "SELECT session_id, entry_seq, type, payload FROM entry_materialized WHERE session_id = ? ORDER BY entry_seq, type", - ) - .all<{ - session_id: string; - entry_seq: number; - type: string; - payload: string; - }>("session-1"); expect( - entryRows.some((entryRow) => entryRow.type === "label" && JSON.parse(entryRow.payload).targetId === userId), - ).toBe(true); - expect(entryRows.some((entryRow) => entryRow.type === "thinking")).toBe(false); - expect(entryRows.some((entryRow) => entryRow.type === "model")).toBe(false); + await db + .prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'session_materialized'") + .get(), + ).toBeUndefined(); + expect( + await db + .prepare("SELECT type, COUNT(*) AS count FROM records WHERE session_id = ? AND type = ? GROUP BY type") + .get<{ type: string; count: number }>("session-1", "usage"), + ).toEqual({ type: "usage", count: 3 }); + expect( + await db + .prepare( + `SELECT message_count, cached_tokens, uncached_tokens, total_tokens, cost_total + FROM session_stats + WHERE session_id = ?`, + ) + .get("session-1"), + ).toEqual({ + message_count: 2, + cached_tokens: 50, + uncached_tokens: 128, + total_tokens: 211, + cost_total: 0.73, + }); + const nameFact = await db + .prepare("SELECT value FROM facts WHERE session_id = ? AND kind = 'name' ORDER BY seq DESC LIMIT 1") + .get<{ value: string }>("session-1"); + expect(JSON.parse(nameFact?.value ?? "null")).toBe("My Session"); + const labelFact = await db + .prepare("SELECT key, value FROM facts WHERE session_id = ? AND kind = 'label' ORDER BY seq DESC LIMIT 1") + .get<{ key: string; value: string }>("session-1"); + expect(labelFact?.key).toBe(userId); + expect(JSON.parse(labelFact?.value ?? "null")).toBe("checkpoint"); } finally { await db.close(); } diff --git a/packages/agent/test/harness/sqlite-node.test.ts b/packages/agent/test/harness/sqlite-node.test.ts index 2fda00539..d7cc7c9e7 100644 --- a/packages/agent/test/harness/sqlite-node.test.ts +++ b/packages/agent/test/harness/sqlite-node.test.ts @@ -9,8 +9,8 @@ import { import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; import { JsonlSessionRepository } from "../../src/harness/session/jsonl-repo.ts"; import { createScanningSessionSearch } from "../../src/harness/session/search.ts"; -import type { SessionSearch, SessionSearchHit, SessionSearchOptions } from "../../src/harness/types.ts"; -import { createTempDir, createUserMessage } from "./session-test-utils.ts"; +import type { SessionSearchOptions } from "../../src/harness/types.ts"; +import { createTempDir, createUserMessage, getSqliteEntries } from "./session-test-utils.ts"; const ownedRepositories: AsyncDisposable[] = []; @@ -46,6 +46,42 @@ describe("JsonlSessionBackend with scanning search", () => { }); }); +describe("SqliteSessionRepository writer leases", () => { + it("shares one storage queue for repeated opens in the same repository", async () => { + const root = createTempDir(); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const databasePath = join(root, "sessions.sqlite"); + const { repository: repo } = createSqliteFixture({ env, sqlite, databasePath }); + const session = await repo.create({ cwd: root, id: "session" }); + const reopened = await repo.open(await session.getMetadata()); + + const [first, second] = await Promise.all([ + session.appendMessage(createUserMessage("first")), + reopened.appendMessage(createUserMessage("second")), + ]); + + expect((await getSqliteEntries(session)).map((entry) => entry.id)).toEqual([first, second]); + }); + + it("rejects a second repository while a session lease is active", async () => { + const root = createTempDir(); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const databasePath = join(root, "sessions.sqlite"); + const { repository: firstRepo } = createSqliteFixture({ env, sqlite, databasePath }); + const session = await firstRepo.create({ cwd: root, id: "session" }); + const metadata = await session.getMetadata(); + const secondRepo = new SqliteSessionRepository({ env, sqlite, databasePath }); + ownedRepositories.push(secondRepo); + + await expect(secondRepo.open(metadata)).rejects.toMatchObject({ code: "storage" }); + + await firstRepo.close(); + await expect(secondRepo.open(metadata)).resolves.toBeDefined(); + }); +}); + describe("SqliteSessionBackend with explicit SQLite FTS5 search", () => { it("uses SQLite FTS5 when composed with its search implementation", async () => { const root = createTempDir(); @@ -118,7 +154,7 @@ describe("SqliteSessionBackend with explicit SQLite FTS5 search", () => { } await expect(session.appendMessage(createUserMessage("must roll back"))).rejects.toThrow(); - await expect(session.getEntries()).resolves.toEqual([]); + await expect(getSqliteEntries(session)).resolves.toEqual([]); }); it("rolls back canonical deletion when co-located FTS cleanup fails", async () => { @@ -141,7 +177,7 @@ describe("SqliteSessionBackend with explicit SQLite FTS5 search", () => { await expect(repo.delete(metadata)).rejects.toThrow(); const reopened = await repo.open(metadata); - await expect(reopened.getEntries()).resolves.toHaveLength(1); + await expect(getSqliteEntries(reopened)).resolves.toHaveLength(1); }); it("initializes canonical storage when searched before the first session is created", async () => { @@ -168,8 +204,10 @@ describe("SqliteSessionRepository with custom search", () => { it("uses an independently supplied search implementation", async () => { const root = createTempDir(); const searches: SessionSearchOptions[] = []; - const search: SessionSearch = { - async search(options): Promise[]> { + const search: { + search(options: SessionSearchOptions): Promise<{ metadata: SqliteSessionMetadata; entryId: string }[]>; + } = { + async search(options) { searches.push(options); return []; }, diff --git a/packages/agent/vitest.config.ts b/packages/agent/vitest.config.ts index ed0ab064a..5ca2a36fa 100644 --- a/packages/agent/vitest.config.ts +++ b/packages/agent/vitest.config.ts @@ -4,6 +4,7 @@ import { defineConfig } from "vitest/config"; const aiSrcIndex = fileURLToPath(new URL("../ai/src/index.ts", import.meta.url)); const aiSrcCompat = fileURLToPath(new URL("../ai/src/compat.ts", import.meta.url)); const agentSrcIndex = fileURLToPath(new URL("./src/index.ts", import.meta.url)); +const agentSrcExperimental = fileURLToPath(new URL("./src/experimental.ts", import.meta.url)); export default defineConfig({ test: { @@ -16,6 +17,7 @@ export default defineConfig({ resolve: { alias: [ { find: /^@earendil-works\/pi-agent-core$/, replacement: agentSrcIndex }, + { find: /^@earendil-works\/pi-agent-core\/experimental$/, replacement: agentSrcExperimental }, { find: /^@earendil-works\/pi-ai$/, replacement: aiSrcIndex }, { find: /^@earendil-works\/pi-ai\/compat$/, replacement: aiSrcCompat }, ], diff --git a/packages/agent/vitest.harness.config.ts b/packages/agent/vitest.harness.config.ts index 8045bb1a6..be1e8f800 100644 --- a/packages/agent/vitest.harness.config.ts +++ b/packages/agent/vitest.harness.config.ts @@ -4,6 +4,7 @@ import { defineConfig } from "vitest/config"; const aiSrcIndex = fileURLToPath(new URL("../ai/src/index.ts", import.meta.url)); const aiSrcCompat = fileURLToPath(new URL("../ai/src/compat.ts", import.meta.url)); const agentSrcIndex = fileURLToPath(new URL("../agent/src/index.ts", import.meta.url)); +const agentSrcExperimental = fileURLToPath(new URL("../agent/src/experimental.ts", import.meta.url)); export default defineConfig({ test: { @@ -23,6 +24,7 @@ export default defineConfig({ resolve: { alias: [ { find: /^@earendil-works\/pi-agent-core$/, replacement: agentSrcIndex }, + { find: /^@earendil-works\/pi-agent-core\/experimental$/, replacement: agentSrcExperimental }, { find: /^@earendil-works\/pi-ai$/, replacement: aiSrcIndex }, { find: /^@earendil-works\/pi-ai\/compat$/, replacement: aiSrcCompat }, ], diff --git a/packages/storage/sqlite-node/src/index.ts b/packages/storage/sqlite-node/src/index.ts index 05f822e3b..98c7fdfa4 100644 --- a/packages/storage/sqlite-node/src/index.ts +++ b/packages/storage/sqlite-node/src/index.ts @@ -8,6 +8,10 @@ function isNamedParameters(value: unknown): value is Record; @@ -15,7 +19,7 @@ class NodeSqliteStatement implements SqliteStatement { this.statement = statement; } - async run(...params: unknown[]): Promise { + run(...params: unknown[]): SqliteRunResult { const [first, ...rest] = params; const result = isNamedParameters(first) ? this.statement.run(first, ...(rest as SQLInputValue[])) @@ -26,7 +30,7 @@ class NodeSqliteStatement implements SqliteStatement { }; } - async get(...params: unknown[]): Promise { + get(...params: unknown[]): TRow | undefined { const [first, ...rest] = params; return ( isNamedParameters(first) @@ -35,7 +39,7 @@ class NodeSqliteStatement implements SqliteStatement { ) as TRow | undefined; } - async all(...params: unknown[]): Promise { + all(...params: unknown[]): TRow[] { const [first, ...rest] = params; return ( isNamedParameters(first) @@ -52,7 +56,7 @@ class NodeSqliteDatabase implements SqliteDatabase { this.db = db; } - async exec(sql: string): Promise { + exec(sql: string): void { this.db.exec(sql); } @@ -60,10 +64,13 @@ class NodeSqliteDatabase implements SqliteDatabase { return new NodeSqliteStatement(this.db.prepare(sql)); } - async transaction(fn: () => Promise): Promise { - this.db.exec("BEGIN"); + transaction(fn: () => T): T { + this.db.exec("BEGIN IMMEDIATE"); try { - const result = await fn(); + const result = fn(); + if (isAsyncResult(result)) { + throw new TypeError("SQLite transaction callbacks must be synchronous"); + } this.db.exec("COMMIT"); return result; } catch (error) { @@ -76,7 +83,7 @@ class NodeSqliteDatabase implements SqliteDatabase { } } - async close(): Promise { + close(): void { this.db.close(); } } diff --git a/packages/storage/sqlite-node/src/sqlite/branch-cache.ts b/packages/storage/sqlite-node/src/sqlite/branch-cache.ts new file mode 100644 index 000000000..babb2c1a3 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/branch-cache.ts @@ -0,0 +1,105 @@ +import { SessionError } from "@earendil-works/pi-agent-core/experimental"; +import { uuidv7 } from "@earendil-works/pi-ai"; +import { + copyBranchEntriesThroughSeq, + deleteBranchEntries, + insertBranchEntriesForPath, + insertBranchEntry, + readBranchContainingEntry, +} from "./storage/branch-entries.ts"; + +import { deleteBranchTips, insertBranchTip, readBranchTipBranchId, updateBranchTip } from "./storage/branch-tips.ts"; +import type { SqliteDatabase } from "./types.ts"; + +export function deleteBranchCache(db: SqliteDatabase, sessionId: string) { + deleteBranchTips(db, sessionId); + deleteBranchEntries(db, sessionId); +} + +export function rebuildBranchCache(db: SqliteDatabase, sessionId: string) { + const tips = db + .prepare( + `SELECT leaf.id + FROM entries AS leaf + WHERE leaf.session_id = ? + AND NOT EXISTS ( + SELECT 1 FROM entries AS child WHERE child.session_id = leaf.session_id AND child.parent_id = leaf.id + ) + ORDER BY leaf.seq`, + ) + .all<{ id: string }>(sessionId); + deleteBranchCache(db, sessionId); + for (const tip of tips) buildCachedBranch(db, sessionId, tip.id); +} + +export function buildCachedBranch(db: SqliteDatabase, sessionId: string, leafId: string) { + db.exec("SAVEPOINT build_branch_cache"); + try { + const branchId = uuidv7(); + insertBranchEntriesForPath(db, sessionId, branchId, leafId); + insertBranchTip(db, sessionId, leafId, branchId); + db.exec("RELEASE SAVEPOINT build_branch_cache"); + } catch (error) { + try { + db.exec("ROLLBACK TO SAVEPOINT build_branch_cache"); + db.exec("RELEASE SAVEPOINT build_branch_cache"); + } catch { + // Preserve the original build failure. + } + if (error instanceof SessionError) throw error; + throw new SessionError( + "storage", + `Failed to build SQLite branch cache at entry ${leafId}`, + error instanceof Error ? error : undefined, + ); + } +} + +function extendBranch( + db: SqliteDatabase, + sessionId: string, + branchId: string, + parentId: string, + entryId: string, + entrySeq: number, + entryType: string, + customType: string | null, +) { + insertBranchEntry(db, sessionId, branchId, entryId, entrySeq, entryType, customType); + if (!updateBranchTip(db, sessionId, branchId, parentId, entryId)) { + throw new SessionError("invalid_entry", `Branch tip ${parentId} changed during append`); + } +} + +export function appendEntryToBranchCache( + db: SqliteDatabase, + sessionId: string, + entryId: string, + entrySeq: number, + entryType: string, + customType: string | null, + parentId: string | null, +) { + if (parentId === null) { + const branchId = uuidv7(); + insertBranchEntry(db, sessionId, branchId, entryId, entrySeq, entryType, customType); + insertBranchTip(db, sessionId, entryId, branchId); + return; + } + + const tipBranchId = readBranchTipBranchId(db, sessionId, parentId); + if (tipBranchId !== undefined) { + extendBranch(db, sessionId, tipBranchId, parentId, entryId, entrySeq, entryType, customType); + return; + } + + const source = readBranchContainingEntry(db, sessionId, parentId); + if (!source) { + throw new SessionError("invalid_entry", `Branch cache has no branch containing parent entry ${parentId}`); + } + + const branchId = uuidv7(); + copyBranchEntriesThroughSeq(db, sessionId, branchId, source.branchId, source.entrySeq); + insertBranchEntry(db, sessionId, branchId, entryId, entrySeq, entryType, customType); + insertBranchTip(db, sessionId, entryId, branchId); +} diff --git a/packages/storage/sqlite-node/src/sqlite/index.ts b/packages/storage/sqlite-node/src/sqlite/index.ts index 1c2b85199..8b5fbbaa0 100644 --- a/packages/storage/sqlite-node/src/sqlite/index.ts +++ b/packages/storage/sqlite-node/src/sqlite/index.ts @@ -1,7 +1,16 @@ export * from "./migrations.ts"; export { + type SqliteSessionCreateOptions, + type SqliteSessionListOptions, + type SqliteSessionMetadata, SqliteSessionRepository, type SqliteSessionRepositoryOptions, + type SqliteWriterLeaseOptions, } from "./repo.ts"; export * from "./search-backend.ts"; -export * from "./types.ts"; +export type { + SqliteDatabase, + SqliteDatabaseFactory, + SqliteRunResult, + SqliteStatement, +} from "./types.ts"; diff --git a/packages/storage/sqlite-node/src/sqlite/migrations.ts b/packages/storage/sqlite-node/src/sqlite/migrations.ts index 08414c0d0..a0debe3f0 100644 --- a/packages/storage/sqlite-node/src/sqlite/migrations.ts +++ b/packages/storage/sqlite-node/src/sqlite/migrations.ts @@ -19,16 +19,11 @@ export async function loadMigrations(): Promise { order: 1, sql: await loadMigrationSql("./migrations/001_initial.sql"), }, - { - id: "002_branch_tips.sql", - order: 2, - sql: await loadMigrationSql("./migrations/002_branch_tips.sql"), - }, ]; } -async function ensureMigrationsTable(db: SqliteDatabase): Promise { - await db.exec(` +function ensureMigrationsTable(db: SqliteDatabase): void { + db.exec(` CREATE TABLE IF NOT EXISTS migrations ( id TEXT PRIMARY KEY, applied_at TEXT NOT NULL @@ -37,18 +32,19 @@ CREATE TABLE IF NOT EXISTS migrations ( } export async function applyMigrations(db: SqliteDatabase): Promise { - await ensureMigrationsTable(db); + ensureMigrationsTable(db); const migrations = await loadMigrations(); - const appliedRows = await db.prepare("SELECT id FROM migrations ORDER BY applied_at, id").all<{ id: string }>(); + const appliedRows = db.prepare("SELECT id FROM migrations ORDER BY applied_at, id").all<{ id: string }>(); const applied = new Set(appliedRows.map((row) => row.id)); for (const migration of migrations) { if (applied.has(migration.id)) continue; - await db.transaction(async () => { - await db.exec(migration.sql); - await db - .prepare("INSERT INTO migrations (id, applied_at) VALUES (?, ?)") - .run(migration.id, new Date().toISOString()); + db.transaction(() => { + db.exec(migration.sql); + db.prepare("INSERT INTO migrations (id, applied_at) VALUES (?, ?)").run( + migration.id, + new Date().toISOString(), + ); }); applied.add(migration.id); } diff --git a/packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql b/packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql index 6e6f397d9..3937389a4 100644 --- a/packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql +++ b/packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql @@ -3,57 +3,122 @@ CREATE TABLE IF NOT EXISTS sessions ( created_at TEXT NOT NULL, cwd TEXT NOT NULL, parent_session_id TEXT NULL, - metadata TEXT NULL, - active_leaf_id TEXT NULL + metadata TEXT NULL ) WITHOUT ROWID; CREATE INDEX IF NOT EXISTS idx_sessions_created_at ON sessions(created_at DESC); CREATE INDEX IF NOT EXISTS idx_sessions_cwd ON sessions(cwd); CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_session_id); -CREATE TABLE IF NOT EXISTS session_entries ( +CREATE TABLE IF NOT EXISTS entries ( session_id TEXT NOT NULL, + seq INTEGER NOT NULL, id TEXT NOT NULL, - entry_seq INTEGER NOT NULL, parent_id TEXT NULL, type TEXT NOT NULL, timestamp TEXT NOT NULL, payload TEXT NOT NULL, - PRIMARY KEY (session_id, id) + PRIMARY KEY (session_id, id), + UNIQUE (session_id, seq) ); -CREATE UNIQUE INDEX IF NOT EXISTS idx_session_entries_session_seq ON session_entries(session_id, entry_seq); -CREATE INDEX IF NOT EXISTS idx_session_entries_session_parent ON session_entries(session_id, parent_id); -CREATE INDEX IF NOT EXISTS idx_session_entries_session_type ON session_entries(session_id, type); +CREATE INDEX IF NOT EXISTS idx_entries_session_seq ON entries(session_id, seq); +CREATE INDEX IF NOT EXISTS idx_entries_session_parent ON entries(session_id, parent_id); +CREATE INDEX IF NOT EXISTS idx_entries_session_type_seq ON entries(session_id, type, seq); CREATE TABLE IF NOT EXISTS session_sequences ( session_id TEXT PRIMARY KEY, next_seq INTEGER NOT NULL ) WITHOUT ROWID; +CREATE TABLE IF NOT EXISTS session_stats ( + session_id TEXT PRIMARY KEY, + message_count INTEGER NOT NULL, + cached_tokens REAL NOT NULL, + uncached_tokens REAL NOT NULL, + total_tokens REAL NOT NULL, + cost_total REAL NOT NULL +) WITHOUT ROWID; + +-- Derived branch cache. Parent links in entries remain canonical; this cache +-- exists only to make branch scans cheap. CREATE TABLE IF NOT EXISTS branch_entries ( session_id TEXT NOT NULL, branch_id TEXT NOT NULL, entry_id TEXT NOT NULL, entry_seq INTEGER NOT NULL, + entry_type TEXT NULL, + custom_type TEXT NULL, PRIMARY KEY (session_id, branch_id, entry_id) ) WITHOUT ROWID; -CREATE INDEX IF NOT EXISTS idx_branch_entries_session_branch ON branch_entries(session_id, branch_id); CREATE INDEX IF NOT EXISTS idx_branch_entries_session_branch_seq ON branch_entries(session_id, branch_id, entry_seq); CREATE INDEX IF NOT EXISTS idx_branch_entries_session_entry ON branch_entries(session_id, entry_id); +CREATE INDEX IF NOT EXISTS idx_branch_entries_session_branch_type_seq ON branch_entries(session_id, branch_id, entry_type, entry_seq); +CREATE INDEX IF NOT EXISTS idx_branch_entries_session_branch_custom_seq ON branch_entries(session_id, branch_id, custom_type, entry_seq); -CREATE TABLE IF NOT EXISTS session_materialized ( - session_id TEXT PRIMARY KEY, - payload TEXT NOT NULL -) WITHOUT ROWID; - -CREATE TABLE IF NOT EXISTS entry_materialized ( +CREATE TABLE IF NOT EXISTS lanes ( session_id TEXT NOT NULL, - entry_seq INTEGER NOT NULL, - type TEXT NOT NULL, - payload TEXT NOT NULL, - PRIMARY KEY (session_id, entry_seq, type) + lane TEXT NOT NULL, + leaf_id TEXT NULL, + PRIMARY KEY (session_id, lane) ) WITHOUT ROWID; -CREATE INDEX IF NOT EXISTS idx_entry_materialized_session_type_seq ON entry_materialized(session_id, type, entry_seq); +CREATE INDEX IF NOT EXISTS idx_lanes_session_leaf ON lanes(session_id, leaf_id); + +CREATE TABLE IF NOT EXISTS records ( + session_id TEXT NOT NULL, + seq INTEGER NOT NULL, + id TEXT NOT NULL, + lane TEXT NOT NULL, + run_id TEXT NULL, + type TEXT NOT NULL, + op_kind TEXT NULL, + timestamp TEXT NOT NULL, + payload TEXT NOT NULL, + PRIMARY KEY (session_id, id), + UNIQUE (session_id, seq) +) WITHOUT ROWID; + +CREATE INDEX IF NOT EXISTS idx_records_session_seq ON records(session_id, seq); +CREATE INDEX IF NOT EXISTS idx_records_session_lane_type_seq ON records(session_id, lane, type, seq); +CREATE INDEX IF NOT EXISTS idx_records_session_lane_type_op_kind_seq ON records(session_id, lane, type, op_kind, seq); +CREATE INDEX IF NOT EXISTS idx_records_session_run_id_seq ON records(session_id, run_id, seq); + +CREATE TABLE IF NOT EXISTS lane_moves ( + session_id TEXT NOT NULL, + seq INTEGER NOT NULL, + lane TEXT NOT NULL, + leaf_id TEXT NULL, + PRIMARY KEY (session_id, seq) +) WITHOUT ROWID; + +CREATE INDEX IF NOT EXISTS idx_lane_moves_session_lane_seq ON lane_moves(session_id, lane, seq); + +CREATE TABLE IF NOT EXISTS facts ( + session_id TEXT NOT NULL, + seq INTEGER NOT NULL, + kind TEXT NOT NULL, + key TEXT NULL, + value TEXT NULL, + PRIMARY KEY (session_id, seq) +) WITHOUT ROWID; + +CREATE INDEX IF NOT EXISTS idx_facts_session_kind_key_seq ON facts(session_id, kind, key, seq); + +CREATE TABLE IF NOT EXISTS branch_tips ( + session_id TEXT NOT NULL, + tip_id TEXT NOT NULL, + branch_id TEXT NOT NULL, + PRIMARY KEY (session_id, tip_id), + UNIQUE (session_id, branch_id) +) WITHOUT ROWID; + +-- Per-session writer claim. The fence prevents an expired owner from writing +-- after a new owner takes over the session. +CREATE TABLE IF NOT EXISTS leases ( + session_id TEXT PRIMARY KEY, + owner_id TEXT NOT NULL, + fence INTEGER NOT NULL, + expires_at_ms INTEGER NOT NULL +) WITHOUT ROWID; diff --git a/packages/storage/sqlite-node/src/sqlite/migrations/002_branch_tips.sql b/packages/storage/sqlite-node/src/sqlite/migrations/002_branch_tips.sql deleted file mode 100644 index aba7c9397..000000000 --- a/packages/storage/sqlite-node/src/sqlite/migrations/002_branch_tips.sql +++ /dev/null @@ -1,12 +0,0 @@ -CREATE TABLE IF NOT EXISTS branch_tips ( - session_id TEXT NOT NULL, - tip_id TEXT NOT NULL, - branch_id TEXT NOT NULL, - PRIMARY KEY (session_id, tip_id), - UNIQUE (session_id, branch_id) -) WITHOUT ROWID; - -DELETE FROM branch_tips; -DELETE FROM branch_entries; - -DROP INDEX IF EXISTS idx_branch_entries_session_branch; diff --git a/packages/storage/sqlite-node/src/sqlite/repo.ts b/packages/storage/sqlite-node/src/sqlite/repo.ts index f38e48dc5..abf78d2e2 100644 --- a/packages/storage/sqlite-node/src/sqlite/repo.ts +++ b/packages/storage/sqlite-node/src/sqlite/repo.ts @@ -1,51 +1,142 @@ -import type { - SessionForkOptions, - SessionForkSelection, - SessionRepository, - SessionStorage, - SessionTreeEntry, -} from "@earendil-works/pi-agent-core"; +import type { FileError, FileSystem, Result } from "@earendil-works/pi-agent-core"; import { - createSession, - createSessionForkSelection, - createSessionId, - getFileSystemResultOrThrow, - readSessionEntriesForFork, - type Session, - type SessionContextBuildOptions, + type BranchBounds, + type Entry, + type EntryQuery, + type ForkOptions, + type LaneRecord, + type LogItem, + type LogOptions, + type NewRecord, + type ProvisionedEntry, + type RecordQuery, + Session, + type SessionCreateOptions, SessionError, -} from "@earendil-works/pi-agent-core"; + type SessionMetadata, + type SessionRepo as SessionRepository, + type SessionStats, + type SessionStorage, +} from "@earendil-works/pi-agent-core/experimental"; +import { uuidv7 } from "@earendil-works/pi-ai"; +import { appendEntryToBranchCache, buildCachedBranch, deleteBranchCache, rebuildBranchCache } from "./branch-cache.ts"; import { applyMigrations } from "./migrations.ts"; -import { SqliteSessionConnection } from "./storage/index.ts"; -import { rowToMetadata, type SessionRow } from "./storage/sessions.ts"; -import type { - SqliteDatabase, - SqliteDatabaseFactory, - SqliteSessionCreateOptions, - SqliteSessionListOptions, - SqliteSessionMetadata, - SqliteSessionRepositoryEnv, -} from "./types.ts"; +import { type CachedBranchEntryRow, queryCachedBranchRows, readCachedBranch } from "./storage/branch-entries.ts"; +import { readBranchTipIds } from "./storage/branch-tips.ts"; +import { + deleteEntryRows, + type EntryRow, + entryPayload, + idExistsInEntries, + insertEntryRow, + readEntryRow, + readEntryRows, +} from "./storage/entries.ts"; +import { appendFact, deleteFactRows, readFactRows, readLatestFact, readLatestLabelFacts } from "./storage/facts.ts"; +import { + createInitialLane, + deleteLaneRows, + createLane as insertLane, + readLane, + readLaneHead, + readLaneMoveRows, + readLanes, + setLaneLeaf, + moveLane as updateLane, +} from "./storage/lanes.ts"; +import { + acquireSessionLease, + deleteSessionLease, + releaseSessionLease, + renewSessionLease, + type SessionLease, +} from "./storage/leases.ts"; +import { appendRecordRow, deleteRecordRows, idExistsInRecords, readRecordRows } from "./storage/records.ts"; +import { + advanceSequence, + createSequence, + deleteSequence, + getNextSequence, + setNextSequence, +} from "./storage/session-sequences.ts"; +import { + addUsageToStats, + createStats, + deleteStats, + incrementMessageCount, + readStats, +} from "./storage/session-stats.ts"; +import { + deleteSessionRow, + insertSessionRow, + readSessionRow, + readSessionRows, + rowToMetadata, + type SessionRow, + sessionExists, +} from "./storage/sessions.ts"; +import type { SqliteDatabase, SqliteDatabaseFactory } from "./types.ts"; -function getParentPath(path: string): string { - const normalized = path.replace(/[\\/]+$/, ""); - const lastSlash = Math.max(normalized.lastIndexOf("/"), normalized.lastIndexOf("\\")); - if (lastSlash < 0) return "."; - if (lastSlash === 0) return normalized.slice(0, 1); - return normalized.slice(0, lastSlash); +export interface SqliteSessionMetadata extends SessionMetadata { + cwd: string; + path: string; + metadata?: Record; } -async function configureSqliteDatabase(db: SqliteDatabase): Promise { - await db.exec("PRAGMA journal_mode=WAL"); - await db.exec("PRAGMA synchronous=FULL"); - await db.exec("PRAGMA busy_timeout=5000"); +export interface SqliteSessionCreateOptions extends SessionCreateOptions { + cwd: string; + metadata?: Record; } -export type SqliteSessionBackendOptions = { +export interface SqliteSessionListOptions { + cwd?: string; +} + +export interface SqliteWriterLeaseOptions { + /** Time without a successful heartbeat before another writer may take over. Default: 30 seconds. */ + ttlMs?: number; + /** Idle heartbeat cadence. Default: 10 seconds. Must be less than ttlMs. */ + heartbeatIntervalMs?: number; +} + +export type SqliteSessionRepositoryEnv = Pick; + +export interface SqliteSessionRepositoryOptions { env: SqliteSessionRepositoryEnv; sqlite: SqliteDatabaseFactory; databasePath: string; -}; + writerLease?: SqliteWriterLeaseOptions; +} + +interface ResolvedWriterLeaseOptions { + ttlMs: number; + heartbeatIntervalMs: number; +} + +function resolveWriterLeaseOptions(options: SqliteWriterLeaseOptions | undefined): ResolvedWriterLeaseOptions { + const ttlMs = options?.ttlMs ?? 30_000; + const heartbeatIntervalMs = options?.heartbeatIntervalMs ?? 10_000; + if (!Number.isSafeInteger(ttlMs) || ttlMs <= 0) throw new RangeError("writerLease.ttlMs must be positive"); + if (!Number.isSafeInteger(heartbeatIntervalMs) || heartbeatIntervalMs <= 0 || heartbeatIntervalMs >= ttlMs) { + throw new RangeError("writerLease.heartbeatIntervalMs must be positive and less than ttlMs"); + } + return { ttlMs, heartbeatIntervalMs }; +} + +function activeWriterError(sessionId: string): SessionError { + return new SessionError("storage", `SQLite session ${sessionId} already has an active writer`); +} + +function lostWriterError(sessionId: string): SessionError { + return new SessionError("storage", `SQLite session ${sessionId} writer lease was lost`); +} + +function acquireWriterLease(db: SqliteDatabase, sessionId: string, options: ResolvedWriterLeaseOptions): SessionLease { + const now = Date.now(); + const lease = acquireSessionLease(db, sessionId, uuidv7(), now, now + options.ttlMs); + if (!lease) throw activeWriterError(sessionId); + return lease; +} class SerialOperationQueue { private tail: Promise = Promise.resolve(); @@ -64,191 +155,747 @@ class SerialOperationQueue { } } -class SqliteSessionBackend { - private readonly env: SqliteSessionRepositoryEnv; - private readonly sqlite: SqliteDatabaseFactory; - private readonly databasePathInput: string; - private databasePath: string | undefined; - private databasePromise: Promise | undefined; - private database: SqliteDatabase | undefined; - private disposed = false; - private disposePromise: Promise | undefined; - private readonly operations = new SerialOperationQueue(); - private readonly writers = new Map(); +function resultOrThrow(result: Result, message: string): T { + if (!result.ok) { + const code = result.error.code === "not_found" ? "not_found" : "storage"; + throw new SessionError(code, `${message}: ${result.error.message}`, result.error); + } + return result.value; +} - constructor(options: SqliteSessionBackendOptions) { - this.env = options.env; - this.sqlite = options.sqlite; - this.databasePathInput = options.databasePath; +function getParentPath(path: string): string { + const normalized = path.replace(/[\\/]+$/, ""); + const lastSlash = Math.max(normalized.lastIndexOf("/"), normalized.lastIndexOf("\\")); + if (lastSlash < 0) return "."; + if (lastSlash === 0) return normalized.slice(0, 1); + return normalized.slice(0, lastSlash); +} + +function configureSqliteDatabase(db: SqliteDatabase): void { + db.exec("PRAGMA journal_mode=WAL"); + db.exec("PRAGMA synchronous=FULL"); + db.exec("PRAGMA busy_timeout=5000"); +} + +function timestampToText(timestamp: number): string { + return new Date(timestamp).toISOString(); +} + +function timestampFromText(timestamp: string): number { + return Date.parse(timestamp); +} + +function entryRowFromCached(row: CachedBranchEntryRow): EntryRow { + return { ...row, seq: row.entry_seq, type: row.type as Entry["type"] }; +} + +function readObjectPayload(row: EntryRow): Record { + const payload = JSON.parse(row.payload) as unknown; + if (typeof payload !== "object" || payload === null || Array.isArray(payload)) { + throw new Error("Payload is not an object"); + } + return payload as Record; +} + +function decodeEntry(row: EntryRow): Entry { + try { + const payload = readObjectPayload(row); + const timestamp = timestampFromText(row.timestamp); + if (!Number.isFinite(timestamp)) throw new Error(`Invalid timestamp ${row.timestamp}`); + const base = { id: row.id, seq: row.seq, parentId: row.parent_id, timestamp }; + switch (row.type) { + case "message": + if (typeof payload.message !== "object" || payload.message === null) throw new Error("Missing message"); + return { + ...base, + type: "message", + message: payload.message as Extract["message"], + ...(payload.terminate === true ? { terminate: true as const } : {}), + }; + case "model_change": + if (typeof payload.provider !== "string" || typeof payload.modelId !== "string") { + throw new Error("Invalid model_change payload"); + } + return { ...base, type: "model_change", provider: payload.provider, modelId: payload.modelId }; + case "thinking_level_change": + if (typeof payload.thinkingLevel !== "string") throw new Error("Invalid thinking_level_change payload"); + return { ...base, type: "thinking_level_change", thinkingLevel: payload.thinkingLevel }; + case "active_tools_change": + if (!Array.isArray(payload.activeToolNames)) throw new Error("Invalid active_tools_change payload"); + if (payload.activeToolNames.some((value) => typeof value !== "string")) { + throw new Error("Invalid active_tools_change payload"); + } + return { ...base, type: "active_tools_change", activeToolNames: payload.activeToolNames }; + case "compaction": + if ( + typeof payload.summary !== "string" || + !Array.isArray(payload.retainedTail) || + typeof payload.tokensBefore !== "number" + ) { + throw new Error("Invalid compaction payload"); + } + return { + ...base, + type: "compaction", + summary: payload.summary, + retainedTail: payload.retainedTail as Extract["retainedTail"], + tokensBefore: payload.tokensBefore, + ...(Object.hasOwn(payload, "details") ? { details: payload.details } : {}), + ...(Object.hasOwn(payload, "usage") + ? { usage: payload.usage as Extract["usage"] } + : {}), + }; + case "branch_summary": + if (typeof payload.fromId !== "string" || typeof payload.summary !== "string") { + throw new Error("Invalid branch_summary payload"); + } + return { + ...base, + type: "branch_summary", + fromId: payload.fromId, + summary: payload.summary, + ...(Object.hasOwn(payload, "details") ? { details: payload.details } : {}), + ...(Object.hasOwn(payload, "usage") + ? { usage: payload.usage as Extract["usage"] } + : {}), + }; + case "custom": + if (typeof payload.customType !== "string") throw new Error("Invalid custom payload"); + return { + ...base, + type: "custom", + customType: payload.customType, + ...(Object.hasOwn(payload, "data") ? { data: payload.data } : {}), + }; + } + } catch (error) { + throw new SessionError( + "invalid_entry", + `Invalid SQLite session entry ${row.id}: failed to decode entry ${row.id}`, + error instanceof Error ? error : undefined, + ); + } +} + +function recordRunId(record: NewRecord): string | undefined { + return record.type === "operation_started" ? record.id : "runId" in record ? record.runId : undefined; +} + +function recordOpKind(record: NewRecord): string | undefined { + return record.type === "operation_started" ? record.intent.kind : undefined; +} + +function decodeRecord(row: { seq: number; timestamp: string; payload: string }): LaneRecord { + try { + const timestamp = timestampFromText(row.timestamp); + if (!Number.isFinite(timestamp)) throw new Error(`Invalid timestamp ${row.timestamp}`); + return { + ...(JSON.parse(row.payload) as object), + seq: row.seq, + timestamp, + } as LaneRecord; + } catch (error) { + throw new SessionError( + "storage", + `Invalid SQLite session record at sequence ${row.seq}: failed to decode payload`, + error instanceof Error ? error : undefined, + ); + } +} + +function validateCachedBranchRows(rows: readonly CachedBranchEntryRow[], query: BranchBounds): void { + if (rows.length === 0) return; + const path = [...rows].sort((left, right) => left.entry_seq - right.entry_seq); + if (query.stopAtId === undefined && query.stopAtType === undefined && path[0]?.parent_id !== null) { + throw new SessionError("invalid_entry", `Entry ${path[0]?.parent_id} not found`); + } + for (let index = 1; index < path.length; index++) { + const previous = path[index - 1]!; + const current = path[index]!; + if (current.parent_id !== previous.id) { + throw new SessionError("invalid_entry", `Entry ${current.parent_id} not found`); + } + } +} + +function matchesEntryQuery(entry: Entry, query: EntryQuery): boolean { + return ( + (query.type === undefined || entry.type === query.type) && + (query.customType === undefined || (entry.type === "custom" && entry.customType === query.customType)) && + (query.cursor === undefined || + (query.order === "oldestFirst" ? entry.seq > query.cursor.afterSeq : entry.seq < query.cursor.afterSeq)) + ); +} + +function assertUnusedId(db: SqliteDatabase, sessionId: string, id: string): void { + if (idExistsInEntries(db, sessionId, id) || idExistsInRecords(db, sessionId, id)) { + throw new SessionError("already_exists", `ID already exists: ${id}`); + } +} + +function requireSessionRow(db: SqliteDatabase, sessionId: string): SessionRow { + const row = readSessionRow(db, sessionId); + if (!row) throw new SessionError("not_found", `Session not found: ${sessionId}`); + return row; +} + +class SqliteSessionStorage implements SessionStorage { + private readonly db: SqliteDatabase; + private readonly metadata: SqliteSessionMetadata; + private readonly lease: SessionLease; + private readonly leaseOptions: ResolvedWriterLeaseOptions; + private readonly onRelease: () => void; + private readonly operations = new SerialOperationQueue(); + private heartbeatTimer: ReturnType | undefined; + private leaseError: SessionError | undefined; + private closing = false; + private releasePromise: Promise | undefined; + + constructor( + db: SqliteDatabase, + metadata: SqliteSessionMetadata, + lease: SessionLease, + leaseOptions: ResolvedWriterLeaseOptions, + onRelease: () => void, + ) { + this.db = db; + this.metadata = metadata; + this.lease = lease; + this.leaseOptions = leaseOptions; + this.onRelease = onRelease; + this.scheduleHeartbeat(); } - create(options: SqliteSessionCreateOptions): Promise> { - this.assertOpen(); + async release(): Promise { + this.releasePromise ??= this.finishRelease(); + await this.releasePromise; + } + + private async finishRelease(): Promise { + this.closing = true; + if (this.heartbeatTimer !== undefined) clearTimeout(this.heartbeatTimer); + try { + await this.operations.enqueue(() => + this.db.transaction(() => releaseSessionLease(this.db, this.metadata.id, this.lease)), + ); + } finally { + this.onRelease(); + } + } + + private enqueueWrite(operation: () => T): Promise { + if (this.closing) + return Promise.reject(new SessionError("storage", `SQLite session ${this.metadata.id} is closed`)); + return this.operations.enqueue(() => { + if (this.leaseError) throw this.leaseError; + return this.db.transaction(() => { + const now = Date.now(); + if (!renewSessionLease(this.db, this.metadata.id, this.lease, now, now + this.leaseOptions.ttlMs)) { + this.leaseError = lostWriterError(this.metadata.id); + if (this.heartbeatTimer !== undefined) clearTimeout(this.heartbeatTimer); + throw this.leaseError; + } + return operation(); + }); + }); + } + + private scheduleHeartbeat(): void { + if (this.closing || this.leaseError) return; + this.heartbeatTimer = setTimeout(async () => { + this.heartbeatTimer = undefined; + try { + await this.operations.enqueue(() => { + if (this.closing || this.leaseError) return; + this.db.transaction(() => { + const now = Date.now(); + if (!renewSessionLease(this.db, this.metadata.id, this.lease, now, now + this.leaseOptions.ttlMs)) { + this.leaseError = lostWriterError(this.metadata.id); + } + }); + }); + } catch { + // A transient heartbeat failure is retried. Every write still verifies ownership transactionally. + } finally { + this.scheduleHeartbeat(); + } + }, this.leaseOptions.heartbeatIntervalMs); + this.heartbeatTimer.unref(); + } + + async getMetadata(): Promise { + return structuredClone(this.metadata); + } + + isForSession(sessionId: string): boolean { + return this.metadata.id === sessionId; + } + + async getLanes(): Promise<{ lane: string; leafId: string | null }[]> { + return readLanes(this.db, this.metadata.id).map((row) => ({ lane: row.lane, leafId: row.leaf_id })); + } + + async createLane(lane: string, at: string | null): Promise { + return this.enqueueWrite(() => { + if (readLane(this.db, this.metadata.id, lane)) { + throw new SessionError("already_exists", `Lane already exists: ${lane}`); + } + if (at !== null && !readEntryRow(this.db, this.metadata.id, at)) { + throw new SessionError("not_found", `Entry not found: ${at}`); + } + const seq = getNextSequence(this.db, this.metadata.id); + insertLane(this.db, this.metadata.id, seq, lane, at); + advanceSequence(this.db, this.metadata.id, seq); + }); + } + + async moveLane(lane: string, to: string | null): Promise { + return this.enqueueWrite(() => { + if (!readLane(this.db, this.metadata.id, lane)) + throw new SessionError("invalid_lane", `Lane not found: ${lane}`); + if (to !== null && !readEntryRow(this.db, this.metadata.id, to)) { + throw new SessionError("not_found", `Entry not found: ${to}`); + } + const seq = getNextSequence(this.db, this.metadata.id); + updateLane(this.db, this.metadata.id, seq, lane, to); + advanceSequence(this.db, this.metadata.id, seq); + }); + } + + async appendEntry(entry: ProvisionedEntry, lane: string): Promise { + return this.enqueueWrite(() => { + const parentId = readLaneHead(this.db, this.metadata.id, lane).leafId; + assertUnusedId(this.db, this.metadata.id, entry.id); + const seq = getNextSequence(this.db, this.metadata.id); + const committed = { ...entry, parentId, seq, timestamp: Date.now() } as Entry; + insertEntryRow(this.db, this.metadata.id, { + seq, + id: committed.id, + parentId: committed.parentId, + type: committed.type, + timestamp: timestampToText(committed.timestamp), + payload: JSON.stringify(entryPayload(committed)), + }); + setLaneLeaf(this.db, this.metadata.id, lane, committed.id); + appendEntryToBranchCache( + this.db, + this.metadata.id, + committed.id, + seq, + committed.type, + committed.type === "custom" ? committed.customType : null, + committed.parentId, + ); + if (committed.type === "message") incrementMessageCount(this.db, this.metadata.id); + advanceSequence(this.db, this.metadata.id, seq); + return structuredClone(committed as TEntry); + }); + } + + async appendRecord(record: NewRecord): Promise; + async appendRecord(record: NewRecord): Promise { + return this.enqueueWrite(() => { + if (!readLane(this.db, this.metadata.id, record.lane)) { + throw new SessionError("invalid_lane", `Lane not found: ${record.lane}`); + } + assertUnusedId(this.db, this.metadata.id, record.id); + const seq = getNextSequence(this.db, this.metadata.id); + const committed: LaneRecord = { ...record, seq, timestamp: Date.now() }; + appendRecordRow(this.db, this.metadata.id, { + seq, + id: record.id, + lane: record.lane, + runId: recordRunId(record), + type: record.type, + opKind: recordOpKind(record), + timestamp: timestampToText(committed.timestamp), + payload: JSON.stringify(record), + }); + if (record.type === "usage") addUsageToStats(this.db, this.metadata.id, record.usage); + advanceSequence(this.db, this.metadata.id, seq); + return structuredClone(committed); + }); + } + + async getEntry(id: string): Promise { + const row = readEntryRow(this.db, this.metadata.id, id); + return row ? decodeEntry(row) : undefined; + } + + async findEntries(query: EntryQuery = {}): Promise { + const rows = readEntryRows(this.db, this.metadata.id, { order: query.order }); + const entries = rows.map(decodeEntry).filter((entry) => matchesEntryQuery(entry, query)); + return structuredClone(query.limit === undefined ? entries : entries.slice(0, query.limit)); + } + + async findEntriesOnBranch(query: EntryQuery & BranchBounds & { start: string }): Promise { + const cached = readCachedBranch(this.db, this.metadata.id, query.start); + if (!cached) { + if (!readEntryRow(this.db, this.metadata.id, query.start)) + throw new SessionError("not_found", `Entry not found: ${query.start}`); + throw new SessionError("invalid_entry", `Branch cache missing entry ${query.start}`); + } + const rows = queryCachedBranchRows(this.db, this.metadata.id, cached, query); + validateCachedBranchRows(rows, query); + const entries = rows + .map(entryRowFromCached) + .map(decodeEntry) + .filter((entry) => matchesEntryQuery(entry, query)); + return structuredClone(query.limit === undefined ? entries : entries.slice(0, query.limit)); + } + + async findRecords(query: RecordQuery = {}): Promise { + const rows = readRecordRows(this.db, this.metadata.id, query); + return structuredClone(rows.map(decodeRecord)); + } + + async getLog(options: LogOptions = {}): Promise { + const afterSeq = options.afterSeq ?? 0; + const entryRows = readEntryRows(this.db, this.metadata.id, { afterSeq, order: "oldestFirst" }); + const recordRows = readRecordRows(this.db, this.metadata.id, { afterSeq }); + const laneRows = readLaneMoveRows(this.db, this.metadata.id, { afterSeq }); + const factRows = readFactRows(this.db, this.metadata.id, { afterSeq }); + + const log: LogItem[] = [ + ...entryRows.map((row) => ({ kind: "entry" as const, seq: row.seq, entry: decodeEntry(row) })), + ...recordRows.map((row) => ({ kind: "record" as const, seq: row.seq, record: decodeRecord(row) })), + ...laneRows.map((row) => ({ kind: "lane" as const, seq: row.seq, lane: row.lane, leafId: row.leaf_id })), + ...factRows.map((row) => { + if (row.kind === "name") + return { + kind: "fact" as const, + seq: row.seq, + fact: "name" as const, + name: JSON.parse(row.value ?? "null") as string, + }; + return { + kind: "fact" as const, + seq: row.seq, + fact: "label" as const, + targetId: row.key ?? "", + label: row.value === null ? undefined : (JSON.parse(row.value) as string), + }; + }), + ].sort((left, right) => left.seq - right.seq); + return structuredClone(options.limit === undefined ? log : log.slice(0, options.limit)); + } + + async getName(): Promise { + const row = readLatestFact(this.db, this.metadata.id, "name", null); + return row?.value === undefined || row.value === null ? undefined : (JSON.parse(row.value) as string); + } + + async setName(name: string): Promise { + return this.enqueueWrite(() => { + const seq = getNextSequence(this.db, this.metadata.id); + appendFact(this.db, this.metadata.id, seq, "name", null, JSON.stringify(name)); + advanceSequence(this.db, this.metadata.id, seq); + }); + } + + async getLabel(id: string): Promise { + const row = readLatestFact(this.db, this.metadata.id, "label", id); + return row?.value === undefined || row.value === null ? undefined : (JSON.parse(row.value) as string); + } + + async setLabel(id: string, label: string | undefined): Promise { + return this.enqueueWrite(() => { + if (!readEntryRow(this.db, this.metadata.id, id)) { + throw new SessionError("not_found", `Entry not found: ${id}`); + } + const seq = getNextSequence(this.db, this.metadata.id); + appendFact(this.db, this.metadata.id, seq, "label", id, label === undefined ? null : JSON.stringify(label)); + advanceSequence(this.db, this.metadata.id, seq); + }); + } + + async getStats(): Promise { + return readStats(this.db, this.metadata.id); + } +} + +function claimStorage( + db: SqliteDatabase, + metadata: SqliteSessionMetadata, + leaseOptions: ResolvedWriterLeaseOptions, + onRelease: () => void, +): SqliteSessionStorage { + requireSessionRow(db, metadata.id); + const claimed = db.transaction(() => { + const lease = acquireWriterLease(db, metadata.id, leaseOptions); + const row = requireSessionRow(db, metadata.id); + readLanes(db, metadata.id); + return { lease, row }; + }); + return new SqliteSessionStorage( + db, + metadataFromRow(claimed.row, metadata.path), + claimed.lease, + leaseOptions, + onRelease, + ); +} + +function metadataFromRow(row: SessionRow, path: string): SqliteSessionMetadata { + const base = rowToMetadata(row, path); + return { ...base, createdAt: Date.parse(base.createdAt) }; +} + +export class SqliteSessionRepository + implements SessionRepository +{ + private databasePath: string | undefined; + private database: SqliteDatabase | undefined; + private databasePromise: Promise | undefined; + private readonly operations = new SerialOperationQueue(); + private readonly activeStorages = new Set(); + private readonly options: SqliteSessionRepositoryOptions; + private readonly leaseOptions: ResolvedWriterLeaseOptions; + + constructor(options: SqliteSessionRepositoryOptions) { + this.options = options; + this.leaseOptions = resolveWriterLeaseOptions(options.writerLease); + } + + private async releaseStoragesForSession(sessionId: string): Promise { + for (const storage of [...this.activeStorages]) { + if (storage.isForSession(sessionId)) await storage.release(); + } + } + + private sessionFromLease( + db: SqliteDatabase, + metadata: SqliteSessionMetadata, + lease: SessionLease, + ): Session { + let storage: SqliteSessionStorage; + storage = new SqliteSessionStorage(db, metadata, lease, this.leaseOptions, () => { + this.activeStorages.delete(storage); + }); + this.activeStorages.add(storage); + return new Session(storage); + } + + private claimSession(db: SqliteDatabase, metadata: SqliteSessionMetadata): Session { + const active = [...this.activeStorages].find((storage) => storage.isForSession(metadata.id)); + if (active) { + readLanes(db, metadata.id); + return new Session(active); + } + let storage: SqliteSessionStorage; + storage = claimStorage(db, metadata, this.leaseOptions, () => { + this.activeStorages.delete(storage); + }); + this.activeStorages.add(storage); + return new Session(storage); + } + + async create(options: SqliteSessionCreateOptions): Promise> { return this.operations.enqueue(async () => { const db = await this.getDatabase(); const path = await this.getDatabasePath(); - const connection = await db.transaction(() => - SqliteSessionConnection.create(db, path, { + const id = options.id ?? uuidv7(); + if (sessionExists(db, id)) throw new SessionError("already_exists", `Session already exists: ${id}`); + const createdAt = Date.now(); + const lease = db.transaction(() => { + insertSessionRow(db, { + id, + createdAt: timestampToText(createdAt), cwd: options.cwd, - sessionId: options.id ?? createSessionId(), parentSessionId: options.parentSessionId, metadata: options.metadata, - }), - ); - this.writers.set(connection.metadata.id, connection); - return this.storage(connection); - }); - } - - open(metadata: SqliteSessionMetadata): Promise> { - this.assertOpen(); - return this.operations.enqueue(() => this.loadSession(metadata)); - } - - private async loadSession(metadata: SqliteSessionMetadata): Promise> { - if ( - !getFileSystemResultOrThrow(await this.env.exists(metadata.path), `Failed to check database ${metadata.path}`) - ) { - throw new SessionError("not_found", `Session not found: ${metadata.id}`); - } - const connection = - this.writers.get(metadata.id) ?? (await SqliteSessionConnection.open(await this.getDatabase(), metadata)); - this.writers.set(metadata.id, connection); - return this.storage(connection); - } - - list(options: SqliteSessionListOptions = {}): Promise { - this.assertOpen(); - return this.operations.enqueue(() => this.listSessions(options)); - } - - private async listSessions(options: SqliteSessionListOptions): Promise { - const path = await this.getDatabasePath(); - if (!getFileSystemResultOrThrow(await this.env.exists(path), `Failed to check database ${path}`)) return []; - const db = await this.getDatabase(); - const rows = options.cwd - ? await db - .prepare( - "SELECT id, created_at, metadata, cwd, parent_session_id, active_leaf_id FROM sessions WHERE cwd = ? ORDER BY created_at DESC", - ) - .all(options.cwd) - : await db - .prepare( - "SELECT id, created_at, metadata, cwd, parent_session_id, active_leaf_id FROM sessions ORDER BY created_at DESC", - ) - .all(); - return rows.map((row) => rowToMetadata(row, path)); - } - - private appendEntry(metadata: SqliteSessionMetadata, entry: SessionTreeEntry): Promise { - this.assertOpen(); - return this.operations.enqueue(async () => { - const connection = - this.writers.get(metadata.id) ?? (await SqliteSessionConnection.open(await this.getDatabase(), metadata)); - this.writers.set(metadata.id, connection); - await connection.appendEntry(entry); - }); - } - - delete(metadata: SqliteSessionMetadata): Promise { - this.assertOpen(); - return this.operations.enqueue(async () => { - const db = await this.getDatabase(); - await db.transaction(async () => { - await db.prepare("DELETE FROM branch_tips WHERE session_id = ?").run(metadata.id); - await db.prepare("DELETE FROM branch_entries WHERE session_id = ?").run(metadata.id); - await db.prepare("DELETE FROM session_entries WHERE session_id = ?").run(metadata.id); - await db.prepare("DELETE FROM entry_materialized WHERE session_id = ?").run(metadata.id); - await db.prepare("DELETE FROM session_materialized WHERE session_id = ?").run(metadata.id); - await db.prepare("DELETE FROM session_sequences WHERE session_id = ?").run(metadata.id); - const result = await db.prepare("DELETE FROM sessions WHERE id = ?").run(metadata.id); - if (result.changes === 0) throw new SessionError("not_found", `Session not found: ${metadata.id}`); - }); - this.writers.delete(metadata.id); - }); - } - - fork( - source: SqliteSessionMetadata, - options: SqliteSessionCreateOptions, - selection: SessionForkSelection, - ): Promise> { - this.assertOpen(); - return this.operations.enqueue(async () => { - const db = await this.getDatabase(); - const connection = await db.transaction(async () => { - const sourceConnection = this.writers.get(source.id) ?? (await SqliteSessionConnection.open(db, source)); - this.writers.set(source.id, sourceConnection); - const entries = await readSessionEntriesForFork(sourceConnection, selection); - const connection = await SqliteSessionConnection.create(db, await this.getDatabasePath(), { - cwd: options.cwd, - sessionId: options.id ?? createSessionId(), - parentSessionId: options.parentSessionId ?? source.id, - metadata: options.metadata ?? source.metadata, }); - for (const entry of entries) await connection.appendEntry(entry, { transaction: false }); - return connection; + createSequence(db, id); + createStats(db, id); + createInitialLane(db, id); + return acquireWriterLease(db, id, this.leaseOptions); }); - this.writers.set(connection.metadata.id, connection); - return this.storage(connection); + const row = requireSessionRow(db, id); + return this.sessionFromLease(db, metadataFromRow(row, path), lease); }); } + async open(metadata: SqliteSessionMetadata): Promise> { + return this.operations.enqueue(async () => this.claimSession(await this.getDatabase(), metadata)); + } + + /** Rebuilds this session's private branch-read cache from canonical entry parent links. */ + async repairBranchCache(metadata: SqliteSessionMetadata): Promise { + return this.operations.enqueue(async () => { + await this.releaseStoragesForSession(metadata.id); + const db = await this.getDatabase(); + db.transaction(() => { + const lease = acquireWriterLease(db, metadata.id, this.leaseOptions); + requireSessionRow(db, metadata.id); + rebuildBranchCache(db, metadata.id); + releaseSessionLease(db, metadata.id, lease); + }); + }); + } + + async list(options: SqliteSessionListOptions = {}): Promise { + return this.operations.enqueue(async () => { + const path = await this.getDatabasePath(); + if (!resultOrThrow(await this.options.env.exists(path), `Failed to check database ${path}`)) return []; + const db = await this.getDatabase(); + const rows = readSessionRows(db, options); + return rows.map((row) => metadataFromRow(row, path)); + }); + } + + async delete(metadata: SqliteSessionMetadata): Promise { + return this.operations.enqueue(async () => { + await this.releaseStoragesForSession(metadata.id); + const db = await this.getDatabase(); + db.transaction(() => { + if (!sessionExists(db, metadata.id)) { + deleteSessionLease(db, metadata.id); + return; + } + acquireWriterLease(db, metadata.id, this.leaseOptions); + deleteBranchCache(db, metadata.id); + deleteFactRows(db, metadata.id); + deleteLaneRows(db, metadata.id); + deleteRecordRows(db, metadata.id); + deleteEntryRows(db, metadata.id); + deleteSessionLease(db, metadata.id); + deleteStats(db, metadata.id); + deleteSequence(db, metadata.id); + deleteSessionRow(db, metadata.id); + }); + }); + } + + async fork( + source: SqliteSessionMetadata, + options: ForkOptions & SqliteSessionCreateOptions, + ): Promise> { + return this.operations.enqueue(async () => { + const db = await this.getDatabase(); + const path = await this.getDatabasePath(); + const sourceMetadata = metadataFromRow(requireSessionRow(db, source.id), path); + const id = options.id ?? uuidv7(); + if (sessionExists(db, id)) throw new SessionError("already_exists", `Session already exists: ${id}`); + + const entries: EntryRow[] = []; + const lanes: { lane: string; leafId: string | null }[] = []; + const branchTips: string[] = []; + let branchForkTargetId: string | null = null; + + if (options.scope === "tree") { + entries.push(...readEntryRows(db, source.id, { order: "oldestFirst" })); + lanes.push(...readLanes(db, source.id).map((row) => ({ lane: row.lane, leafId: row.leaf_id }))); + branchTips.push(...readBranchTipIds(db, source.id)); + } else { + const main = readLane(db, source.id, "main"); + if (!main) throw new SessionError("invalid_lane", "Lane not found: main"); + const selectedEntryId = options.entryId ?? main.leaf_id; + if (selectedEntryId !== null) { + const target = readEntryRow(db, source.id, selectedEntryId); + if (!target || target.type !== "message") { + throw new SessionError( + "invalid_fork_target", + `Fork target is not a message entry: ${selectedEntryId}`, + ); + } + const position = options.position ?? (options.entryId === undefined ? "at" : "before"); + branchForkTargetId = position === "at" ? target.id : target.parent_id; + } + lanes.push({ lane: "main", leafId: branchForkTargetId }); + if (branchForkTargetId !== null) { + const cached = readCachedBranch(db, source.id, branchForkTargetId); + if (!cached) { + throw new SessionError( + "invalid_fork_target", + `Fork target is not on a cached branch: ${branchForkTargetId}`, + ); + } + const rows = queryCachedBranchRows(db, source.id, cached, { order: "oldestFirst" }); + entries.push(...rows.map(entryRowFromCached)); + branchTips.push(branchForkTargetId); + } + } + + const copiedIds = new Set(entries.map((entry) => entry.id)); + const latestName = readLatestFact(db, source.id, "name", null); + const latestLabels = readLatestLabelFacts(db, source.id); + const labelsToCopy = latestLabels.filter( + (row) => options.scope === "tree" || (row.key !== null && copiedIds.has(row.key)), + ); + const createdAt = Date.now(); + const metadata = options.metadata ?? sourceMetadata.metadata; + let lease: SessionLease; + + try { + lease = db.transaction(() => { + insertSessionRow(db, { + id, + createdAt: timestampToText(createdAt), + cwd: options.cwd, + parentSessionId: options.parentSessionId ?? source.id, + metadata, + }); + createSequence(db, id); + createStats(db, id); + + let nextSeq = 1; + const allocateSeq = () => nextSeq++; + for (const entry of entries) { + insertEntryRow(db, id, { + seq: allocateSeq(), + id: entry.id, + parentId: entry.parent_id, + type: entry.type, + timestamp: entry.timestamp, + payload: entry.payload, + }); + } + + if (options.scope === "tree") { + for (const lane of lanes) insertLane(db, id, allocateSeq(), lane.lane, lane.leafId); + } else { + createInitialLane(db, id, "main", branchForkTargetId); + } + + if (latestName?.value !== undefined && latestName.value !== null) { + appendFact(db, id, allocateSeq(), "name", null, latestName.value); + } + for (const label of labelsToCopy) appendFact(db, id, allocateSeq(), "label", label.key, label.value); + + setNextSequence(db, id, nextSeq); + for (const tip of branchTips) buildCachedBranch(db, id, tip); + return acquireWriterLease(db, id, this.leaseOptions); + }); + } catch (error) { + if (error instanceof SessionError) throw error; + throw new SessionError( + "storage", + `Failed to fork SQLite session ${id}`, + error instanceof Error ? error : undefined, + ); + } + + const row = requireSessionRow(db, id); + return this.sessionFromLease(db, metadataFromRow(row, path), lease); + }); + } + + async close(): Promise { + await this.operations.drain(); + for (const storage of [...this.activeStorages]) await storage.release(); + if (this.database) this.database.close(); + this.database = undefined; + this.databasePromise = undefined; + } + async [Symbol.asyncDispose](): Promise { - if (!this.disposePromise) { - this.disposed = true; - this.disposePromise = this.finishDisposal(); - } - await this.disposePromise; - } - - private async finishDisposal(): Promise { - await this.operations.drain(); - const db = this.database ?? (this.databasePromise ? await this.databasePromise : undefined); - this.database = undefined; - this.databasePromise = undefined; - this.writers.clear(); - if (db) await db.close(); - } - - private assertOpen(): void { - if (this.disposed) throw new SessionError("storage", "SQLite session repository is disposed"); - } - - private storage(connection: SqliteSessionConnection): SessionStorage { - const metadata = connection.metadata; - return { - metadata, - readHead: () => this.read(metadata, (current) => current.readHead()), - readEntry: (id) => this.read(metadata, (current) => current.readEntry(id)), - readEntries: (options) => this.read(metadata, (current) => current.readEntries(options)), - appendEntry: (entry) => this.appendEntry(metadata, entry), - findEntriesOnBranch: (query) => this.read(metadata, (current) => current.findEntriesOnBranch(query)), - readPathToRootOrCompaction: (leafId) => - this.read(metadata, (current) => current.readPathToRootOrCompaction(leafId)), - getLabel: (id) => this.read(metadata, (current) => current.getLabel(id)), - getName: () => this.read(metadata, (current) => current.getName()), - getStats: () => this.read(metadata, (current) => current.getStats()), - }; - } - - private read( - metadata: SqliteSessionMetadata, - read: (connection: SqliteSessionConnection) => Promise, - ): Promise { - this.assertOpen(); - return this.operations.enqueue(async () => { - const connection = - this.writers.get(metadata.id) ?? (await SqliteSessionConnection.open(await this.getDatabase(), metadata)); - this.writers.set(metadata.id, connection); - return read(connection); - }); + await this.close(); } private async getDatabasePath(): Promise { - this.databasePath ??= getFileSystemResultOrThrow( - await this.env.absolutePath(this.databasePathInput), - `Failed to resolve SQLite sessions database ${this.databasePathInput}`, + this.databasePath ??= resultOrThrow( + await this.options.env.absolutePath(this.options.databasePath), + `Failed to resolve SQLite sessions database ${this.options.databasePath}`, ); return this.databasePath; } @@ -261,67 +908,18 @@ class SqliteSessionBackend { private async openDatabase(): Promise { const path = await this.getDatabasePath(); - const directory = getParentPath(path); - getFileSystemResultOrThrow( - await this.env.createDir(directory, { recursive: true }), - `Failed to create SQLite sessions directory ${directory}`, + resultOrThrow( + await this.options.env.createDir(getParentPath(path), { recursive: true }), + `Failed to create SQLite sessions directory ${path}`, ); - const db = await this.sqlite.open(path); + const db = await this.options.sqlite.open(path); try { - await configureSqliteDatabase(db); + configureSqliteDatabase(db); await applyMigrations(db); return db; } catch (error) { - await db.close(); + db.close(); throw error; } } } - -export interface SqliteSessionRepositoryOptions extends SqliteSessionBackendOptions { - contextBuildOptions?: SessionContextBuildOptions; -} - -export class SqliteSessionRepository - implements SessionRepository -{ - private readonly backend: SqliteSessionBackend; - private readonly contextBuildOptions: SessionContextBuildOptions; - - constructor(options: SqliteSessionRepositoryOptions) { - const { contextBuildOptions, ...backendOptions } = options; - this.backend = new SqliteSessionBackend(backendOptions); - this.contextBuildOptions = contextBuildOptions ?? {}; - } - - async create(options: SqliteSessionCreateOptions): Promise> { - return createSession(await this.backend.create(options), this.contextBuildOptions); - } - - async open(metadata: SqliteSessionMetadata): Promise> { - return createSession(await this.backend.open(metadata), this.contextBuildOptions); - } - - async list(options?: SqliteSessionListOptions): Promise { - return await this.backend.list(options); - } - - async delete(metadata: SqliteSessionMetadata): Promise { - await this.backend.delete(metadata); - } - - async fork( - source: SqliteSessionMetadata, - options: SessionForkOptions & SqliteSessionCreateOptions, - ): Promise> { - const { entryId: _entryId, position: _position, ...createOptions } = options; - return createSession( - await this.backend.fork(source, createOptions, createSessionForkSelection(options)), - this.contextBuildOptions, - ); - } - - async [Symbol.asyncDispose](): Promise { - await this.backend[Symbol.asyncDispose](); - } -} diff --git a/packages/storage/sqlite-node/src/sqlite/search-backend.ts b/packages/storage/sqlite-node/src/sqlite/search-backend.ts index 6ba2e865f..73f2cd985 100644 --- a/packages/storage/sqlite-node/src/sqlite/search-backend.ts +++ b/packages/storage/sqlite-node/src/sqlite/search-backend.ts @@ -17,10 +17,10 @@ function getParentPath(path: string): string { return normalized.slice(0, lastSlash); } -async function configureSqliteDatabase(db: SqliteDatabase): Promise { - await db.exec("PRAGMA journal_mode=WAL"); - await db.exec("PRAGMA synchronous=FULL"); - await db.exec("PRAGMA busy_timeout=5000"); +function configureSqliteDatabase(db: SqliteDatabase): void { + db.exec("PRAGMA journal_mode=WAL"); + db.exec("PRAGMA synchronous=FULL"); + db.exec("PRAGMA busy_timeout=5000"); } export interface SqliteSessionSearchOptions { @@ -29,33 +29,33 @@ export interface SqliteSessionSearchOptions { databasePath: string; } -async function tableExists(db: SqliteDatabase, name: string): Promise { - return !!(await db +function tableExists(db: SqliteDatabase, name: string): boolean { + return !!db .prepare("SELECT 1 AS found FROM sqlite_master WHERE type = 'table' AND name = ? LIMIT 1") - .get<{ found: number }>(name)); + .get<{ found: number }>(name); } -async function ensureSearchSchema(db: SqliteDatabase): Promise { - const ftsExists = await tableExists(db, "session_search_fts"); - await db.exec(` +function ensureSearchSchema(db: SqliteDatabase): void { + const ftsExists = tableExists(db, "session_search_fts"); + db.exec(` CREATE VIRTUAL TABLE IF NOT EXISTS session_search_fts USING fts5( payload, - content = 'session_entries', + content = 'entries', content_rowid = 'rowid', tokenize = 'trigram remove_diacritics 1' ); -CREATE TRIGGER IF NOT EXISTS session_search_fts_ai AFTER INSERT ON session_entries BEGIN +CREATE TRIGGER IF NOT EXISTS session_search_fts_ai AFTER INSERT ON entries BEGIN INSERT INTO session_search_fts(rowid, payload) VALUES (new.rowid, new.payload); END; -CREATE TRIGGER IF NOT EXISTS session_search_fts_ad AFTER DELETE ON session_entries BEGIN +CREATE TRIGGER IF NOT EXISTS session_search_fts_ad AFTER DELETE ON entries BEGIN INSERT INTO session_search_fts(session_search_fts, rowid, payload) VALUES('delete', old.rowid, old.payload); END; -CREATE TRIGGER IF NOT EXISTS session_search_fts_au AFTER UPDATE OF payload ON session_entries BEGIN +CREATE TRIGGER IF NOT EXISTS session_search_fts_au AFTER UPDATE OF payload ON entries BEGIN INSERT INTO session_search_fts(session_search_fts, rowid, payload) VALUES('delete', old.rowid, old.payload); INSERT INTO session_search_fts(rowid, payload) VALUES (new.rowid, new.payload); END; `); - if (!ftsExists) await db.exec("INSERT INTO session_search_fts(session_search_fts) VALUES('rebuild')"); + if (!ftsExists) db.exec("INSERT INTO session_search_fts(session_search_fts) VALUES('rebuild')"); } /** SQLite FTS search over a co-located canonical session database. */ @@ -86,12 +86,12 @@ class SqliteSessionSearch implements SessionSearch { ); const db = await this.options.sqlite.open(path); try { - await configureSqliteDatabase(db); + configureSqliteDatabase(db); await applyMigrations(db); - await ensureSearchSchema(db); + ensureSearchSchema(db); return db; } catch (error) { - await db.close(); + db.close(); throw error; } } @@ -102,9 +102,9 @@ class SqliteSessionSearch implements SessionSearch { const db = await this.openDatabase(); try { const query = `"${text.replaceAll('"', '""')}"`; - const rows = await db + const rows = db .prepare( - "SELECT s.id, s.created_at, s.metadata, s.cwd, s.parent_session_id, s.active_leaf_id, se.id AS entry_id, se.timestamp, bm25(session_search_fts) AS score FROM session_search_fts JOIN session_entries se ON se.rowid = session_search_fts.rowid JOIN sessions s ON s.id = se.session_id WHERE session_search_fts MATCH ? AND (? IS NULL OR s.cwd = ?) ORDER BY score", + "SELECT s.id, s.created_at, s.metadata, s.cwd, s.parent_session_id, se.id AS entry_id, se.timestamp, bm25(session_search_fts) AS score FROM session_search_fts JOIN entries se ON se.rowid = session_search_fts.rowid JOIN sessions s ON s.id = se.session_id WHERE session_search_fts MATCH ? AND (? IS NULL OR s.cwd = ?) ORDER BY score", ) .all( query, @@ -119,7 +119,7 @@ class SqliteSessionSearch implements SessionSearch { score: row.score, })); } finally { - await db.close(); + db.close(); } } } diff --git a/packages/storage/sqlite-node/src/sqlite/storage/branch-cache.ts b/packages/storage/sqlite-node/src/sqlite/storage/branch-cache.ts deleted file mode 100644 index 9f448927b..000000000 --- a/packages/storage/sqlite-node/src/sqlite/storage/branch-cache.ts +++ /dev/null @@ -1,326 +0,0 @@ -import { uuidv7 } from "@earendil-works/pi-ai"; -import type { SqliteDatabase } from "../types.ts"; -import type { SessionEntryRow } from "./session-entries.ts"; -import { invalidSession } from "./shared.ts"; - -/** Derived root-to-tip paths. Canonical parent links remain authoritative. */ -export interface CachedBranch { - branchId: string; - leafSeq: number; -} - -export interface CachedBranchQuery { - stopAtType?: SessionEntryRow["type"]; - stopAtId?: string; - order?: "newestFirst" | "oldestFirst"; -} - -export async function readCachedBranch( - db: SqliteDatabase, - sessionId: string, - leafId: string, -): Promise { - const membership = await db - .prepare( - "SELECT branch_id, entry_seq FROM branch_entries WHERE session_id = ? AND entry_id = ? ORDER BY branch_id LIMIT 1", - ) - .get<{ branch_id: string; entry_seq: number }>(sessionId, leafId); - if (!membership) return undefined; - return { branchId: membership.branch_id, leafSeq: membership.entry_seq }; -} - -export async function isCachedBranchValid( - db: SqliteDatabase, - sessionId: string, - branch: CachedBranch, - leafId: string, - startSeq = 0, -): Promise { - const result = await db - .prepare( - `WITH path AS ( - SELECT - b.entry_id, - b.entry_seq, - e.id AS stored_entry_id, - e.parent_id, - LAG(b.entry_id) OVER (ORDER BY b.entry_seq) AS previous_entry_id - FROM branch_entries AS b - LEFT JOIN session_entries AS e ON e.session_id = b.session_id AND e.id = b.entry_id - WHERE b.session_id = ? AND b.branch_id = ? AND b.entry_seq BETWEEN ? AND ? - ) - SELECT - COUNT(*) AS row_count, - COALESCE(SUM( - stored_entry_id IS NULL OR - (? = 0 AND previous_entry_id IS NULL AND parent_id IS NOT NULL) OR - (previous_entry_id IS NOT NULL AND parent_id IS NOT previous_entry_id) - ), 0) AS invalid_count, - COALESCE(MAX(entry_seq = ? AND entry_id = ?), 0) AS contains_leaf - FROM path`, - ) - .get<{ row_count: number; invalid_count: number; contains_leaf: number }>( - sessionId, - branch.branchId, - startSeq, - branch.leafSeq, - startSeq, - branch.leafSeq, - leafId, - ); - return result?.row_count !== 0 && result?.invalid_count === 0 && result.contains_leaf === 1; -} - -export async function readNewestCachedStopSeq( - db: SqliteDatabase, - sessionId: string, - branch: CachedBranch, - stopAtType: SessionEntryRow["type"] | undefined, - stopAtId: string | undefined, -): Promise { - const predicates: string[] = []; - const params: unknown[] = [sessionId, branch.branchId, branch.leafSeq]; - if (stopAtType !== undefined) { - predicates.push("e.type = ?"); - params.push(stopAtType); - } - if (stopAtId !== undefined) { - predicates.push("b.entry_id = ?"); - params.push(stopAtId); - } - if (predicates.length === 0) return undefined; - const row = await db - .prepare( - `SELECT MAX(b.entry_seq) AS entry_seq - FROM branch_entries AS b - JOIN session_entries AS e ON e.session_id = b.session_id AND e.id = b.entry_id - WHERE b.session_id = ? AND b.branch_id = ? AND b.entry_seq <= ? - AND (${predicates.join(" OR ")})`, - ) - .get<{ entry_seq: number | null }>(...params); - return row?.entry_seq ?? undefined; -} - -export async function readCachedBranchRows( - db: SqliteDatabase, - sessionId: string, - branch: CachedBranch, - startSeq: number, -): Promise { - return db - .prepare( - `SELECT e.session_id, e.id, e.entry_seq, e.parent_id, e.type, e.timestamp, e.payload - FROM branch_entries AS b - JOIN session_entries AS e ON e.session_id = b.session_id AND e.id = b.entry_id - WHERE b.session_id = ? AND b.branch_id = ? AND b.entry_seq BETWEEN ? AND ? - ORDER BY b.entry_seq`, - ) - .all(sessionId, branch.branchId, startSeq, branch.leafSeq); -} - -export async function queryCachedBranchRows( - db: SqliteDatabase, - sessionId: string, - branch: CachedBranch, - query: CachedBranchQuery, -): Promise { - const oldestFirst = query.order === "oldestFirst"; - const boundaryParams: unknown[] = [sessionId, branch.branchId, branch.leafSeq]; - const stopPredicates: string[] = []; - if (query.stopAtType !== undefined) { - stopPredicates.push("stop_entry.type = ?"); - boundaryParams.push(query.stopAtType); - } - if (query.stopAtId !== undefined) { - stopPredicates.push("stop.entry_id = ?"); - boundaryParams.push(query.stopAtId); - } - - const boundary = stopPredicates.length - ? `WITH boundary AS ( - SELECT ${oldestFirst ? "MIN" : "MAX"}(stop.entry_seq) AS entry_seq - FROM branch_entries AS stop - JOIN session_entries AS stop_entry - ON stop_entry.session_id = stop.session_id AND stop_entry.id = stop.entry_id - WHERE stop.session_id = ? AND stop.branch_id = ? AND stop.entry_seq <= ? - AND (${stopPredicates.join(" OR ")}) - )` - : ""; - const range = stopPredicates.length - ? `AND b.entry_seq ${oldestFirst ? "<=" : ">="} COALESCE( - (SELECT entry_seq FROM boundary), ${oldestFirst ? branch.leafSeq : 0} - )` - : ""; - const sql = `${boundary} - SELECT e.session_id, e.id, e.entry_seq, e.parent_id, e.type, e.timestamp, e.payload - FROM branch_entries AS b - JOIN session_entries AS e ON e.session_id = b.session_id AND e.id = b.entry_id - WHERE b.session_id = ? AND b.branch_id = ? AND b.entry_seq <= ? - ${range} - ORDER BY b.entry_seq ${oldestFirst ? "ASC" : "DESC"}`; - - const params = [...(stopPredicates.length === 0 ? [] : boundaryParams), sessionId, branch.branchId, branch.leafSeq]; - return db.prepare(sql).all(...params); -} - -export async function readCachedEntryRowsByType( - db: SqliteDatabase, - sessionId: string, - branch: CachedBranch, - type: SessionEntryRow["type"], -): Promise { - // Drive the join from the usually sparse entry type. Ordering from branch_entries - // makes SQLite scan the complete cached path before filtering by type. - return db - .prepare( - `SELECT e.session_id, e.id, e.entry_seq, e.parent_id, e.type, e.timestamp, e.payload - FROM session_entries AS e INDEXED BY idx_session_entries_session_type - CROSS JOIN branch_entries AS b - WHERE e.session_id = ? AND e.type = ? - AND b.session_id = e.session_id AND b.entry_id = e.id - AND b.branch_id = ? AND b.entry_seq <= ? - ORDER BY e.entry_seq DESC`, - ) - .all(sessionId, type, branch.branchId, branch.leafSeq); -} - -export async function readCachedEntrySeq( - db: SqliteDatabase, - sessionId: string, - branchId: string, - entryId: string, -): Promise { - const row = await db - .prepare("SELECT entry_seq FROM branch_entries WHERE session_id = ? AND branch_id = ? AND entry_id = ?") - .get<{ entry_seq: number }>(sessionId, branchId, entryId); - return row?.entry_seq; -} - -export async function rebuildCachedBranch( - db: SqliteDatabase, - sessionId: string, - leafId: string, - branchIdToReplace?: string, -): Promise { - await db.exec("SAVEPOINT rebuild_branch_cache"); - try { - const tip = await db - .prepare("SELECT branch_id FROM branch_tips WHERE session_id = ? AND tip_id = ?") - .get<{ branch_id: string }>(sessionId, leafId); - const branchIds = new Set([branchIdToReplace, tip?.branch_id].filter((id): id is string => id !== undefined)); - for (const branchId of branchIds) { - await db.prepare("DELETE FROM branch_tips WHERE session_id = ? AND branch_id = ?").run(sessionId, branchId); - await db.prepare("DELETE FROM branch_entries WHERE session_id = ? AND branch_id = ?").run(sessionId, branchId); - } - - const branchId = uuidv7(); - await db - .prepare( - `WITH RECURSIVE path(id, entry_seq, parent_id) AS ( - SELECT id, entry_seq, parent_id - FROM session_entries - WHERE session_id = ? AND id = ? - UNION ALL - SELECT parent.id, parent.entry_seq, parent.parent_id - FROM session_entries AS parent - JOIN path AS child ON child.parent_id = parent.id - WHERE parent.session_id = ? - ) - INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) - SELECT ?, ?, id, entry_seq FROM path`, - ) - .run(sessionId, leafId, sessionId, sessionId, branchId); - await db - .prepare("INSERT INTO branch_tips (session_id, tip_id, branch_id) VALUES (?, ?, ?)") - .run(sessionId, leafId, branchId); - await db.exec("RELEASE SAVEPOINT rebuild_branch_cache"); - } catch (error) { - try { - await db.exec("ROLLBACK TO SAVEPOINT rebuild_branch_cache"); - await db.exec("RELEASE SAVEPOINT rebuild_branch_cache"); - } catch { - // Preserve the original repair failure. - } - throw error; - } -} - -async function extendBranch( - db: SqliteDatabase, - sessionId: string, - branchId: string, - parentId: string, - entryId: string, - entrySeq: number, -): Promise { - await db - .prepare("INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) VALUES (?, ?, ?, ?)") - .run(sessionId, branchId, entryId, entrySeq); - const result = await db - .prepare("UPDATE branch_tips SET tip_id = ? WHERE session_id = ? AND branch_id = ? AND tip_id = ?") - .run(entryId, sessionId, branchId, parentId); - if (result.changes !== 1) throw invalidSession(`branch tip ${parentId} changed during append`); -} - -export async function appendEntryToBranchCache( - db: SqliteDatabase, - sessionId: string, - entryId: string, - entrySeq: number, - parentId: string | null, - repairParent: (parentId: string) => Promise, -): Promise { - if (parentId === null) { - const branchId = uuidv7(); - await db - .prepare("INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) VALUES (?, ?, ?, ?)") - .run(sessionId, branchId, entryId, entrySeq); - await db - .prepare("INSERT INTO branch_tips (session_id, tip_id, branch_id) VALUES (?, ?, ?)") - .run(sessionId, entryId, branchId); - return; - } - - let tip = await db - .prepare("SELECT branch_id FROM branch_tips WHERE session_id = ? AND tip_id = ?") - .get<{ branch_id: string }>(sessionId, parentId); - if (tip) { - await extendBranch(db, sessionId, tip.branch_id, parentId, entryId, entrySeq); - return; - } - - const source = await db - .prepare( - `SELECT b.branch_id, b.entry_seq - FROM branch_entries AS b - WHERE b.session_id = ? AND b.entry_id = ? - ORDER BY b.branch_id - LIMIT 1`, - ) - .get<{ branch_id: string; entry_seq: number }>(sessionId, parentId); - if (!source) { - await repairParent(parentId); - tip = await db - .prepare("SELECT branch_id FROM branch_tips WHERE session_id = ? AND tip_id = ?") - .get<{ branch_id: string }>(sessionId, parentId); - if (!tip) throw invalidSession(`branch cache repair did not create tip ${parentId}`); - await extendBranch(db, sessionId, tip.branch_id, parentId, entryId, entrySeq); - return; - } - - const branchId = uuidv7(); - await db - .prepare( - `INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) - SELECT session_id, ?, entry_id, entry_seq - FROM branch_entries - WHERE session_id = ? AND branch_id = ? AND entry_seq <= ?`, - ) - .run(branchId, sessionId, source.branch_id, source.entry_seq); - await db - .prepare("INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) VALUES (?, ?, ?, ?)") - .run(sessionId, branchId, entryId, entrySeq); - await db - .prepare("INSERT INTO branch_tips (session_id, tip_id, branch_id) VALUES (?, ?, ?)") - .run(sessionId, entryId, branchId); -} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts b/packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts new file mode 100644 index 000000000..dfdadf64c --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts @@ -0,0 +1,146 @@ +import type { Entry } from "@earendil-works/pi-agent-core/experimental"; +import type { SqliteDatabase } from "../types.ts"; + +/** Derived root-to-tip branch cache membership. Canonical parent links remain in entries. */ +export interface CachedBranch { + branchId: string; + leafSeq: number; +} + +export interface CachedBranchEntryRow { + session_id: string; + id: string; + entry_seq: number; + parent_id: string | null; + type: Entry["type"]; + timestamp: string; + payload: string; +} + +export interface CachedBranchQuery { + stopAtType?: Entry["type"]; + stopAtId?: string; + order?: "newestFirst" | "oldestFirst"; +} + +export function readCachedBranch(db: SqliteDatabase, sessionId: string, leafId: string) { + const membership = db + .prepare( + "SELECT branch_id, entry_seq FROM branch_entries WHERE session_id = ? AND entry_id = ? ORDER BY branch_id LIMIT 1", + ) + .get<{ branch_id: string; entry_seq: number }>(sessionId, leafId); + if (!membership) return undefined; + return { branchId: membership.branch_id, leafSeq: membership.entry_seq }; +} + +export function queryCachedBranchRows( + db: SqliteDatabase, + sessionId: string, + branch: CachedBranch, + query: CachedBranchQuery, +) { + const oldestFirst = query.order === "oldestFirst"; + const boundaryParams: unknown[] = [sessionId, branch.branchId, branch.leafSeq]; + const stopPredicates: string[] = []; + if (query.stopAtType !== undefined) { + stopPredicates.push("stop_entry.type = ?"); + boundaryParams.push(query.stopAtType); + } + if (query.stopAtId !== undefined) { + stopPredicates.push("stop.entry_id = ?"); + boundaryParams.push(query.stopAtId); + } + + const boundary = stopPredicates.length + ? `WITH boundary AS ( + SELECT ${oldestFirst ? "MIN" : "MAX"}(stop.entry_seq) AS entry_seq + FROM branch_entries AS stop + JOIN entries AS stop_entry + ON stop_entry.session_id = stop.session_id AND stop_entry.id = stop.entry_id + WHERE stop.session_id = ? AND stop.branch_id = ? AND stop.entry_seq <= ? + AND (${stopPredicates.join(" OR ")}) + )` + : ""; + const range = stopPredicates.length + ? `AND b.entry_seq ${oldestFirst ? "<=" : ">="} COALESCE( + (SELECT entry_seq FROM boundary), ${oldestFirst ? branch.leafSeq : 0} + )` + : ""; + const sql = `${boundary} + SELECT e.session_id, e.id, e.seq AS entry_seq, e.parent_id, e.type, e.timestamp, e.payload + FROM branch_entries AS b + JOIN entries AS e ON e.session_id = b.session_id AND e.id = b.entry_id + WHERE b.session_id = ? AND b.branch_id = ? AND b.entry_seq <= ? + ${range} + ORDER BY b.entry_seq ${oldestFirst ? "ASC" : "DESC"}`; + + const params = [...(stopPredicates.length === 0 ? [] : boundaryParams), sessionId, branch.branchId, branch.leafSeq]; + return db.prepare(sql).all(...params); +} + +export function deleteBranchEntries(db: SqliteDatabase, sessionId: string) { + db.prepare("DELETE FROM branch_entries WHERE session_id = ?").run(sessionId); +} + +export function insertBranchEntry( + db: SqliteDatabase, + sessionId: string, + branchId: string, + entryId: string, + entrySeq: number, + entryType: string, + customType: string | null, +) { + db.prepare( + `INSERT INTO branch_entries + (session_id, branch_id, entry_id, entry_seq, entry_type, custom_type) + VALUES (?, ?, ?, ?, ?, ?)`, + ).run(sessionId, branchId, entryId, entrySeq, entryType, customType); +} + +export function insertBranchEntriesForPath(db: SqliteDatabase, sessionId: string, branchId: string, leafId: string) { + db.prepare( + `WITH RECURSIVE path(id, entry_seq, parent_id, type, custom_type) AS ( + SELECT id, seq, parent_id, type, + CASE WHEN type = 'custom' THEN json_extract(payload, '$.customType') ELSE NULL END + FROM entries + WHERE session_id = ? AND id = ? + UNION ALL + SELECT parent.id, parent.seq, parent.parent_id, parent.type, + CASE WHEN parent.type = 'custom' THEN json_extract(parent.payload, '$.customType') ELSE NULL END + FROM entries AS parent + JOIN path AS child ON child.parent_id = parent.id + WHERE parent.session_id = ? + ) + INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq, entry_type, custom_type) + SELECT ?, ?, id, entry_seq, type, custom_type FROM path`, + ).run(sessionId, leafId, sessionId, sessionId, branchId); +} + +export function readBranchContainingEntry(db: SqliteDatabase, sessionId: string, entryId: string) { + const row = db + .prepare( + `SELECT b.branch_id, b.entry_seq + FROM branch_entries AS b + WHERE b.session_id = ? AND b.entry_id = ? + ORDER BY b.branch_id + LIMIT 1`, + ) + .get<{ branch_id: string; entry_seq: number }>(sessionId, entryId); + return row === undefined ? undefined : { branchId: row.branch_id, entrySeq: row.entry_seq }; +} + +export function copyBranchEntriesThroughSeq( + db: SqliteDatabase, + sessionId: string, + targetBranchId: string, + sourceBranchId: string, + throughSeq: number, +) { + db.prepare( + `INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq, entry_type, custom_type) + SELECT session_id, ?, entry_id, entry_seq, entry_type, custom_type + FROM branch_entries + WHERE session_id = ? AND branch_id = ? AND entry_seq <= ?`, + ).run(targetBranchId, sessionId, sourceBranchId, throughSeq); +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/branch-tips.ts b/packages/storage/sqlite-node/src/sqlite/storage/branch-tips.ts new file mode 100644 index 000000000..e4dffc488 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/branch-tips.ts @@ -0,0 +1,40 @@ +import type { SqliteDatabase } from "../types.ts"; + +export function readBranchTipIds(db: SqliteDatabase, sessionId: string) { + return db + .prepare("SELECT tip_id FROM branch_tips WHERE session_id = ? ORDER BY tip_id") + .all<{ tip_id: string }>(sessionId) + .map((row) => row.tip_id); +} + +export function readBranchTipBranchId(db: SqliteDatabase, sessionId: string, tipId: string) { + const tip = db + .prepare("SELECT branch_id FROM branch_tips WHERE session_id = ? AND tip_id = ?") + .get<{ branch_id: string }>(sessionId, tipId); + return tip?.branch_id; +} + +export function insertBranchTip(db: SqliteDatabase, sessionId: string, tipId: string, branchId: string) { + db.prepare("INSERT INTO branch_tips (session_id, tip_id, branch_id) VALUES (?, ?, ?)").run( + sessionId, + tipId, + branchId, + ); +} + +export function updateBranchTip( + db: SqliteDatabase, + sessionId: string, + branchId: string, + oldTipId: string, + newTipId: string, +) { + const result = db + .prepare("UPDATE branch_tips SET tip_id = ? WHERE session_id = ? AND branch_id = ? AND tip_id = ?") + .run(newTipId, sessionId, branchId, oldTipId); + return result.changes === 1; +} + +export function deleteBranchTips(db: SqliteDatabase, sessionId: string) { + db.prepare("DELETE FROM branch_tips WHERE session_id = ?").run(sessionId); +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/entries.ts b/packages/storage/sqlite-node/src/sqlite/storage/entries.ts new file mode 100644 index 000000000..fc227d45d --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/entries.ts @@ -0,0 +1,75 @@ +import type { Entry, EntryOrder } from "@earendil-works/pi-agent-core/experimental"; +import type { SqliteDatabase } from "../types.ts"; + +export interface EntryRow { + session_id: string; + seq: number; + id: string; + parent_id: string | null; + type: Entry["type"]; + timestamp: string; + payload: string; +} + +export interface NewEntryRow { + seq: number; + id: string; + parentId: string | null; + type: Entry["type"]; + timestamp: string; + payload: string; +} + +export function entryPayload(entry: Entry): Record { + const { type: _type, id: _id, seq: _seq, parentId: _parentId, timestamp: _timestamp, ...payload } = entry; + return payload; +} + +function orderedSql(order: EntryOrder | undefined): string { + return order === "oldestFirst" ? "ASC" : "DESC"; +} + +export function insertEntryRow(db: SqliteDatabase, sessionId: string, entry: NewEntryRow) { + db.prepare( + "INSERT INTO entries (session_id, id, seq, parent_id, type, timestamp, payload) VALUES (?, ?, ?, ?, ?, ?, ?)", + ).run(sessionId, entry.id, entry.seq, entry.parentId, entry.type, entry.timestamp, entry.payload); +} + +export function readEntryRow(db: SqliteDatabase, sessionId: string, entryId: string) { + return db + .prepare( + "SELECT session_id, seq, id, parent_id, type, timestamp, payload FROM entries WHERE session_id = ? AND id = ?", + ) + .get(sessionId, entryId); +} + +export function readEntryRows( + db: SqliteDatabase, + sessionId: string, + options: { afterSeq?: number; order?: EntryOrder } = {}, +) { + const predicates = ["session_id = ?"]; + const params: unknown[] = [sessionId]; + if (options.afterSeq !== undefined) { + predicates.push("seq > ?"); + params.push(options.afterSeq); + } + return db + .prepare( + `SELECT session_id, seq, id, parent_id, type, timestamp, payload + FROM entries + WHERE ${predicates.join(" AND ")} + ORDER BY seq ${orderedSql(options.order)}`, + ) + .all(...params); +} + +export function idExistsInEntries(db: SqliteDatabase, sessionId: string, id: string) { + return !!db + .prepare("SELECT 1 AS found FROM entries WHERE session_id = ? AND id = ? LIMIT 1") + .get<{ found: number }>(sessionId, id); +} + +export function deleteEntryRows(db: SqliteDatabase, sessionId: string) { + db.prepare("DELETE FROM entries WHERE session_id = ?").run(sessionId); +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/facts.ts b/packages/storage/sqlite-node/src/sqlite/storage/facts.ts new file mode 100644 index 000000000..50d056f47 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/facts.ts @@ -0,0 +1,73 @@ +import type { SqliteDatabase } from "../types.ts"; + +export interface FactRow { + session_id: string; + seq: number; + kind: string; + key: string | null; + value: string | null; +} + +export function appendFact( + db: SqliteDatabase, + sessionId: string, + seq: number, + kind: string, + key: string | null, + value: string | null, +) { + db.prepare("INSERT INTO facts (session_id, seq, kind, key, value) VALUES (?, ?, ?, ?, ?)").run( + sessionId, + seq, + kind, + key, + value, + ); +} + +export function readLatestFact(db: SqliteDatabase, sessionId: string, kind: string, key: string | null) { + return db + .prepare( + `SELECT session_id, seq, kind, key, value + FROM facts + WHERE session_id = ? AND kind = ? AND key IS ? + ORDER BY seq DESC + LIMIT 1`, + ) + .get(sessionId, kind, key); +} + +export function readLatestLabelFacts(db: SqliteDatabase, sessionId: string) { + return db + .prepare( + `SELECT key, value FROM ( + SELECT key, value, ROW_NUMBER() OVER (PARTITION BY key ORDER BY seq DESC) AS rank + FROM facts + WHERE session_id = ? AND kind = 'label' + ) + WHERE rank = 1 AND value IS NOT NULL + ORDER BY key`, + ) + .all<{ key: string; value: string }>(sessionId); +} + +export function readFactRows(db: SqliteDatabase, sessionId: string, options: { afterSeq?: number } = {}) { + const predicates = ["session_id = ?"]; + const params: unknown[] = [sessionId]; + if (options.afterSeq !== undefined) { + predicates.push("seq > ?"); + params.push(options.afterSeq); + } + return db + .prepare( + `SELECT session_id, seq, kind, key, value + FROM facts + WHERE ${predicates.join(" AND ")} + ORDER BY seq`, + ) + .all(...params); +} + +export function deleteFactRows(db: SqliteDatabase, sessionId: string) { + db.prepare("DELETE FROM facts WHERE session_id = ?").run(sessionId); +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/index.ts b/packages/storage/sqlite-node/src/sqlite/storage/index.ts deleted file mode 100644 index c60a73d78..000000000 --- a/packages/storage/sqlite-node/src/sqlite/storage/index.ts +++ /dev/null @@ -1,459 +0,0 @@ -import type { - SessionBranchQuery, - SessionEntryCursorOptions, - SessionStats, - SessionTreeEntry, -} from "@earendil-works/pi-agent-core"; -import { SessionError, toError } from "@earendil-works/pi-agent-core"; -import type { SqliteDatabase, SqliteSessionMetadata } from "../types.ts"; -import { - appendEntryToBranchCache, - type CachedBranch, - isCachedBranchValid, - queryCachedBranchRows, - readCachedBranch, - readCachedBranchRows, - readCachedEntryRowsByType, - readCachedEntrySeq, - readNewestCachedStopSeq, - rebuildCachedBranch, -} from "./branch-cache.ts"; -import { decodeEntry, encodeEntry, type SessionEntryRow } from "./session-entries.ts"; -import { - applyEntryToMaterializedState, - createEmptyMaterializedState, - type EntryMaterializedRow, - entryMaterializedValues, - materializedStateFromRows, - materializedStateValues, - type SessionMaterializedRow, - type SessionMaterializedState, - serializeSummary, -} from "./session-materialized.ts"; -import { advanceSequence, getNextSequence } from "./session-sequences.ts"; -import { rowToMetadata, type SessionRow } from "./sessions.ts"; -import { invalidEntry, invalidSession, leafIdAfterEntry } from "./shared.ts"; - -function decodeEntryRows(entryRows: SessionEntryRow[]): SessionTreeEntry[] { - const entries: SessionTreeEntry[] = []; - for (const entryRow of entryRows) { - try { - const entry = decodeEntry(entryRow); - entries.push(entry); - } catch (error) { - throw invalidEntry(`failed to decode entry ${entryRow.id}`, toError(error)); - } - } - return entries; -} - -async function loadSqliteSession( - db: SqliteDatabase, - sessionId: string, -): Promise<{ - row: SessionRow; - materializedState: SessionMaterializedState; -}> { - const row = await db - .prepare("SELECT id, created_at, metadata, cwd, parent_session_id, active_leaf_id FROM sessions WHERE id = ?") - .get(sessionId); - if (!row) throw new SessionError("not_found", `Session not found: ${sessionId}`); - - const materializedRow = await db - .prepare("SELECT session_id, payload FROM session_materialized WHERE session_id = ?") - .get(sessionId); - if (!materializedRow) throw invalidSession(`missing materialized row for session ${sessionId}`); - const entryMaterializedRows = await db - .prepare( - "SELECT session_id, entry_seq, type, payload FROM entry_materialized WHERE session_id = ? ORDER BY entry_seq, type", - ) - .all(sessionId); - return { - row, - materializedState: materializedStateFromRows(materializedRow, entryMaterializedRows), - }; -} - -export class SqliteSessionConnection { - private readonly db: SqliteDatabase; - readonly metadata: SqliteSessionMetadata; - private byId: Map; - private materializedState: SessionMaterializedState; - - async findEntriesOnBranch(query: SessionBranchQuery & { start: string | null }): Promise { - if (query.limit !== undefined && (!Number.isInteger(query.limit) || query.limit <= 0)) { - throw new RangeError("Session branch query limit must be a positive integer"); - } - if (query.start === null) return []; - const startId = query.start; - let cached = await readCachedBranch(this.db, this.metadata.id, startId); - const validationStartSeq = - cached && query.order !== "oldestFirst" - ? await readNewestCachedStopSeq(this.db, this.metadata.id, cached, query.stopAtType, query.stopAtId) - : undefined; - if (!cached || !(await isCachedBranchValid(this.db, this.metadata.id, cached, startId, validationStartSeq))) { - if (query.order !== "oldestFirst" && (query.stopAtId !== undefined || query.stopAtType !== undefined)) { - return this.findEntriesOnCanonicalBranch({ ...query, start: startId }); - } - cached = await this.repairBranchCacheForQuery(startId, cached?.branchId); - } - const decoded = decodeEntryRows(await queryCachedBranchRows(this.db, this.metadata.id, cached, query)); - const filtered = decoded.filter( - (entry) => - (query.type === undefined || entry.type === query.type) && - (query.customType === undefined || (entry.type === "custom" && entry.customType === query.customType)), - ); - const entries = query.limit === undefined ? filtered : filtered.slice(0, query.limit); - for (const entry of entries) this.byId.set(entry.id, entry); - return entries; - } - - private async findEntriesOnCanonicalBranch( - query: SessionBranchQuery & { start: string }, - ): Promise { - const rows: SessionEntryRow[] = []; - const visited = new Set(); - let currentId: string | null = query.start; - while (currentId !== null) { - if (visited.has(currentId)) throw invalidSession(`cycle in parent chain at entry ${currentId}`); - visited.add(currentId); - const row: SessionEntryRow | undefined = await this.db - .prepare( - "SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND id = ?", - ) - .get(this.metadata.id, currentId); - if (!row) { - if (currentId === query.start) throw new SessionError("not_found", `Entry ${query.start} not found`); - throw invalidSession(`Entry ${currentId} not found`); - } - rows.push(row); - if (row.id === query.stopAtId || row.type === query.stopAtType) break; - currentId = row.parent_id; - } - const entries = decodeEntryRows(rows).filter( - (entry) => - (query.type === undefined || entry.type === query.type) && - (query.customType === undefined || (entry.type === "custom" && entry.customType === query.customType)), - ); - const limited = query.limit === undefined ? entries : entries.slice(0, query.limit); - for (const entry of limited) this.byId.set(entry.id, entry); - return limited; - } - - private async repairBranchCacheForQuery(leafId: string, branchIdToReplace?: string): Promise { - const visited = new Set(); - let currentId: string | null = leafId; - while (currentId !== null) { - if (visited.has(currentId)) throw invalidSession(`cycle in parent chain at entry ${currentId}`); - visited.add(currentId); - const row: { parent_id: string | null } | undefined = await this.db - .prepare("SELECT parent_id FROM session_entries WHERE session_id = ? AND id = ?") - .get<{ parent_id: string | null }>(this.metadata.id, currentId); - if (!row) { - if (currentId === leafId) throw new SessionError("not_found", `Entry ${leafId} not found`); - throw invalidSession(`Entry ${currentId} not found`); - } - currentId = row.parent_id; - } - try { - await rebuildCachedBranch(this.db, this.metadata.id, leafId, branchIdToReplace); - } catch (error) { - if (error instanceof SessionError) throw error; - throw new SessionError("storage", `Failed to rebuild SQLite branch cache at entry ${leafId}`, toError(error)); - } - const cached = await readCachedBranch(this.db, this.metadata.id, leafId); - if (!cached || !(await isCachedBranchValid(this.db, this.metadata.id, cached, leafId))) { - throw invalidSession(`branch cache repair did not produce a valid path to entry ${leafId}`); - } - return cached; - } - - async readPathToRootOrCompaction(leafId: string | null): Promise { - if (leafId === null) return []; - const cached = await readCachedBranch(this.db, this.metadata.id, leafId); - if (cached) { - const compactionRows = await readCachedEntryRowsByType(this.db, this.metadata.id, cached, "compaction"); - let startSeq = 0; - let expectedStartId: string | null = null; - let pendingStop: { id: string; seq: number } | undefined; - for (const compactionRow of compactionRows) { - if (pendingStop && pendingStop.seq >= compactionRow.entry_seq) { - startSeq = pendingStop.seq; - expectedStartId = pendingStop.id; - break; - } - const compaction = decodeEntryRows([compactionRow])[0]!; - if (compaction.type !== "compaction") throw invalidSession(`entry ${compaction.id} is not a compaction`); - if (compaction.retainedTail) { - startSeq = compactionRow.entry_seq; - expectedStartId = compaction.id; - pendingStop = undefined; - break; - } else if (compaction.firstKeptEntryId !== undefined) { - const firstKeptSeq = await readCachedEntrySeq( - this.db, - this.metadata.id, - cached.branchId, - compaction.firstKeptEntryId, - ); - pendingStop = - firstKeptSeq !== undefined && firstKeptSeq < compactionRow.entry_seq - ? { id: compaction.firstKeptEntryId, seq: firstKeptSeq } - : undefined; - } else { - pendingStop = undefined; - } - } - if (startSeq === 0 && pendingStop) { - startSeq = pendingStop.seq; - expectedStartId = pendingStop.id; - } - const entries = decodeEntryRows(await readCachedBranchRows(this.db, this.metadata.id, cached, startSeq)); - if (this.isValidCachedPath(entries, leafId, expectedStartId)) { - for (const entry of entries) this.byId.set(entry.id, entry); - return entries; - } - } - - const entries = await this.repairBranchCache(leafId, cached?.branchId); - return this.trimPathToRootOrCompaction(entries); - } - - private async repairBranchCache(leafId: string, branchIdToReplace?: string): Promise { - const entries = await this.readCanonicalPathToRoot(leafId); - try { - await rebuildCachedBranch(this.db, this.metadata.id, leafId, branchIdToReplace); - } catch (error) { - if (error instanceof SessionError) throw error; - throw new SessionError("storage", `Failed to rebuild SQLite branch cache at entry ${leafId}`, toError(error)); - } - return entries; - } - - private async readCanonicalPathToRoot(leafId: string): Promise { - const path: SessionTreeEntry[] = []; - let current = await this.readEntry(leafId); - if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`); - const visited = new Set(); - while (current) { - if (visited.has(current.id)) throw invalidSession(`cycle in parent chain at entry ${current.id}`); - visited.add(current.id); - path.push(current); - if (!current.parentId) break; - const parent = await this.readEntry(current.parentId); - if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); - current = parent; - } - return path.reverse(); - } - - private isValidCachedPath( - entries: readonly SessionTreeEntry[], - leafId: string, - expectedStartId: string | null, - ): boolean { - if (entries.length === 0 || entries.at(-1)!.id !== leafId) return false; - if (expectedStartId === null ? entries[0]!.parentId !== null : entries[0]!.id !== expectedStartId) return false; - for (let index = 1; index < entries.length; index++) { - if (entries[index]!.parentId !== entries[index - 1]!.id) return false; - } - return true; - } - - private trimPathToRootOrCompaction(entries: readonly SessionTreeEntry[]): SessionTreeEntry[] { - const path: SessionTreeEntry[] = []; - let stopAtEntryId: string | null = null; - for (let index = entries.length - 1; index >= 0; index--) { - const entry = entries[index]!; - path.push(entry); - if (stopAtEntryId !== null && entry.id === stopAtEntryId) break; - if (entry.type === "compaction") { - if (entry.retainedTail) break; - stopAtEntryId = entry.firstKeptEntryId ?? null; - } - } - return path.reverse(); - } - - private constructor( - db: SqliteDatabase, - metadata: SqliteSessionMetadata, - materializedState: SessionMaterializedState, - ) { - this.db = db; - this.metadata = metadata; - this.byId = new Map(); - this.materializedState = materializedState; - } - - static async open(db: SqliteDatabase, metadata: SqliteSessionMetadata): Promise { - const loaded = await loadSqliteSession(db, metadata.id); - return new SqliteSessionConnection(db, rowToMetadata(loaded.row, metadata.path), loaded.materializedState); - } - - static async create( - db: SqliteDatabase, - path: string, - options: { - cwd: string; - sessionId: string; - parentSessionId?: string; - metadata?: Record; - }, - ): Promise { - const createdAt = new Date().toISOString(); - await db - .prepare( - "INSERT INTO sessions (id, created_at, metadata, cwd, parent_session_id, active_leaf_id) VALUES (?, ?, ?, ?, ?, ?)", - ) - .run( - options.sessionId, - createdAt, - options.metadata === undefined ? null : JSON.stringify(options.metadata), - options.cwd, - options.parentSessionId ?? null, - null, - ); - await db.prepare("INSERT INTO session_sequences (session_id, next_seq) VALUES (?, ?)").run(options.sessionId, 1); - await db - .prepare("INSERT INTO session_materialized (session_id, payload) VALUES (?, ?)") - .run(...materializedStateValues(options.sessionId, createEmptyMaterializedState())); - return new SqliteSessionConnection( - db, - { - id: options.sessionId, - createdAt, - cwd: options.cwd, - path, - parentSessionId: options.parentSessionId, - metadata: options.metadata, - }, - createEmptyMaterializedState(), - ); - } - - async getLabel(id: string): Promise { - return this.materializedState.labelsById.get(id); - } - - async getName(): Promise { - return this.materializedState.name; - } - - async getStats(): Promise { - const { messageCount, cachedTokens, uncachedTokens, totalTokens, costTotal } = this.materializedState; - return { messageCount, cachedTokens, uncachedTokens, totalTokens, costTotal }; - } - - async readHead(): Promise<{ leafId: string | null }> { - const row = await this.db - .prepare( - `SELECT - s.active_leaf_id, - (s.active_leaf_id IS NULL OR EXISTS ( - SELECT 1 FROM session_entries AS e WHERE e.session_id = s.id AND e.id = s.active_leaf_id - )) AS active_leaf_exists - FROM sessions AS s - WHERE s.id = ?`, - ) - .get<{ active_leaf_id: string | null; active_leaf_exists: number }>(this.metadata.id); - if (!row) throw new SessionError("not_found", `Session not found: ${this.metadata.id}`); - if (row.active_leaf_exists === 0) { - throw new SessionError("invalid_session", `Entry ${row.active_leaf_id} not found`); - } - return { leafId: row.active_leaf_id }; - } - - async appendEntry(entry: SessionTreeEntry, options: { transaction?: boolean } = {}): Promise { - if (entry.type === "leaf" && entry.targetId !== null && !(await this.readEntry(entry.targetId))) { - throw new SessionError("not_found", `Entry ${entry.targetId} not found`); - } - const encoded = encodeEntry(entry); - const nextMaterializedState: SessionMaterializedState = { - ...this.materializedState, - labelsById: new Map(this.materializedState.labelsById), - modelThinkingConfigs: [...this.materializedState.modelThinkingConfigs], - currentModel: this.materializedState.currentModel ? { ...this.materializedState.currentModel } : null, - }; - const nextLeafId = leafIdAfterEntry(entry); - try { - applyEntryToMaterializedState(nextMaterializedState, entry); - const write = async () => { - const nextSeq = await getNextSequence(this.db, this.metadata.id); - await this.db - .prepare( - "INSERT INTO session_entries (session_id, id, entry_seq, parent_id, type, timestamp, payload) VALUES (?, ?, ?, ?, ?, ?, ?)", - ) - .run(this.metadata.id, entry.id, nextSeq, entry.parentId, entry.type, entry.timestamp, encoded.payload); - await advanceSequence(this.db, this.metadata.id, nextSeq); - await this.db - .prepare("UPDATE session_materialized SET payload = ? WHERE session_id = ?") - .run(serializeSummary(nextMaterializedState), this.metadata.id); - for (const materializedEntry of entryMaterializedValues(entry)) { - await this.db - .prepare("INSERT INTO entry_materialized (session_id, entry_seq, type, payload) VALUES (?, ?, ?, ?)") - .run(this.metadata.id, nextSeq, materializedEntry.type, materializedEntry.payload); - } - await this.db - .prepare("UPDATE sessions SET active_leaf_id = ? WHERE id = ?") - .run(nextLeafId, this.metadata.id); - await appendEntryToBranchCache( - this.db, - this.metadata.id, - entry.id, - nextSeq, - entry.parentId, - async (parentId) => { - await this.repairBranchCache(parentId); - }, - ); - }; - if (options.transaction === false) await write(); - else await this.db.transaction(write); - this.materializedState = nextMaterializedState; - this.byId.set(entry.id, entry); - } catch (error) { - if (error instanceof SessionError) throw error; - throw new SessionError("storage", `Failed to append SQLite session entry ${entry.id}`, toError(error)); - } - } - - async readEntry(id: string): Promise { - const cached = this.byId.get(id); - if (cached) return cached; - const row = await this.db - .prepare( - "SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND id = ?", - ) - .get(this.metadata.id, id); - if (!row) return undefined; - try { - const entry = decodeEntry(row); - this.byId.set(entry.id, entry); - return entry; - } catch (error) { - throw invalidEntry(`failed to decode entry ${row.id}`, toError(error)); - } - } - - async readEntries(options?: SessionEntryCursorOptions): Promise { - const afterEntrySeq = options?.afterEntrySeq ?? 0; - const rows = - options?.limit === undefined - ? await this.db - .prepare( - "SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND entry_seq > ? ORDER BY entry_seq", - ) - .all(this.metadata.id, afterEntrySeq) - : await this.db - .prepare( - "SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND entry_seq > ? ORDER BY entry_seq LIMIT ?", - ) - .all(this.metadata.id, afterEntrySeq, options.limit); - const entries = decodeEntryRows(rows); - for (const entry of entries) { - this.byId.set(entry.id, entry); - } - return entries; - } -} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/lanes.ts b/packages/storage/sqlite-node/src/sqlite/storage/lanes.ts new file mode 100644 index 000000000..0b8ca7916 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/lanes.ts @@ -0,0 +1,116 @@ +import { SessionError } from "@earendil-works/pi-agent-core/experimental"; +import type { SqliteDatabase } from "../types.ts"; + +export interface LaneRow { + session_id: string; + lane: string; + leaf_id: string | null; +} + +export interface LaneMoveRow { + session_id: string; + seq: number; + lane: string; + leaf_id: string | null; +} + +export function createInitialLane(db: SqliteDatabase, sessionId: string, lane = "main", leafId: string | null = null) { + db.prepare("INSERT INTO lanes (session_id, lane, leaf_id) VALUES (?, ?, ?)").run(sessionId, lane, leafId); +} + +export function readLanes(db: SqliteDatabase, sessionId: string) { + const rows = db + .prepare( + `SELECT + l.session_id, + l.lane, + l.leaf_id, + (l.leaf_id IS NULL OR EXISTS ( + SELECT 1 FROM entries AS e WHERE e.session_id = l.session_id AND e.id = l.leaf_id + )) AS leaf_exists + FROM lanes AS l + WHERE l.session_id = ? + ORDER BY l.lane`, + ) + .all(sessionId); + for (const row of rows) { + if (row.leaf_exists === 0) { + throw new SessionError("storage", `Lane ${row.lane} points at missing entry ${row.leaf_id}`); + } + } + return rows.map(({ session_id, lane, leaf_id }) => ({ session_id, lane, leaf_id })); +} + +export function readLane(db: SqliteDatabase, sessionId: string, lane: string) { + return db + .prepare("SELECT session_id, lane, leaf_id FROM lanes WHERE session_id = ? AND lane = ?") + .get(sessionId, lane); +} + +export function readLaneHead(db: SqliteDatabase, sessionId: string, lane: string) { + const row = db + .prepare( + `SELECT + l.leaf_id, + (l.leaf_id IS NULL OR EXISTS ( + SELECT 1 FROM entries AS e WHERE e.session_id = l.session_id AND e.id = l.leaf_id + )) AS leaf_exists + FROM lanes AS l + WHERE l.session_id = ? AND l.lane = ?`, + ) + .get<{ leaf_id: string | null; leaf_exists: number }>(sessionId, lane); + if (!row) throw new SessionError("invalid_lane", `Lane not found: ${lane}`); + if (row.leaf_exists === 0) throw new SessionError("storage", `Entry ${row.leaf_id} not found`); + return { leafId: row.leaf_id }; +} + +export function createLane(db: SqliteDatabase, sessionId: string, seq: number, lane: string, leafId: string | null) { + db.prepare("INSERT INTO lanes (session_id, lane, leaf_id) VALUES (?, ?, ?)").run(sessionId, lane, leafId); + appendLaneMove(db, sessionId, seq, lane, leafId); +} + +export function moveLane(db: SqliteDatabase, sessionId: string, seq: number, lane: string, leafId: string | null) { + const result = db + .prepare("UPDATE lanes SET leaf_id = ? WHERE session_id = ? AND lane = ?") + .run(leafId, sessionId, lane); + if (result.changes !== 1) throw new SessionError("invalid_lane", `Lane not found: ${lane}`); + appendLaneMove(db, sessionId, seq, lane, leafId); +} + +export function setLaneLeaf(db: SqliteDatabase, sessionId: string, lane: string, leafId: string | null) { + const result = db + .prepare("UPDATE lanes SET leaf_id = ? WHERE session_id = ? AND lane = ?") + .run(leafId, sessionId, lane); + if (result.changes !== 1) throw new SessionError("invalid_lane", `Lane not found: ${lane}`); +} + +export function readLaneMoveRows(db: SqliteDatabase, sessionId: string, options: { afterSeq?: number } = {}) { + const predicates = ["session_id = ?"]; + const params: unknown[] = [sessionId]; + if (options.afterSeq !== undefined) { + predicates.push("seq > ?"); + params.push(options.afterSeq); + } + return db + .prepare( + `SELECT session_id, seq, lane, leaf_id + FROM lane_moves + WHERE ${predicates.join(" AND ")} + ORDER BY seq`, + ) + .all(...params); +} + +export function deleteLaneRows(db: SqliteDatabase, sessionId: string) { + db.prepare("DELETE FROM lane_moves WHERE session_id = ?").run(sessionId); + db.prepare("DELETE FROM lanes WHERE session_id = ?").run(sessionId); +} + +function appendLaneMove(db: SqliteDatabase, sessionId: string, seq: number, lane: string, leafId: string | null) { + db.prepare("INSERT INTO lane_moves (session_id, seq, lane, leaf_id) VALUES (?, ?, ?, ?)").run( + sessionId, + seq, + lane, + leafId, + ); +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/leases.ts b/packages/storage/sqlite-node/src/sqlite/storage/leases.ts new file mode 100644 index 000000000..c00d6c82d --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/leases.ts @@ -0,0 +1,65 @@ +import type { SqliteDatabase } from "../types.ts"; + +export interface SessionLease { + ownerId: string; + fence: number; + expiresAtMs: number; +} + +interface SessionLeaseRow { + owner_id: string; + fence: number; + expires_at_ms: number; +} + +export function acquireSessionLease( + db: SqliteDatabase, + sessionId: string, + ownerId: string, + now: number, + expiresAtMs: number, +) { + const row = db + .prepare( + `INSERT INTO leases (session_id, owner_id, fence, expires_at_ms) + VALUES (?, ?, 1, ?) + ON CONFLICT(session_id) DO UPDATE SET + owner_id = excluded.owner_id, + fence = leases.fence + 1, + expires_at_ms = excluded.expires_at_ms + WHERE leases.expires_at_ms <= ? + RETURNING owner_id, fence, expires_at_ms`, + ) + .get(sessionId, ownerId, expiresAtMs, now); + return row === undefined ? undefined : { ownerId: row.owner_id, fence: row.fence, expiresAtMs: row.expires_at_ms }; +} + +export function renewSessionLease( + db: SqliteDatabase, + sessionId: string, + lease: SessionLease, + now: number, + expiresAtMs: number, +) { + const result = db + .prepare( + `UPDATE leases + SET expires_at_ms = ? + WHERE session_id = ? AND owner_id = ? AND fence = ? AND expires_at_ms > ?`, + ) + .run(expiresAtMs, sessionId, lease.ownerId, lease.fence, now); + if (result.changes === 1) lease.expiresAtMs = expiresAtMs; + return result.changes === 1; +} + +export function releaseSessionLease(db: SqliteDatabase, sessionId: string, lease: SessionLease) { + db.prepare("DELETE FROM leases WHERE session_id = ? AND owner_id = ? AND fence = ?").run( + sessionId, + lease.ownerId, + lease.fence, + ); +} + +export function deleteSessionLease(db: SqliteDatabase, sessionId: string) { + db.prepare("DELETE FROM leases WHERE session_id = ?").run(sessionId); +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/records.ts b/packages/storage/sqlite-node/src/sqlite/storage/records.ts new file mode 100644 index 000000000..c615a5096 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/records.ts @@ -0,0 +1,95 @@ +import type { SqliteDatabase } from "../types.ts"; + +export interface RecordRow { + session_id: string; + seq: number; + id: string; + lane: string; + run_id: string | null; + type: string; + op_kind: string | null; + timestamp: string; + payload: string; +} + +export interface NewRecordRow { + seq: number; + id: string; + lane: string; + runId?: string; + type: string; + opKind?: string; + timestamp: string; + payload: string; +} + +export function appendRecordRow(db: SqliteDatabase, sessionId: string, record: NewRecordRow) { + db.prepare( + `INSERT INTO records + (session_id, seq, id, lane, run_id, type, op_kind, timestamp, payload) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ).run( + sessionId, + record.seq, + record.id, + record.lane, + record.runId ?? null, + record.type, + record.opKind ?? null, + record.timestamp, + record.payload, + ); +} + +export function idExistsInRecords(db: SqliteDatabase, sessionId: string, id: string) { + return !!db + .prepare("SELECT 1 AS found FROM records WHERE session_id = ? AND id = ? LIMIT 1") + .get<{ found: number }>(sessionId, id); +} + +export function deleteRecordRows(db: SqliteDatabase, sessionId: string) { + db.prepare("DELETE FROM records WHERE session_id = ?").run(sessionId); +} + +export function readRecordRows( + db: SqliteDatabase, + sessionId: string, + query: { + lane?: string; + type?: string; + runId?: string; + afterSeq?: number; + order?: "newestFirst" | "oldestFirst"; + limit?: number; + } = {}, +) { + const predicates = ["session_id = ?"]; + const params: unknown[] = [sessionId]; + if (query.lane !== undefined) { + predicates.push("lane = ?"); + params.push(query.lane); + } + if (query.type !== undefined) { + predicates.push("type = ?"); + params.push(query.type); + } + if (query.runId !== undefined) { + predicates.push("run_id = ?"); + params.push(query.runId); + } + if (query.afterSeq !== undefined) { + predicates.push("seq > ?"); + params.push(query.afterSeq); + } + const limit = query.limit === undefined ? "" : " LIMIT ?"; + if (query.limit !== undefined) params.push(query.limit); + const direction = query.order === "oldestFirst" ? "ASC" : "DESC"; + return db + .prepare( + `SELECT session_id, seq, id, lane, run_id, type, op_kind, timestamp, payload + FROM records + WHERE ${predicates.join(" AND ")} + ORDER BY seq ${direction}${limit}`, + ) + .all(...params); +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/session-entries.ts b/packages/storage/sqlite-node/src/sqlite/storage/session-entries.ts deleted file mode 100644 index 1b1cd75a8..000000000 --- a/packages/storage/sqlite-node/src/sqlite/storage/session-entries.ts +++ /dev/null @@ -1,217 +0,0 @@ -import type { SessionTreeEntry, SessionTreeEntryBase } from "@earendil-works/pi-agent-core"; -import { invalidEntry, isRecord } from "./shared.ts"; - -export interface SessionEntryRow { - session_id: string; - id: string; - entry_seq: number; - parent_id: string | null; - type: SessionTreeEntry["type"]; - timestamp: string; - payload: string; -} - -export type EncodedEntry = { - payload: string; -}; - -type EntryPayload = Omit; - -type MessagePayload = EntryPayload>; -type ThinkingLevelChangePayload = EntryPayload>; -type ModelChangePayload = EntryPayload>; -type ActiveToolsChangePayload = EntryPayload>; -type CompactionPayload = EntryPayload>; -type BranchSummaryPayload = EntryPayload>; -type CustomPayload = EntryPayload>; -type CustomMessagePayload = EntryPayload>; -type LabelPayload = EntryPayload>; -type SessionInfoPayload = EntryPayload>; -type LeafPayload = EntryPayload>; - -function parsePayload(row: SessionEntryRow): unknown { - try { - return JSON.parse(row.payload); - } catch (error) { - throw invalidEntry(`entry ${row.id} payload is not valid JSON`, error instanceof Error ? error : undefined); - } -} - -function isTextImageContentArray(value: unknown): boolean { - return ( - Array.isArray(value) && - value.every( - (item) => - isRecord(item) && typeof item.type === "string" && (item.type !== "text" || typeof item.text === "string"), - ) - ); -} - -export function validateSessionTreeEntry(entry: SessionTreeEntry): void { - if (typeof entry.id !== "string" || !entry.id) throw invalidEntry("entry is missing id"); - if (entry.parentId !== null && typeof entry.parentId !== "string") { - throw invalidEntry(`entry ${entry.id} has invalid parentId`); - } - if (typeof entry.timestamp !== "string" || !entry.timestamp) { - throw invalidEntry(`entry ${entry.id} is missing timestamp`); - } - - switch (entry.type) { - case "message": - if (!isRecord(entry.message) || typeof entry.message.role !== "string") { - throw invalidEntry(`entry ${entry.id} is missing message payload`); - } - break; - case "thinking_level_change": - if (typeof entry.thinkingLevel !== "string") throw invalidEntry(`entry ${entry.id} is missing thinkingLevel`); - break; - case "model_change": - if (typeof entry.provider !== "string" || typeof entry.modelId !== "string") { - throw invalidEntry(`entry ${entry.id} has invalid model_change payload`); - } - break; - case "active_tools_change": - if ( - !Array.isArray(entry.activeToolNames) || - entry.activeToolNames.some((value) => typeof value !== "string") - ) { - throw invalidEntry(`entry ${entry.id} has invalid active_tools_change payload`); - } - break; - case "compaction": - if ( - typeof entry.summary !== "string" || - typeof entry.firstKeptEntryId !== "string" || - typeof entry.tokensBefore !== "number" || - (entry.retainedTail !== undefined && !Array.isArray(entry.retainedTail)) - ) { - throw invalidEntry(`entry ${entry.id} has invalid compaction payload`); - } - break; - case "branch_summary": - if (typeof entry.fromId !== "string" || typeof entry.summary !== "string") { - throw invalidEntry(`entry ${entry.id} has invalid branch_summary payload`); - } - break; - case "custom": - if (typeof entry.customType !== "string") throw invalidEntry(`entry ${entry.id} has invalid custom payload`); - break; - case "custom_message": - if ( - typeof entry.customType !== "string" || - typeof entry.display !== "boolean" || - !(typeof entry.content === "string" || isTextImageContentArray(entry.content)) - ) { - throw invalidEntry(`entry ${entry.id} has invalid custom_message payload`); - } - break; - case "label": - if (typeof entry.targetId !== "string" || (entry.label !== undefined && typeof entry.label !== "string")) { - throw invalidEntry(`entry ${entry.id} has invalid label payload`); - } - break; - case "session_info": - if (entry.name !== undefined && typeof entry.name !== "string") { - throw invalidEntry(`entry ${entry.id} has invalid session_info payload`); - } - break; - case "leaf": - if (entry.targetId !== null && typeof entry.targetId !== "string") { - throw invalidEntry(`entry ${entry.id} has invalid leaf payload`); - } - break; - default: { - const exhaustive: never = entry; - throw invalidEntry(`unknown entry type ${(exhaustive as { type?: string }).type ?? "unknown"}`); - } - } -} - -function entryToPayload(entry: TEntry): EntryPayload { - const { type: _type, id: _id, parentId: _parentId, timestamp: _timestamp, ...payload } = entry; - return payload as EntryPayload; -} - -export function encodeEntry(entry: SessionTreeEntry): EncodedEntry { - validateSessionTreeEntry(entry); - return { payload: JSON.stringify(entryToPayload(entry)) }; -} - -export function decodeEntry(row: SessionEntryRow): SessionTreeEntry { - const payload = parsePayload(row); - if (!isRecord(payload)) throw invalidEntry(`entry ${row.id} payload is not an object`); - const base = { - id: row.id, - parentId: row.parent_id, - timestamp: row.timestamp, - }; - - switch (row.type) { - case "message": { - if (!("message" in payload)) throw invalidEntry(`entry ${row.id} is missing message payload`); - const messagePayload = payload as MessagePayload; - return { ...base, type: "message", ...messagePayload }; - } - case "thinking_level_change": - if (typeof payload.thinkingLevel !== "string") throw invalidEntry(`entry ${row.id} is missing thinkingLevel`); - return { ...base, type: "thinking_level_change", ...(payload as ThinkingLevelChangePayload) }; - case "model_change": - if (typeof payload.provider !== "string" || typeof payload.modelId !== "string") { - throw invalidEntry(`entry ${row.id} has invalid model_change payload`); - } - return { ...base, type: "model_change", ...(payload as ModelChangePayload) }; - case "active_tools_change": - if ( - !Array.isArray(payload.activeToolNames) || - payload.activeToolNames.some((value) => typeof value !== "string") - ) { - throw invalidEntry(`entry ${row.id} has invalid active_tools_change payload`); - } - return { ...base, type: "active_tools_change", ...(payload as ActiveToolsChangePayload) }; - case "compaction": - if ( - typeof payload.summary !== "string" || - typeof payload.firstKeptEntryId !== "string" || - typeof payload.tokensBefore !== "number" || - (payload.retainedTail !== undefined && !Array.isArray(payload.retainedTail)) - ) { - throw invalidEntry(`entry ${row.id} has invalid compaction payload`); - } - return { ...base, type: "compaction", ...(payload as CompactionPayload) }; - case "branch_summary": - if (typeof payload.fromId !== "string" || typeof payload.summary !== "string") { - throw invalidEntry(`entry ${row.id} has invalid branch_summary payload`); - } - return { ...base, type: "branch_summary", ...(payload as BranchSummaryPayload) }; - case "custom": - if (typeof payload.customType !== "string") throw invalidEntry(`entry ${row.id} has invalid custom payload`); - return { ...base, type: "custom", ...(payload as CustomPayload) }; - case "custom_message": - if ( - typeof payload.customType !== "string" || - typeof payload.display !== "boolean" || - !("content" in payload) - ) { - throw invalidEntry(`entry ${row.id} has invalid custom_message payload`); - } - return { ...base, type: "custom_message", ...(payload as CustomMessagePayload) }; - case "label": - if (typeof payload.targetId !== "string") throw invalidEntry(`entry ${row.id} has invalid label payload`); - if (payload.label !== undefined && typeof payload.label !== "string") { - throw invalidEntry(`entry ${row.id} has invalid label payload`); - } - return { ...base, type: "label", ...(payload as LabelPayload) }; - case "session_info": - if (payload.name !== undefined && typeof payload.name !== "string") { - throw invalidEntry(`entry ${row.id} has invalid session_info payload`); - } - return { ...base, type: "session_info", ...(payload as SessionInfoPayload) }; - case "leaf": - if (payload.targetId !== null && typeof payload.targetId !== "string") { - throw invalidEntry(`entry ${row.id} has invalid leaf payload`); - } - return { ...base, type: "leaf", ...(payload as LeafPayload) }; - default: - throw invalidEntry(`unknown entry type ${row.type}`); - } -} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts b/packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts deleted file mode 100644 index e46b543d3..000000000 --- a/packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts +++ /dev/null @@ -1,355 +0,0 @@ -import type { SessionTreeEntry, ThinkingLevel } from "@earendil-works/pi-agent-core"; -import { invalidSession, isRecord } from "./shared.ts"; - -export interface SessionMaterializedRow { - session_id: string; - payload: string; -} - -export interface EntryMaterializedRow { - session_id: string; - entry_seq: number; - type: string; - payload: string; -} - -export interface ModelThinkingConfig { - provider: string; - modelId: string; - thinkingLevel: ThinkingLevel; -} - -export interface SessionMaterializedState { - name: string | undefined; - messageCount: number; - cachedTokens: number; - uncachedTokens: number; - totalTokens: number; - costTotal: number; - labelsById: Map; - modelThinkingConfigs: ModelThinkingConfig[]; - currentModel: { provider: string; modelId: string } | null; - currentThinkingLevel: ThinkingLevel | null; -} - -interface SessionMaterializedSummary { - name?: string; - messageCount: number; - cachedTokens: number; - uncachedTokens: number; - totalTokens: number; - costTotal: number; - currentModel?: { provider: string; modelId: string } | null; - currentThinkingLevel?: ThinkingLevel | null; -} - -function compareModelThinkingConfig(left: ModelThinkingConfig, right: ModelThinkingConfig): number { - return ( - left.provider.localeCompare(right.provider) || - left.modelId.localeCompare(right.modelId) || - left.thinkingLevel.localeCompare(right.thinkingLevel) - ); -} - -function normalizeModelThinkingConfigs(configs: readonly ModelThinkingConfig[]): ModelThinkingConfig[] { - const unique = new Map(); - for (const config of configs) { - unique.set(`${config.provider}\u0000${config.modelId}\u0000${config.thinkingLevel}`, config); - } - return [...unique.values()].sort(compareModelThinkingConfig); -} - -function addModelThinkingConfig( - state: SessionMaterializedState, - provider: string, - modelId: string, - thinkingLevel: ThinkingLevel, -): void { - state.modelThinkingConfigs = normalizeModelThinkingConfigs([ - ...state.modelThinkingConfigs, - { provider, modelId, thinkingLevel }, - ]); -} - -export function isThinkingLevel(value: unknown): value is ThinkingLevel { - return ( - value === "off" || - value === "minimal" || - value === "low" || - value === "medium" || - value === "high" || - value === "xhigh" - ); -} - -function getAssistantUsage(message: unknown): - | { - provider: string; - modelId: string; - input: number; - output: number; - cacheRead: number; - cacheWrite: number; - costTotal: number; - } - | undefined { - if (!isRecord(message) || message.role !== "assistant") return undefined; - if (typeof message.provider !== "string" || typeof message.model !== "string") return undefined; - if (!isRecord(message.usage) || !isRecord(message.usage.cost)) return undefined; - const { input, output, cacheRead, cacheWrite } = message.usage; - const costTotal = message.usage.cost.total; - if ( - typeof input !== "number" || - typeof output !== "number" || - typeof cacheRead !== "number" || - typeof cacheWrite !== "number" || - typeof costTotal !== "number" - ) { - return undefined; - } - return { - provider: message.provider, - modelId: message.model, - input, - output, - cacheRead, - cacheWrite, - costTotal, - }; -} - -export function createEmptyMaterializedState(): SessionMaterializedState { - return { - name: undefined, - messageCount: 0, - cachedTokens: 0, - uncachedTokens: 0, - totalTokens: 0, - costTotal: 0, - labelsById: new Map(), - modelThinkingConfigs: [], - currentModel: null, - currentThinkingLevel: null, - }; -} - -export function applyEntryToMaterializedState(state: SessionMaterializedState, entry: SessionTreeEntry): void { - switch (entry.type) { - case "session_info": - state.name = entry.name?.trim() || undefined; - break; - case "label": { - const label = entry.label?.trim(); - if (label) { - state.labelsById.set(entry.targetId, label); - } else { - state.labelsById.delete(entry.targetId); - } - break; - } - case "model_change": - state.currentModel = { provider: entry.provider, modelId: entry.modelId }; - if (state.currentThinkingLevel) { - addModelThinkingConfig(state, entry.provider, entry.modelId, state.currentThinkingLevel); - } - break; - case "thinking_level_change": - if (!isThinkingLevel(entry.thinkingLevel)) break; - state.currentThinkingLevel = entry.thinkingLevel; - if (state.currentModel) { - addModelThinkingConfig(state, state.currentModel.provider, state.currentModel.modelId, entry.thinkingLevel); - } - break; - case "message": { - state.messageCount += 1; - const usage = getAssistantUsage(entry.message); - if (!usage) break; - state.cachedTokens += usage.cacheRead; - state.uncachedTokens += usage.input + usage.cacheWrite; - state.totalTokens += usage.input + usage.output + usage.cacheRead + usage.cacheWrite; - state.costTotal += usage.costTotal; - state.currentModel = { provider: usage.provider, modelId: usage.modelId }; - if (state.currentThinkingLevel) { - addModelThinkingConfig(state, usage.provider, usage.modelId, state.currentThinkingLevel); - } - break; - } - case "compaction": - case "branch_summary": { - const usage = entry.usage; - if ( - !isRecord(usage) || - !isRecord(usage.cost) || - typeof usage.input !== "number" || - typeof usage.output !== "number" || - typeof usage.cacheRead !== "number" || - typeof usage.cacheWrite !== "number" || - typeof usage.cost.total !== "number" - ) { - break; - } - state.cachedTokens += usage.cacheRead; - state.uncachedTokens += usage.input + usage.cacheWrite; - state.totalTokens += usage.input + usage.output + usage.cacheRead + usage.cacheWrite; - state.costTotal += usage.cost.total; - break; - } - case "active_tools_change": - case "custom": - case "custom_message": - case "leaf": - break; - default: { - const exhaustive: never = entry; - void exhaustive; - break; - } - } -} - -export function serializeSummary(state: SessionMaterializedState): string { - const summary: SessionMaterializedSummary = { - name: state.name, - messageCount: state.messageCount, - cachedTokens: state.cachedTokens, - uncachedTokens: state.uncachedTokens, - totalTokens: state.totalTokens, - costTotal: state.costTotal, - currentModel: state.currentModel, - currentThinkingLevel: state.currentThinkingLevel, - }; - return JSON.stringify(summary); -} - -function parseSummary(json: string): SessionMaterializedSummary { - let parsed: unknown; - try { - parsed = JSON.parse(json); - } catch (error) { - throw invalidSession( - `materialized session summary is not valid JSON`, - error instanceof Error ? error : undefined, - ); - } - if (!isRecord(parsed) || Array.isArray(parsed)) { - throw invalidSession("materialized session summary is not an object"); - } - const currentModel = parsed.currentModel; - const currentThinkingLevel = parsed.currentThinkingLevel; - if ( - (parsed.name !== undefined && typeof parsed.name !== "string") || - typeof parsed.messageCount !== "number" || - typeof parsed.cachedTokens !== "number" || - typeof parsed.uncachedTokens !== "number" || - typeof parsed.totalTokens !== "number" || - typeof parsed.costTotal !== "number" || - (currentModel !== undefined && - currentModel !== null && - (!isRecord(currentModel) || - typeof currentModel.provider !== "string" || - typeof currentModel.modelId !== "string")) || - (currentThinkingLevel !== undefined && currentThinkingLevel !== null && !isThinkingLevel(currentThinkingLevel)) - ) { - throw invalidSession("materialized session summary has invalid fields"); - } - return { - name: parsed.name?.trim() || undefined, - messageCount: parsed.messageCount, - cachedTokens: parsed.cachedTokens, - uncachedTokens: parsed.uncachedTokens, - totalTokens: parsed.totalTokens, - costTotal: parsed.costTotal, - currentModel: - currentModel && isRecord(currentModel) - ? { provider: currentModel.provider as string, modelId: currentModel.modelId as string } - : (currentModel ?? undefined), - currentThinkingLevel: (currentThinkingLevel as ThinkingLevel | null | undefined) ?? undefined, - }; -} - -function parseEntryMaterializedPayload(row: EntryMaterializedRow): unknown { - try { - return JSON.parse(row.payload); - } catch (error) { - throw invalidSession( - `materialized entry row ${row.entry_seq} is not valid JSON`, - error instanceof Error ? error : undefined, - ); - } -} - -export function materializedStateFromRows( - summaryRow: SessionMaterializedRow, - entryRows: EntryMaterializedRow[], -): SessionMaterializedState { - const summary = parseSummary(summaryRow.payload); - const state: SessionMaterializedState = { - name: summary.name, - messageCount: summary.messageCount, - cachedTokens: summary.cachedTokens, - uncachedTokens: summary.uncachedTokens, - totalTokens: summary.totalTokens, - costTotal: summary.costTotal, - labelsById: new Map(), - modelThinkingConfigs: [], - currentModel: summary.currentModel ?? null, - currentThinkingLevel: summary.currentThinkingLevel ?? null, - }; - for (const row of entryRows) { - const payload = parseEntryMaterializedPayload(row); - if (!isRecord(payload)) throw invalidSession(`materialized entry row ${row.entry_seq} is not an object`); - if (row.type === "label") { - if (typeof payload.targetId !== "string") { - throw invalidSession(`materialized label row ${row.entry_seq} is missing targetId`); - } - if (payload.label !== null && payload.label !== undefined && typeof payload.label !== "string") { - throw invalidSession(`materialized label row ${row.entry_seq} has invalid label`); - } - const label = typeof payload.label === "string" ? payload.label.trim() : ""; - if (label) { - state.labelsById.set(payload.targetId, label); - } else { - state.labelsById.delete(payload.targetId); - } - } - } - return state; -} - -export function materializedStateValues( - sessionId: string, - state: SessionMaterializedState, -): [sessionId: string, payload: string] { - return [sessionId, serializeSummary(state)]; -} - -export function entryMaterializedValues( - entry: SessionTreeEntry, -): Array<{ type: EntryMaterializedRow["type"]; payload: string }> { - switch (entry.type) { - case "label": - return [ - { - type: "label", - payload: JSON.stringify({ targetId: entry.targetId, label: entry.label ?? null }), - }, - ]; - case "model_change": - case "thinking_level_change": - case "message": - return []; - case "active_tools_change": - case "branch_summary": - case "compaction": - case "custom": - case "custom_message": - case "leaf": - case "session_info": - return []; - default: { - const exhaustive: never = entry; - void exhaustive; - return []; - } - } -} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/session-sequences.ts b/packages/storage/sqlite-node/src/sqlite/storage/session-sequences.ts index a38bb668f..04798f531 100644 --- a/packages/storage/sqlite-node/src/sqlite/storage/session-sequences.ts +++ b/packages/storage/sqlite-node/src/sqlite/storage/session-sequences.ts @@ -1,16 +1,28 @@ +import { SessionError } from "@earendil-works/pi-agent-core/experimental"; import type { SqliteDatabase } from "../types.ts"; -import { invalidSession } from "./shared.ts"; -export async function getNextSequence(db: SqliteDatabase, sessionId: string): Promise { - const sequenceRow = await db +export function createSequence(db: SqliteDatabase, sessionId: string, nextSeq = 1) { + db.prepare("INSERT INTO session_sequences (session_id, next_seq) VALUES (?, ?)").run(sessionId, nextSeq); +} + +export function getNextSequence(db: SqliteDatabase, sessionId: string) { + const sequenceRow = db .prepare("SELECT next_seq FROM session_sequences WHERE session_id = ?") .get<{ next_seq: number }>(sessionId); if (!sequenceRow) { - throw invalidSession(`missing sequence row for session ${sessionId}`); + throw new SessionError("storage", `Missing sequence row for session ${sessionId}`); } return sequenceRow.next_seq; } -export async function advanceSequence(db: SqliteDatabase, sessionId: string, nextSeq: number): Promise { - await db.prepare("UPDATE session_sequences SET next_seq = ? WHERE session_id = ?").run(nextSeq + 1, sessionId); +export function setNextSequence(db: SqliteDatabase, sessionId: string, nextSeq: number) { + db.prepare("UPDATE session_sequences SET next_seq = ? WHERE session_id = ?").run(nextSeq, sessionId); +} + +export function advanceSequence(db: SqliteDatabase, sessionId: string, seq: number) { + setNextSequence(db, sessionId, seq + 1); +} + +export function deleteSequence(db: SqliteDatabase, sessionId: string) { + db.prepare("DELETE FROM session_sequences WHERE session_id = ?").run(sessionId); } diff --git a/packages/storage/sqlite-node/src/sqlite/storage/session-stats.ts b/packages/storage/sqlite-node/src/sqlite/storage/session-stats.ts new file mode 100644 index 000000000..ce180a0fc --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/session-stats.ts @@ -0,0 +1,63 @@ +import { SessionError, type SessionStats } from "@earendil-works/pi-agent-core/experimental"; +import type { Usage } from "@earendil-works/pi-ai"; +import type { SqliteDatabase } from "../types.ts"; + +export interface SessionStatsRow { + session_id: string; + message_count: number; + cached_tokens: number; + uncached_tokens: number; + total_tokens: number; + cost_total: number; +} + +export function createStats(db: SqliteDatabase, sessionId: string): void { + db.prepare( + `INSERT INTO session_stats + (session_id, message_count, cached_tokens, uncached_tokens, total_tokens, cost_total) + VALUES (?, 0, 0, 0, 0, 0)`, + ).run(sessionId); +} + +export function readStats(db: SqliteDatabase, sessionId: string): SessionStats { + const row = db + .prepare( + `SELECT session_id, message_count, cached_tokens, uncached_tokens, total_tokens, cost_total + FROM session_stats + WHERE session_id = ?`, + ) + .get(sessionId); + if (!row) throw new SessionError("storage", `Missing stats row for session ${sessionId}`); + return { + messageCount: row.message_count, + cachedTokens: row.cached_tokens, + uncachedTokens: row.uncached_tokens, + totalTokens: row.total_tokens, + costTotal: row.cost_total, + }; +} + +export function incrementMessageCount(db: SqliteDatabase, sessionId: string): void { + const result = db + .prepare("UPDATE session_stats SET message_count = message_count + 1 WHERE session_id = ?") + .run(sessionId); + if (result.changes !== 1) throw new SessionError("storage", `Missing stats row for session ${sessionId}`); +} + +export function addUsageToStats(db: SqliteDatabase, sessionId: string, usage: Usage): void { + const result = db + .prepare( + `UPDATE session_stats + SET cached_tokens = cached_tokens + ?, + uncached_tokens = uncached_tokens + ?, + total_tokens = total_tokens + ?, + cost_total = cost_total + ? + WHERE session_id = ?`, + ) + .run(usage.cacheRead, usage.input + usage.cacheWrite, usage.totalTokens, usage.cost.total, sessionId); + if (result.changes !== 1) throw new SessionError("storage", `Missing stats row for session ${sessionId}`); +} + +export function deleteStats(db: SqliteDatabase, sessionId: string): void { + db.prepare("DELETE FROM session_stats WHERE session_id = ?").run(sessionId); +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/sessions.ts b/packages/storage/sqlite-node/src/sqlite/storage/sessions.ts index 2cb4d4959..3ea1db3c4 100644 --- a/packages/storage/sqlite-node/src/sqlite/storage/sessions.ts +++ b/packages/storage/sqlite-node/src/sqlite/storage/sessions.ts @@ -1,5 +1,5 @@ -import { SessionError } from "@earendil-works/pi-agent-core"; -import type { SqliteSessionMetadata } from "../types.ts"; +import { assertJsonSerializable, SessionError } from "@earendil-works/pi-agent-core/experimental"; +import type { SqliteDatabase, SqliteSessionMetadata } from "../types.ts"; export interface SessionRow { id: string; @@ -7,7 +7,14 @@ export interface SessionRow { metadata: string | null; cwd: string; parent_session_id: string | null; - active_leaf_id: string | null; +} + +export interface NewSessionRow { + id: string; + createdAt: string; + cwd: string; + parentSessionId?: string; + metadata?: Record; } function parseMetadata(metadata: string | null, sessionId: string): Record | undefined { @@ -17,17 +24,62 @@ function parseMetadata(metadata: string | null, sessionId: string): Record; } +export function sessionExists(db: SqliteDatabase, sessionId: string) { + return !!db.prepare("SELECT 1 AS found FROM sessions WHERE id = ?").get<{ found: number }>(sessionId); +} + +function serializeMetadata(metadata: Record | undefined): string | null { + if (metadata === undefined) return null; + if (typeof metadata !== "object" || metadata === null || Array.isArray(metadata)) { + throw new SessionError("invalid_payload", "SQLite session metadata must be an object"); + } + assertJsonSerializable(metadata); + return JSON.stringify(metadata); +} + +export function insertSessionRow(db: SqliteDatabase, session: NewSessionRow) { + db.prepare("INSERT INTO sessions (id, created_at, metadata, cwd, parent_session_id) VALUES (?, ?, ?, ?, ?)").run( + session.id, + session.createdAt, + serializeMetadata(session.metadata), + session.cwd, + session.parentSessionId ?? null, + ); +} + +export function readSessionRow(db: SqliteDatabase, sessionId: string) { + return db + .prepare("SELECT id, created_at, metadata, cwd, parent_session_id FROM sessions WHERE id = ?") + .get(sessionId); +} + +export function readSessionRows(db: SqliteDatabase, options: { cwd?: string } = {}) { + return options.cwd + ? db + .prepare( + "SELECT id, created_at, metadata, cwd, parent_session_id FROM sessions WHERE cwd = ? ORDER BY created_at DESC", + ) + .all(options.cwd) + : db + .prepare("SELECT id, created_at, metadata, cwd, parent_session_id FROM sessions ORDER BY created_at DESC") + .all(); +} + +export function deleteSessionRow(db: SqliteDatabase, sessionId: string) { + db.prepare("DELETE FROM sessions WHERE id = ?").run(sessionId); +} + export function rowToMetadata(row: SessionRow, path: string): SqliteSessionMetadata { return { id: row.id, diff --git a/packages/storage/sqlite-node/src/sqlite/storage/shared.ts b/packages/storage/sqlite-node/src/sqlite/storage/shared.ts deleted file mode 100644 index 704b06507..000000000 --- a/packages/storage/sqlite-node/src/sqlite/storage/shared.ts +++ /dev/null @@ -1,17 +0,0 @@ -import type { SessionTreeEntry } from "@earendil-works/pi-agent-core"; -import { SessionError } from "@earendil-works/pi-agent-core"; -export function isRecord(value: unknown): value is Record { - return typeof value === "object" && value !== null; -} - -export function invalidSession(message: string, cause?: Error): SessionError { - return new SessionError("invalid_session", `Invalid SQLite session: ${message}`, cause); -} - -export function invalidEntry(message: string, cause?: Error): SessionError { - return new SessionError("invalid_entry", `Invalid SQLite session entry: ${message}`, cause); -} - -export function leafIdAfterEntry(entry: SessionTreeEntry): string | null { - return entry.type === "leaf" ? entry.targetId : entry.id; -} diff --git a/packages/storage/sqlite-node/src/sqlite/types.ts b/packages/storage/sqlite-node/src/sqlite/types.ts index 5624ccc43..800443784 100644 --- a/packages/storage/sqlite-node/src/sqlite/types.ts +++ b/packages/storage/sqlite-node/src/sqlite/types.ts @@ -10,17 +10,18 @@ export interface SqliteRunResult { /** Prepared SQLite statement capability used by the SQLite session backend. */ export interface SqliteStatement { - run(...params: unknown[]): Promise; - get(...params: unknown[]): Promise; - all(...params: unknown[]): Promise; + run(...params: unknown[]): SqliteRunResult; + get(...params: unknown[]): TRow | undefined; + all(...params: unknown[]): TRow[]; } /** SQLite database capability used by the SQLite session backend. */ export interface SqliteDatabase { - exec(sql: string): Promise; + exec(sql: string): void; prepare(sql: string): SqliteStatement; - transaction(fn: () => Promise): Promise; - close(): Promise; + /** Runs a synchronous write transaction. The callback must not return a promise. */ + transaction(fn: () => T): T; + close(): void; } export interface SqliteDatabaseFactory { diff --git a/tsconfig.json b/tsconfig.json index a409cdf11..c139f41c8 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -11,6 +11,7 @@ "@earendil-works/pi-ai/*": ["./packages/ai/src/*.ts", "./packages/ai/src/providers/*.ts"], "@earendil-works/pi-ai/dist/*": ["./packages/ai/src/*"], "@earendil-works/pi-agent-core": ["./packages/agent/src/index.ts"], + "@earendil-works/pi-agent-core/experimental": ["./packages/agent/src/experimental.ts"], "@earendil-works/pi-agent-core/*": ["./packages/agent/src/*"], "@earendil-works/pi-agent-sqlite-node": ["./packages/storage/sqlite-node/src/index.ts"], "@earendil-works/pi-coding-agent": ["./packages/coding-agent/src/index.ts"],