diff --git a/packages/session-backends/sqlite-node/src/sqlite/session/session-sequences.ts b/packages/session-backends/sqlite-node/src/sqlite/session/session-sequences.ts new file mode 100644 index 000000000..3f20dfe1f --- /dev/null +++ b/packages/session-backends/sqlite-node/src/sqlite/session/session-sequences.ts @@ -0,0 +1,13 @@ +import { sql } from "../sql.ts"; +import type { SqliteDatabase } from "../types.ts"; + +export function readNextSeq(db: SqliteDatabase): number { + const rows = sql`SELECT next_seq FROM session`.all<{ next_seq: number }>(db); + if (rows.length !== 1) throw new Error(`Expected exactly one session row, found ${rows.length}`); + return rows[0]!.next_seq; +} + +export function advanceNextSeq(db: SqliteDatabase, nextSeq: number): void { + const result = sql`UPDATE session SET next_seq = ${nextSeq}`.run(db); + if (result.changes !== 1) throw new Error(`Expected to update exactly one session row, updated ${result.changes}`); +} diff --git a/packages/session-backends/sqlite-node/test/storage.test.ts b/packages/session-backends/sqlite-node/test/storage.test.ts index 0ee2d53cf..08e24b769 100644 --- a/packages/session-backends/sqlite-node/test/storage.test.ts +++ b/packages/session-backends/sqlite-node/test/storage.test.ts @@ -1,6 +1,8 @@ +import { prepareStorageCommit } from "@earendil-works/pi-agent-core"; import { describe, expect, it } from "vitest"; import { createNodeSqliteFactory, type SqliteDatabase, SqliteStorage, sql } from "../src/index.ts"; import { applyInitialSchema } from "../src/sqlite/migrations.ts"; +import { advanceNextSeq, readNextSeq } from "../src/sqlite/session/session-sequences.ts"; async function withStorage(run: (storage: SqliteStorage, db: SqliteDatabase) => Promise): Promise { const db = await createNodeSqliteFactory().open(":memory:"); @@ -238,6 +240,63 @@ describe("SqliteStorage", () => { }); }); + it("prepares committed writes with assigned sequences and timestamp", () => { + const prepared = prepareStorageCommit( + { + writes: [ + { + kind: "entry", + entry: { + id: "entry", + parentId: null, + type: "message", + message: { role: "user", content: "hi", timestamp: 1 }, + }, + }, + { + kind: "usage", + row: { + id: "usage", + usage: { + input: 1, + output: 2, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 3, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + adjustment: false, + }, + }, + { kind: "register", op: "set", namespace: "fact.name", key: "", value: "name" }, + { kind: "register", op: "delete", namespace: "fact.label", key: "entry" }, + ], + }, + 7, + 1_700_000_000_000, + ); + + expect(prepared.result).toEqual({ firstSeq: 7, seqs: [7, 8, 9, 10], timestamp: 1_700_000_000_000 }); + expect(prepared.writes).toMatchObject([ + { kind: "entry", id: "entry", seq: 7, timestamp: 1_700_000_000_000 }, + { kind: "usage", id: "usage", seq: 8 }, + { kind: "register", op: "set", namespace: "fact.name", key: "", seq: 9, value: "name" }, + { kind: "register", op: "delete", namespace: "fact.label", key: "entry", seq: 10 }, + ]); + }); + + it("reads and advances the next commit sequence", async () => { + await withStorage(async (_storage, db) => { + sql`INSERT INTO session + (created_at, parent_session_id, storage_version, metadata, message_count, usage_payload, next_seq) + VALUES (${1}, ${null}, ${1}, ${null}, ${0}, ${JSON.stringify({})}, ${7})`.run(db); + + expect(readNextSeq(db)).toBe(7); + advanceNextSeq(db, 10); + expect(readNextSeq(db)).toBe(10); + }); + }); + it("gets maintained session stats", async () => { await withStorage(async (storage, db) => { const usage = {