File size: 1,839 Bytes
f500658 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | import { BACKGROUND_CONTEXT } from "@earendil-works/pi-agent-core";
import {
type ConformanceCase,
createStorageConformance,
type StorageFixture,
} from "@earendil-works/pi-agent-core/harness/session/testing";
import { describe, it } from "vitest";
import { createNodeSqliteFactory, SQLITE_STORAGE_VERSION, SqliteStorage, sql } from "../src/index.ts";
import { applyInitialSchema } from "../src/sqlite/migrations.ts";
const SESSION_ID = "session";
const NOW = 1_700_000_000_000;
const EMPTY_USAGE = {
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 0,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
};
function registerConformance(name: string, cases: readonly ConformanceCase[]): void {
describe(name, () => {
for (const group of new Set(cases.map((testCase) => testCase.group))) {
describe(group, () => {
for (const testCase of cases.filter((candidate) => candidate.group === group)) {
it(testCase.name, () => testCase.run());
}
});
}
});
}
registerConformance(
"SqliteStorage conformance",
createStorageConformance(async () => {
const db = await createNodeSqliteFactory().open(":memory:");
try {
await applyInitialSchema(db);
sql`INSERT INTO sessions
(id, created_at, parent_session_id, storage_version, metadata, message_count, usage_payload, next_seq)
VALUES (${SESSION_ID}, ${NOW}, ${null}, ${SQLITE_STORAGE_VERSION}, ${null}, ${0}, ${JSON.stringify(EMPTY_USAGE)}, ${1})`.run(
db,
);
const storage = new SqliteStorage(db, { sessionId: SESSION_ID, now: () => NOW });
return {
storage,
async [Symbol.asyncDispose]() {
try {
await storage.close(BACKGROUND_CONTEXT);
} finally {
db.close();
}
},
} satisfies StorageFixture;
} catch (error) {
db.close();
throw error;
}
}),
);
|