From b3bf0a2a8199ed6dcbd6a45e61952e4f6633e0f3 Mon Sep 17 00:00:00 2001 From: Christian Klotz Date: Wed, 12 Aug 2026 00:22:30 +0300 Subject: [PATCH] feat(agent): add storage backend conformance suite --- packages/agent/package.json | 4 + .../harness/session/testing/conformance.ts | 677 ++++++++++++++++++ .../src/harness/session/testing/index.ts | 6 + .../src/harness/session/testing/types.ts | 16 + .../agent/test/harness/memory-storage.test.ts | 515 ++----------- tsconfig.json | 3 + 6 files changed, 753 insertions(+), 468 deletions(-) create mode 100644 packages/agent/src/harness/session/testing/conformance.ts create mode 100644 packages/agent/src/harness/session/testing/index.ts create mode 100644 packages/agent/src/harness/session/testing/types.ts diff --git a/packages/agent/package.json b/packages/agent/package.json index 7fb9bde12..4c011eaf7 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -14,6 +14,10 @@ "types": "./dist/node.d.ts", "import": "./dist/node.js" }, + "./session/testing": { + "types": "./dist/harness/session/testing/index.d.ts", + "import": "./dist/harness/session/testing/index.js" + }, "./package.json": "./package.json" }, "files": [ diff --git a/packages/agent/src/harness/session/testing/conformance.ts b/packages/agent/src/harness/session/testing/conformance.ts new file mode 100644 index 000000000..d456e685d --- /dev/null +++ b/packages/agent/src/harness/session/testing/conformance.ts @@ -0,0 +1,677 @@ +import { deepStrictEqual, ok, rejects, strictEqual } from "node:assert/strict"; +import type { Usage } from "@earendil-works/pi-ai"; +import type { CompactionEntry, CustomEntry, Entry, MessageEntry, NewEntry, Register, Transaction } from "../types.ts"; +import type { StorageConformanceCase, StorageFixture, StorageFixtureFactory } from "./types.ts"; + +const MESSAGE_TIMESTAMP = 1_650_000_000_000; + +type ConformanceTest = (fixture: StorageFixture) => Promise; + +function createCase( + factory: StorageFixtureFactory, + group: string, + name: string, + test: ConformanceTest, +): StorageConformanceCase { + return { + group, + name, + async run() { + await using fixture = await factory(); + await test(fixture); + }, + }; +} + +function usage(input: number, output: number, options: { cacheWrite1h?: number; reasoning?: number } = {}): Usage { + return { + input, + output, + cacheRead: input + 1, + cacheWrite: output + 1, + ...(options.cacheWrite1h === undefined ? {} : { cacheWrite1h: options.cacheWrite1h }), + ...(options.reasoning === undefined ? {} : { reasoning: options.reasoning }), + totalTokens: input + output, + cost: { + input: input / 100, + output: output / 100, + cacheRead: (input + 1) / 100, + cacheWrite: (output + 1) / 100, + total: (input + output + 2) / 100, + }, + }; +} + +function zeroUsage(): Usage { + return { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +function userEntry(id: string, parentId: string | null = null, text = id): NewEntry { + return { + id, + parentId, + type: "message", + message: { + role: "user", + content: [{ type: "text", text }], + timestamp: MESSAGE_TIMESTAMP, + }, + }; +} + +function customEntry( + id: string, + parentId: string | null, + customType = "note", + data: CustomEntry["data"] = { id }, +): NewEntry { + return { id, parentId, type: "custom", customType, data }; +} + +function compactionEntry(id: string, parentId: string | null): NewEntry { + return { + id, + parentId, + type: "compaction", + summary: `summary:${id}`, + retainedTail: [], + tokensBefore: 10, + fromHook: false, + }; +} + +function ids(entries: readonly Entry[]): string[] { + return entries.map((entry) => entry.id); +} + +function assertStrictlyIncreasing(values: readonly number[]): void { + for (let index = 1; index < values.length; index++) { + ok(values[index - 1]! < values[index]!, `Expected ${values.join(", ")} to be strictly increasing`); + } +} + +function sortRegisters(registers: Register[]): Register[] { + return [...registers].sort((left, right) => left.key.localeCompare(right.key)); +} + +/** Creates fresh, runner-independent cases for the durable Storage contract. */ +export function createStorageConformance(factory: StorageFixtureFactory): readonly StorageConformanceCase[] { + return [ + createCase(factory, "transactions", "commits mixed writes atomically in write order", async ({ storage }) => { + const result = await storage.commit({ + writes: [ + { kind: "entry", entry: userEntry("entry") }, + { kind: "register", op: "set", namespace: "fact.name", key: "", value: "session" }, + { + kind: "usage", + row: { id: "usage", usage: usage(2, 3), adjustment: false, entryId: "entry" }, + }, + ], + }); + + strictEqual(result.seqs.length, 3); + strictEqual(result.firstSeq, result.seqs[0]); + assertStrictlyIncreasing(result.seqs); + ok(Number.isSafeInteger(result.timestamp) && result.timestamp >= 0); + deepStrictEqual( + await storage.getEntries(["entry"]), + new Map([["entry", { ...userEntry("entry"), seq: result.seqs[0], timestamp: result.timestamp }]]), + ); + deepStrictEqual(await storage.getRegister("fact.name", ""), { + namespace: "fact.name", + key: "", + value: "session", + seq: result.seqs[1], + }); + deepStrictEqual(await storage.scanUsage({ order: "asc" }), [ + { id: "usage", seq: result.seqs[2], usage: usage(2, 3), adjustment: false, entryId: "entry" }, + ]); + }), + + createCase( + factory, + "transactions", + "rolls back every store when a mixed transaction fails", + async ({ storage }) => { + await storage.commit({ + writes: [ + { kind: "entry", entry: userEntry("root") }, + { kind: "usage", row: { id: "taken", usage: usage(1, 1), adjustment: false } }, + ], + }); + const entriesBefore = await storage.scanEntries({ order: "asc" }); + const usageBefore = await storage.scanUsage({ order: "asc" }); + const statsBefore = await storage.getStats(); + + await rejects( + storage.commit({ + writes: [ + { kind: "register", op: "set", namespace: "fact.name", key: "", value: "transient" }, + { kind: "entry", entry: customEntry("transient-entry", "root") }, + { kind: "usage", row: { id: "transient-usage", usage: usage(5, 8), adjustment: true } }, + { kind: "entry", entry: customEntry("taken", "root") }, + ], + }), + ); + + deepStrictEqual(await storage.scanEntries({ order: "asc" }), entriesBefore); + deepStrictEqual(await storage.scanUsage({ order: "asc" }), usageBefore); + deepStrictEqual(await storage.getStats(), statsBefore); + strictEqual(await storage.getRegister("fact.name", ""), undefined); + }, + ), + + createCase(factory, "transactions", "enforces one shared entry and usage id namespace", async ({ storage }) => { + await storage.commit({ + writes: [ + { kind: "entry", entry: userEntry("existing-entry") }, + { kind: "usage", row: { id: "existing-usage", usage: usage(1, 1), adjustment: false } }, + ], + }); + + await rejects( + storage.commit({ + writes: [{ kind: "usage", row: { id: "existing-entry", usage: usage(2, 2), adjustment: false } }], + }), + ); + await rejects(storage.commit({ writes: [{ kind: "entry", entry: customEntry("existing-usage", null) }] })); + + for (const [id, writes] of [ + [ + "entry-then-usage", + [ + { kind: "entry", entry: customEntry("entry-then-usage", null) }, + { + kind: "usage", + row: { id: "entry-then-usage", usage: usage(3, 3), adjustment: false }, + }, + ], + ], + [ + "usage-then-entry", + [ + { + kind: "usage", + row: { id: "usage-then-entry", usage: usage(4, 4), adjustment: false }, + }, + { kind: "entry", entry: customEntry("usage-then-entry", null) }, + ], + ], + ] satisfies Array<[string, Transaction["writes"]]>) { + await rejects(storage.commit({ writes }), `Expected duplicate id ${id} to reject`); + } + + deepStrictEqual(ids(await storage.scanEntries({ order: "asc" })), ["existing-entry"]); + deepStrictEqual( + (await storage.scanUsage({ order: "asc" })).map((row) => row.id), + ["existing-usage"], + ); + }), + + createCase( + factory, + "transactions", + "resolves parents only from prior entries and earlier writes", + async ({ storage }) => { + await storage.commit({ writes: [{ kind: "entry", entry: userEntry("root") }] }); + await storage.commit({ + writes: [ + { kind: "entry", entry: customEntry("child", "root") }, + { kind: "entry", entry: customEntry("grandchild", "child") }, + ], + }); + deepStrictEqual(ids(await storage.scanBranch({ start: "grandchild", order: "oldestFirst" })), [ + "root", + "child", + "grandchild", + ]); + + await rejects( + storage.commit({ + writes: [ + { kind: "entry", entry: customEntry("before-parent", "later-parent") }, + { kind: "entry", entry: customEntry("later-parent", "root") }, + { kind: "register", op: "set", namespace: "fact.label", key: "before-parent", value: "transient" }, + ], + }), + ); + await rejects(storage.commit({ writes: [{ kind: "entry", entry: customEntry("orphan", "missing") }] })); + await storage.commit({ + writes: [{ kind: "usage", row: { id: "usage-is-not-parent", usage: usage(1, 1), adjustment: false } }], + }); + await rejects( + storage.commit({ + writes: [{ kind: "entry", entry: customEntry("usage-child", "usage-is-not-parent") }], + }), + ); + + deepStrictEqual( + await storage.getEntries(["before-parent", "later-parent", "orphan", "usage-child"]), + new Map(), + ); + strictEqual(await storage.getRegister("fact.label", "before-parent"), undefined); + }, + ), + + createCase( + factory, + "registers", + "sets, replaces, deletes, and recreates registers without tombstones", + async ({ storage }) => { + const first = await storage.commit({ + writes: [ + { kind: "register", op: "set", namespace: "fact.custom", key: "prefix/b", value: 1 }, + { kind: "register", op: "set", namespace: "fact.custom", key: "prefix/a", value: 2 }, + { kind: "register", op: "set", namespace: "fact.custom", key: "other", value: 3 }, + { kind: "register", op: "set", namespace: "fact.custom", key: "prefix/a", value: null }, + ], + }); + deepStrictEqual(await storage.getRegister("fact.custom", "prefix/a"), { + namespace: "fact.custom", + key: "prefix/a", + value: null, + seq: first.seqs[3], + }); + + const second = await storage.commit({ + writes: [ + { kind: "register", op: "delete", namespace: "fact.custom", key: "prefix/a" }, + { kind: "register", op: "delete", namespace: "fact.custom", key: "absent" }, + { kind: "register", op: "set", namespace: "fact.custom", key: "prefix/a", value: "recreated" }, + ], + }); + + deepStrictEqual(sortRegisters(await storage.listRegisters("fact.custom", "prefix/")), [ + { namespace: "fact.custom", key: "prefix/a", value: "recreated", seq: second.seqs[2] }, + { namespace: "fact.custom", key: "prefix/b", value: 1, seq: first.seqs[0] }, + ]); + strictEqual(await storage.getRegister("fact.custom", "absent"), undefined); + }, + ), + + createCase( + factory, + "registers", + "does not change historical stores during register-only commits", + async ({ storage }) => { + await storage.commit({ + writes: [ + { kind: "entry", entry: userEntry("root") }, + { kind: "usage", row: { id: "historical-usage", usage: usage(2, 3), adjustment: false } }, + ], + }); + const entriesBefore = await storage.scanEntries({ order: "asc" }); + const usageBefore = await storage.scanUsage({ order: "asc" }); + const statsBefore = await storage.getStats(); + + const result = await storage.commit({ + writes: [ + { kind: "register", op: "set", namespace: "fact.name", key: "", value: "first" }, + { kind: "register", op: "set", namespace: "fact.name", key: "", value: "second" }, + ], + }); + + deepStrictEqual(await storage.scanEntries({ order: "asc" }), entriesBefore); + deepStrictEqual(await storage.scanUsage({ order: "asc" }), usageBefore); + deepStrictEqual(await storage.getStats(), statsBefore); + deepStrictEqual(await storage.getRegister("fact.name", ""), { + namespace: "fact.name", + key: "", + value: "second", + seq: result.seqs[1], + }); + }, + ), + + createCase(factory, "entry queries", "stores custom entries with and without data", async ({ storage }) => { + const result = await storage.commit({ + writes: [ + { + kind: "entry", + entry: { id: "without-data", parentId: null, type: "custom", customType: "marker" }, + }, + { + kind: "entry", + entry: customEntry("with-data", "without-data", "note", { nested: [1, 2] }), + }, + ], + }); + + deepStrictEqual( + await storage.getEntries(["without-data", "with-data"]), + new Map([ + [ + "without-data", + { + id: "without-data", + parentId: null, + type: "custom", + customType: "marker", + seq: result.seqs[0], + timestamp: result.timestamp, + }, + ], + [ + "with-data", + { + ...customEntry("with-data", "without-data", "note", { nested: [1, 2] }), + seq: result.seqs[1], + timestamp: result.timestamp, + }, + ], + ]), + ); + }), + + createCase( + factory, + "entry queries", + "scans global entries with explicit ranges, filters, orders, and limits", + async ({ storage }) => { + const result = await storage.commit({ + writes: [ + { kind: "entry", entry: userEntry("root") }, + { kind: "entry", entry: customEntry("note-1", "root", "note") }, + { kind: "entry", entry: customEntry("other", "note-1", "other") }, + { kind: "entry", entry: customEntry("note-2", "other", "note") }, + { kind: "entry", entry: userEntry("tail", "note-2") }, + ], + }); + + deepStrictEqual( + ids( + await storage.scanEntries({ + type: "custom", + customType: "note", + fromSeq: result.seqs[1], + toSeq: result.seqs[3], + order: "desc", + }), + ), + ["note-2", "note-1"], + ); + deepStrictEqual(ids(await storage.scanEntries({ order: "asc", limit: 2 })), ["root", "note-1"]); + deepStrictEqual(ids(await storage.scanEntries({ order: "desc", limit: 2 })), ["tail", "note-2"]); + }, + ), + + createCase( + factory, + "branch queries", + "applies stops before filters and cursors before limits", + async ({ storage }) => { + const result = await storage.commit({ + writes: [ + { kind: "entry", entry: userEntry("root") }, + { kind: "entry", entry: customEntry("marker", "root", "marker") }, + { kind: "entry", entry: userEntry("middle", "marker") }, + { kind: "entry", entry: compactionEntry("compact", "middle") }, + { kind: "entry", entry: customEntry("note", "compact", "note") }, + { kind: "entry", entry: userEntry("leaf", "note") }, + ], + }); + + deepStrictEqual( + ids(await storage.scanBranch({ start: "leaf", stopAtType: "compaction", type: "message" })), + ["leaf"], + ); + deepStrictEqual( + ids( + await storage.scanBranch({ + start: "leaf", + order: "oldestFirst", + stopAtId: "middle", + type: "custom", + }), + ), + ["marker"], + ); + deepStrictEqual( + ids( + await storage.scanBranch({ + start: "leaf", + order: "newestFirst", + cursor: { seq: result.seqs[4] }, + limit: 2, + }), + ), + ["compact", "middle"], + ); + deepStrictEqual( + ids( + await storage.scanBranch({ + start: "leaf", + order: "oldestFirst", + cursor: { seq: result.seqs[1] }, + limit: 2, + }), + ), + ["middle", "compact"], + ); + deepStrictEqual(ids(await storage.scanBranch({ start: "leaf", stopAtId: "leaf", type: "custom" })), []); + deepStrictEqual(ids(await storage.scanBranch({ start: "leaf", customType: "note" })), ["note"]); + await rejects(storage.scanBranch({ start: "missing" })); + }, + ), + + createCase(factory, "branch queries", "returns branch structure without payload fields", async ({ storage }) => { + const result = await storage.commit({ + writes: [ + { kind: "entry", entry: userEntry("root") }, + { kind: "entry", entry: customEntry("child", "root", "note") }, + ], + }); + + deepStrictEqual(await storage.scanBranchStructure({ start: "child", order: "oldestFirst" }), [ + { + id: "root", + parentId: null, + seq: result.seqs[0], + timestamp: result.timestamp, + type: "message", + }, + { + id: "child", + parentId: "root", + seq: result.seqs[1], + timestamp: result.timestamp, + type: "custom", + customType: "note", + }, + ]); + }), + + createCase( + factory, + "usage and stats", + "scans the usage ledger with explicit ranges, orders, and limits", + async ({ storage }) => { + const result = await storage.commit({ + writes: [ + { kind: "usage", row: { id: "usage-1", usage: usage(1, 1), adjustment: false } }, + { kind: "register", op: "set", namespace: "fact.name", key: "", value: "sequence gap" }, + { kind: "usage", row: { id: "usage-2", usage: usage(2, 2), adjustment: false } }, + { kind: "usage", row: { id: "usage-3", usage: usage(3, 3), adjustment: true } }, + ], + }); + + deepStrictEqual( + (await storage.scanUsage({ fromSeq: result.seqs[1], toSeq: result.seqs[2], order: "asc" })).map( + (row) => row.id, + ), + ["usage-2"], + ); + deepStrictEqual( + (await storage.scanUsage({ order: "desc", limit: 2 })).map((row) => row.id), + ["usage-3", "usage-2"], + ); + deepStrictEqual( + (await storage.scanUsage({ order: "asc", limit: 2 })).map((row) => row.id), + ["usage-1", "usage-2"], + ); + }, + ), + + createCase( + factory, + "usage and stats", + "keeps stats equal to message count and ledger totals", + async ({ storage }) => { + deepStrictEqual(await storage.getStats(), { messageCount: 0, usage: zeroUsage() }); + + const firstUsage = usage(2, 3, { cacheWrite1h: 4, reasoning: 1 }); + await storage.commit({ + writes: [ + { kind: "entry", entry: userEntry("message") }, + { kind: "usage", row: { id: "usage-1", usage: firstUsage, adjustment: false } }, + ], + }); + deepStrictEqual(await storage.getStats(), { messageCount: 1, usage: firstUsage }); + + const secondUsage = usage(5, 7, { cacheWrite1h: 6, reasoning: 2 }); + await storage.commit({ + writes: [ + { kind: "entry", entry: customEntry("custom", "message") }, + { kind: "entry", entry: compactionEntry("compaction", "custom") }, + { kind: "usage", row: { id: "usage-2", usage: secondUsage, adjustment: true } }, + ], + }); + deepStrictEqual(await storage.getStats(), { + messageCount: 1, + usage: { + input: 7, + output: 10, + cacheRead: 9, + cacheWrite: 12, + cacheWrite1h: 10, + reasoning: 3, + totalTokens: 17, + cost: { + input: firstUsage.cost.input + secondUsage.cost.input, + output: firstUsage.cost.output + secondUsage.cost.output, + cacheRead: firstUsage.cost.cacheRead + secondUsage.cost.cacheRead, + cacheWrite: firstUsage.cost.cacheWrite + secondUsage.cost.cacheWrite, + total: firstUsage.cost.total + secondUsage.cost.total, + }, + }, + }); + }, + ), + + createCase( + factory, + "immutability", + "detaches nested transaction inputs and read results", + async ({ storage }) => { + const entryData = { nested: { values: [1, 2] } }; + const registerValue = { nested: [3, 4] }; + const details = { source: { name: "original" } }; + const transaction: Transaction = { + writes: [ + { kind: "entry", entry: customEntry("entry", null, "nested", entryData) }, + { kind: "register", op: "set", namespace: "fact.custom", key: "object", value: registerValue }, + { kind: "usage", row: { id: "usage", usage: usage(4, 5), adjustment: false, details } }, + ], + }; + await storage.commit(transaction); + entryData.nested.values[0] = 99; + registerValue.nested[0] = 99; + details.source.name = "mutated"; + + const entryMap = await storage.getEntries(["entry"]); + const entry = entryMap.get("entry"); + ok(entry?.type === "custom"); + (entry.data as { nested: { values: number[] } }).nested.values[0] = 88; + (entryMap as Map).clear(); + const register = await storage.getRegister("fact.custom", "object"); + ok(register !== undefined); + (register.value as { nested: number[] }).nested[0] = 88; + const rows = await storage.scanUsage({ order: "asc" }); + (rows[0]!.details as { source: { name: string } }).source.name = "changed"; + rows[0]!.usage.input = 88; + const scannedEntries = await storage.scanEntries({ order: "asc" }); + (scannedEntries[0] as CustomEntry).data = { replaced: true }; + const branchEntries = await storage.scanBranch({ start: "entry" }); + (branchEntries[0] as CustomEntry).data = { replaced: true }; + const structures = await storage.scanBranchStructure({ start: "entry" }); + structures[0]!.timestamp = 0; + const stats = await storage.getStats(); + stats.usage.input = 88; + + deepStrictEqual((await storage.getEntries(["entry"])).get("entry"), { + ...customEntry("entry", null, "nested", { nested: { values: [1, 2] } }), + seq: entry.seq, + timestamp: entry.timestamp, + }); + deepStrictEqual(await storage.getRegister("fact.custom", "object"), { + namespace: "fact.custom", + key: "object", + value: { nested: [3, 4] }, + seq: register.seq, + }); + deepStrictEqual(await storage.scanUsage({ order: "asc" }), [ + { + id: "usage", + seq: rows[0]!.seq, + usage: usage(4, 5), + adjustment: false, + details: { source: { name: "original" } }, + }, + ]); + deepStrictEqual((await storage.scanEntries({ order: "asc" }))[0]?.type, "custom"); + deepStrictEqual((await storage.scanBranch({ start: "entry" }))[0]?.type, "custom"); + strictEqual((await storage.scanBranchStructure({ start: "entry" }))[0]?.timestamp, entry.timestamp); + strictEqual((await storage.getStats()).usage.input, 4); + }, + ), + + createCase( + factory, + "serialization", + "serializes back-to-back commits in admission order", + async ({ storage }) => { + const first = storage.commit({ writes: [{ kind: "entry", entry: userEntry("first") }] }); + const second = storage.commit({ writes: [{ kind: "entry", entry: userEntry("second", "first") }] }); + const [firstResult, secondResult] = await Promise.all([first, second]); + + ok(firstResult.seqs[0]! < secondResult.seqs[0]!); + deepStrictEqual(ids(await storage.scanEntries({ order: "asc" })), ["first", "second"]); + }, + ), + + createCase( + factory, + "lifecycle", + "seals admission, drains admitted commits, and closes idempotently", + async ({ storage }) => { + const admitted = storage.commit({ writes: [{ kind: "entry", entry: userEntry("admitted") }] }); + const firstClose = storage.close(); + const secondClose = storage.close(); + + await rejects(storage.getStats()); + await rejects(storage.commit({ writes: [] })); + strictEqual((await admitted).seqs.length, 1); + await Promise.all([firstClose, secondClose]); + + const rejectedReads = [ + storage.getEntries([]), + storage.getRegister("fact.name", ""), + storage.listRegisters("fact.name"), + storage.scanBranch({ start: "admitted" }), + storage.scanBranchStructure({ start: "admitted" }), + storage.scanEntries({ order: "asc" }), + storage.scanUsage({ order: "asc" }), + storage.getStats(), + ]; + for (const read of rejectedReads) await rejects(read); + }, + ), + ]; +} diff --git a/packages/agent/src/harness/session/testing/index.ts b/packages/agent/src/harness/session/testing/index.ts new file mode 100644 index 000000000..bf7eccbed --- /dev/null +++ b/packages/agent/src/harness/session/testing/index.ts @@ -0,0 +1,6 @@ +export { createStorageConformance } from "./conformance.ts"; +export type { + StorageConformanceCase, + StorageFixture, + StorageFixtureFactory, +} from "./types.ts"; diff --git a/packages/agent/src/harness/session/testing/types.ts b/packages/agent/src/harness/session/testing/types.ts new file mode 100644 index 000000000..da8dc70d1 --- /dev/null +++ b/packages/agent/src/harness/session/testing/types.ts @@ -0,0 +1,16 @@ +import type { Storage } from "../types.ts"; + +/** A fresh backend storage instance owned by one conformance case. */ +export interface StorageFixture extends AsyncDisposable { + readonly storage: Storage; +} + +/** Creates an isolated storage fixture for one conformance case. */ +export type StorageFixtureFactory = () => Promise; + +/** A runner-independent conformance case that can be registered with any test framework. */ +export interface StorageConformanceCase { + readonly group: string; + readonly name: string; + run(): Promise; +} diff --git a/packages/agent/test/harness/memory-storage.test.ts b/packages/agent/test/harness/memory-storage.test.ts index 3f2137562..5fca182aa 100644 --- a/packages/agent/test/harness/memory-storage.test.ts +++ b/packages/agent/test/harness/memory-storage.test.ts @@ -1,474 +1,53 @@ -import type { Usage } from "@earendil-works/pi-ai"; import { describe, expect, it } from "vitest"; -import { MemoryStorage, type NewEntry, type Transaction } from "../../src/harness/session/index.ts"; +import { MemoryStorage } from "../../src/harness/session/index.ts"; +import { createStorageConformance, type StorageFixture } from "../../src/harness/session/testing/index.ts"; const NOW = 1_700_000_000_000; -function usage(input: number, output: number): Usage { - return { - input, - output, - cacheRead: 0, - cacheWrite: 0, - totalTokens: input + output, - cost: { input: input / 100, output: output / 100, cacheRead: 0, cacheWrite: 0, total: (input + output) / 100 }, - }; -} - -function messageEntry(id: string, parentId: string | null = null, text = id): NewEntry { - return { - id, - parentId, - type: "message", - message: { role: "user", content: text, timestamp: NOW - 1 }, - }; -} - -function customEntry(id: string, parentId: string | null, customType: string): NewEntry { - return { id, parentId, type: "custom", customType, data: { id } }; -} - -function compactionEntry(id: string, parentId: string | null): NewEntry { - return { - id, - parentId, - type: "compaction", - summary: id, - retainedTail: [], - tokensBefore: 10, - fromHook: false, - }; -} - -describe("MemoryStorage", () => { - it("commits entries, registers, and usage as one transaction", async () => { - const storage = new MemoryStorage({ now: () => NOW }); - const result = await storage.commit({ - writes: [ - { kind: "entry", entry: messageEntry("entry") }, - { kind: "register", op: "set", namespace: "fact.name", key: "", value: "session" }, - { kind: "usage", row: { id: "usage", usage: usage(2, 3), adjustment: false, entryId: "entry" } }, - ], - }); - - expect(result).toEqual({ firstSeq: 1, seqs: [1, 2, 3], timestamp: NOW }); - expect(await storage.getEntries(["entry"])).toEqual( - new Map([["entry", { ...messageEntry("entry"), seq: 1, timestamp: NOW }]]), - ); - expect(await storage.getRegister("fact.name", "")).toEqual({ - namespace: "fact.name", - key: "", - value: "session", - seq: 2, - }); - expect(await storage.scanUsage({})).toEqual([ - { id: "usage", seq: 3, usage: usage(2, 3), adjustment: false, entryId: "entry" }, - ]); - }); - - it("rolls back every write when an entry or usage id already exists", async () => { - const storage = new MemoryStorage({ now: () => NOW }); - await storage.commit({ writes: [{ kind: "entry", entry: messageEntry("taken") }] }); - - await expect( - storage.commit({ - writes: [ - { kind: "register", op: "set", namespace: "fact.name", key: "", value: "must roll back" }, - { kind: "usage", row: { id: "taken", usage: usage(1, 1), adjustment: true } }, - ], - }), - ).rejects.toThrow("Duplicate entry or usage id: taken"); - expect(await storage.getRegister("fact.name", "")).toBeUndefined(); - expect(await storage.scanUsage({})).toEqual([]); - - await expect( - storage.commit({ writes: [{ kind: "entry", entry: messageEntry("after-rollback") }] }), - ).resolves.toEqual({ firstSeq: 2, seqs: [2], timestamp: NOW }); - }); - it("sets, replaces, deletes, and recreates registers without tombstones", async () => { - const storage = new MemoryStorage({ now: () => NOW }); - await storage.commit({ - writes: [ - { kind: "register", op: "set", namespace: "fact.custom", key: "prefix/b", value: 1 }, - { kind: "register", op: "set", namespace: "fact.custom", key: "prefix/a", value: 2 }, - { kind: "register", op: "set", namespace: "fact.custom", key: "other", value: 3 }, - { kind: "register", op: "set", namespace: "fact.custom", key: "prefix/a", value: null }, - ], - }); - await storage.commit({ - writes: [ - { kind: "register", op: "delete", namespace: "fact.custom", key: "prefix/a" }, - { kind: "register", op: "delete", namespace: "fact.custom", key: "absent" }, - { kind: "register", op: "set", namespace: "fact.custom", key: "prefix/a", value: "recreated" }, - ], - }); - - expect(await storage.listRegisters("fact.custom", "prefix/")).toEqual([ - { namespace: "fact.custom", key: "prefix/a", value: "recreated", seq: 7 }, - { namespace: "fact.custom", key: "prefix/b", value: 1, seq: 1 }, - ]); - expect(await storage.getRegister("fact.custom", "absent")).toBeUndefined(); - }); - - it("accepts parents inserted earlier in the transaction and rejects missing parents atomically", async () => { - const storage = new MemoryStorage({ now: () => NOW }); - await storage.commit({ - writes: [ - { kind: "entry", entry: messageEntry("root") }, - { kind: "entry", entry: messageEntry("child", "root") }, - ], - }); - expect([...(await storage.getEntries(["child"])).values()][0]?.parentId).toBe("root"); - - await expect( - storage.commit({ - writes: [ - { kind: "entry", entry: messageEntry("orphan", "missing") }, - { kind: "register", op: "set", namespace: "fact.label", key: "orphan", value: "rolled back" }, - ], - }), - ).rejects.toThrow("Missing parent entry: missing"); - expect(await storage.getEntries(["orphan"])).toEqual(new Map()); - expect(await storage.getRegister("fact.label", "orphan")).toBeUndefined(); - }); - - it("scans global entries by inclusive sequence range, filters, order, and limit", async () => { - const storage = new MemoryStorage({ now: () => NOW }); - await storage.commit({ - writes: [ - { kind: "entry", entry: messageEntry("root") }, - { kind: "entry", entry: customEntry("note-1", "root", "note") }, - { kind: "entry", entry: customEntry("other", "note-1", "other") }, - { kind: "entry", entry: customEntry("note-2", "other", "note") }, - { kind: "entry", entry: messageEntry("tail", "note-2") }, - ], - }); - - expect( - (await storage.scanEntries({ type: "custom", customType: "note", fromSeq: 2, toSeq: 4, order: "desc" })).map( - (entry) => entry.id, - ), - ).toEqual(["note-2", "note-1"]); - expect((await storage.scanEntries({ order: "asc", limit: 2 })).map((entry) => entry.id)).toEqual([ - "root", - "note-1", - ]); - expect((await storage.scanEntries({})).map((entry) => entry.id)).toEqual([ - "root", - "note-1", - "other", - "note-2", - "tail", - ]); - expect(await storage.scanEntries({ limit: 0 })).toEqual([]); - }); - - it("scans a branch with inclusive stops before filters, exclusive cursors, ordering, and limits", async () => { - const storage = new MemoryStorage({ now: () => NOW }); - await storage.commit({ - writes: [ - { kind: "entry", entry: messageEntry("root") }, - { kind: "entry", entry: customEntry("marker", "root", "marker") }, - { kind: "entry", entry: messageEntry("middle", "marker") }, - { kind: "entry", entry: compactionEntry("compact", "middle") }, - { kind: "entry", entry: customEntry("note", "compact", "note") }, - { kind: "entry", entry: messageEntry("leaf", "note") }, - ], - }); - - expect( - (await storage.scanBranch({ start: "leaf", stopAtType: "compaction", type: "message" })).map( - (entry) => entry.id, - ), - ).toEqual(["leaf"]); - expect( - ( - await storage.scanBranch({ - start: "leaf", - order: "oldestFirst", - stopAtId: "middle", - type: "custom", - }) - ).map((entry) => entry.id), - ).toEqual(["marker"]); - expect( - (await storage.scanBranch({ start: "leaf", cursor: { seq: 5 }, order: "newestFirst", limit: 2 })).map( - (entry) => entry.id, - ), - ).toEqual(["compact", "middle"]); - expect(await storage.scanBranch({ start: "leaf", customType: "note" })).toMatchObject([{ id: "note" }]); - await expect(storage.scanBranch({ start: "missing" })).rejects.toThrow("Unknown branch start: missing"); - }); - - it("projects branch structure without entry payloads", async () => { - const storage = new MemoryStorage({ now: () => NOW }); - await storage.commit({ - writes: [ - { kind: "entry", entry: messageEntry("root") }, - { kind: "entry", entry: customEntry("child", "root", "note") }, - ], - }); - - expect(await storage.scanBranchStructure({ start: "child", order: "oldestFirst" })).toEqual([ - { id: "root", parentId: null, seq: 1, timestamp: NOW, type: "message" }, - { id: "child", parentId: "root", seq: 2, timestamp: NOW, type: "custom", customType: "note" }, - ]); - }); - - it("keeps stats equal to message entries and the complete usage ledger after every commit", async () => { - const storage = new MemoryStorage({ now: () => NOW }); - expect(await storage.getStats()).toEqual({ messageCount: 0, usage: usage(0, 0) }); - - const firstUsage = { ...usage(2, 3), cacheWrite1h: 4, reasoning: 1 }; - await storage.commit({ - writes: [ - { kind: "entry", entry: messageEntry("message") }, - { kind: "usage", row: { id: "usage-1", usage: firstUsage, adjustment: false } }, - ], - }); - expect(await storage.getStats()).toEqual({ messageCount: 1, usage: firstUsage }); - - const secondUsage = { ...usage(5, 7), cacheWrite1h: 6, reasoning: 2 }; - await storage.commit({ - writes: [ - { kind: "entry", entry: customEntry("custom", "message", "note") }, - { kind: "usage", row: { id: "usage-2", usage: secondUsage, adjustment: true } }, - { kind: "register", op: "set", namespace: "fact.name", key: "", value: "unchanged totals" }, - ], - }); - expect(await storage.getStats()).toEqual({ - messageCount: 1, - usage: { - input: 7, - output: 10, - cacheRead: 0, - cacheWrite: 0, - cacheWrite1h: 10, - reasoning: 3, - totalTokens: 17, - cost: { - input: firstUsage.cost.input + secondUsage.cost.input, - output: firstUsage.cost.output + secondUsage.cost.output, - cacheRead: 0, - cacheWrite: 0, - total: firstUsage.cost.total + secondUsage.cost.total, - }, - }, - }); - }); - - it("serializes reentrant commits in admission order", async () => { - let storage: MemoryStorage; - let secondCommit: ReturnType | undefined; - let clockCalls = 0; - storage = new MemoryStorage({ - now: () => { - clockCalls++; - if (clockCalls === 1) { - secondCommit = storage.commit({ writes: [{ kind: "entry", entry: messageEntry("second") }] }); - } - return NOW + clockCalls; - }, - }); - - const firstCommit = storage.commit({ writes: [{ kind: "entry", entry: messageEntry("first") }] }); - await expect(firstCommit).resolves.toEqual({ firstSeq: 1, seqs: [1], timestamp: NOW + 1 }); - await expect(secondCommit).resolves.toEqual({ firstSeq: 2, seqs: [2], timestamp: NOW + 2 }); - expect((await storage.scanEntries({ order: "asc" })).map((entry) => entry.id)).toEqual(["first", "second"]); - }); - - it("seals admission, drains admitted commits, and closes idempotently", async () => { - const storage = new MemoryStorage({ now: () => NOW }); - const admitted = storage.commit({ writes: [{ kind: "entry", entry: messageEntry("admitted") }] }); - const firstClose = storage.close(); - const secondClose = storage.close(); - - await expect(storage.getStats()).rejects.toThrow("MemoryStorage is closed"); - await expect(storage.commit({ writes: [] })).rejects.toThrow("MemoryStorage is closed"); - await expect(admitted).resolves.toEqual({ firstSeq: 1, seqs: [1], timestamp: NOW }); - await expect(Promise.all([firstClose, secondClose])).resolves.toEqual([undefined, undefined]); - - const rejectedReads = [ - storage.getEntries([]), - storage.getRegister("fact.name", ""), - storage.listRegisters("fact.name"), - storage.scanBranch({ start: "admitted" }), - storage.scanBranchStructure({ start: "admitted" }), - storage.scanEntries({}), - storage.scanUsage({}), - storage.getStats(), - ]; - await Promise.all(rejectedReads.map((read) => expect(read).rejects.toThrow("MemoryStorage is closed"))); - }); - - it.each([ - [ - "entry then usage", - [ - { kind: "entry", entry: messageEntry("shared") }, - { kind: "usage", row: { id: "shared", usage: usage(1, 1), adjustment: false } }, - ], - ], - [ - "usage then entry", - [ - { kind: "usage", row: { id: "shared", usage: usage(1, 1), adjustment: false } }, - { kind: "entry", entry: messageEntry("shared") }, - ], - ], - ] satisfies Array<[string, Transaction["writes"]]>)( - "rejects duplicate ids in the shared entry and usage namespace: %s", - async (_name, writes) => { - const storage = new MemoryStorage({ now: () => NOW }); - await expect(storage.commit({ writes })).rejects.toThrow("Duplicate entry or usage id: shared"); - expect(await storage.getEntries(["shared"])).toEqual(new Map()); - expect(await storage.scanUsage({})).toEqual([]); - }, - ); - - it("scans usage by inclusive sequence range, order, and limit", async () => { - const storage = new MemoryStorage({ now: () => NOW }); - await storage.commit({ - writes: [ - { kind: "usage", row: { id: "usage-1", usage: usage(1, 1), adjustment: false } }, - { kind: "register", op: "set", namespace: "fact.name", key: "", value: "sequence gap" }, - { kind: "usage", row: { id: "usage-2", usage: usage(2, 2), adjustment: false } }, - { kind: "usage", row: { id: "usage-3", usage: usage(3, 3), adjustment: true } }, - ], - }); - - expect((await storage.scanUsage({ fromSeq: 2, toSeq: 3 })).map((row) => row.id)).toEqual(["usage-2"]); - expect((await storage.scanUsage({ order: "desc", limit: 2 })).map((row) => row.id)).toEqual([ - "usage-3", - "usage-2", - ]); - expect((await storage.scanUsage({})).map((row) => row.id)).toEqual(["usage-1", "usage-2", "usage-3"]); - expect(await storage.scanUsage({ limit: 0 })).toEqual([]); - }); - - it("detaches transaction inputs and every read result", async () => { - const storage = new MemoryStorage({ now: () => NOW }); - const transaction: Transaction = { - writes: [ - { kind: "entry", entry: messageEntry("entry", null, "original") }, - { kind: "register", op: "set", namespace: "fact.custom", key: "object", value: { nested: [1] } }, - { - kind: "usage", - row: { id: "usage", usage: usage(4, 5), adjustment: false, details: { source: "original" } }, - }, - ], - }; - const commit = storage.commit(transaction); - const entryWrite = transaction.writes[0]; - if (entryWrite?.kind === "entry" && entryWrite.entry.type === "message") { - entryWrite.entry.message.timestamp = 0; - } - const registerWrite = transaction.writes[1]; - if ( - registerWrite?.kind === "register" && - registerWrite.op === "set" && - registerWrite.namespace === "fact.custom" - ) { - registerWrite.value = { nested: [999] }; - } - const usageWrite = transaction.writes[2]; - if (usageWrite?.kind === "usage") { - usageWrite.row.usage.input = 999; - usageWrite.row.details = { source: "mutated" }; - } - await commit; - - const entries = await storage.getEntries(["entry"]); - const readEntry = entries.get("entry"); - if (readEntry?.type === "message") readEntry.message.timestamp = 0; - (entries as Map).clear(); - const register = await storage.getRegister("fact.custom", "object"); - if (register && typeof register.value === "object" && register.value !== null && !Array.isArray(register.value)) { - register.value.nested = [999]; - } - const registers = await storage.listRegisters("fact.custom"); - registers.length = 0; - const rows = await storage.scanUsage({}); - rows[0]!.usage.input = 999; - const scanned = await storage.scanEntries({}); - if (scanned[0]?.type === "message") scanned[0].message.timestamp = 0; - const branch = await storage.scanBranch({ start: "entry" }); - if (branch[0]?.type === "message") branch[0].message.timestamp = 0; - const structure = await storage.scanBranchStructure({ start: "entry" }); - structure[0]!.timestamp = 0; - const stats = await storage.getStats(); - stats.usage.input = 999; - - const storedEntry = (await storage.getEntries(["entry"])).get("entry"); - expect(storedEntry?.type === "message" ? storedEntry.message.timestamp : undefined).toBe(NOW - 1); - expect(await storage.getRegister("fact.custom", "object")).toMatchObject({ value: { nested: [1] } }); - expect(await storage.listRegisters("fact.custom")).toHaveLength(1); - expect(await storage.scanUsage({})).toMatchObject([{ usage: { input: 4 } }]); - expect(await storage.scanEntries({})).toMatchObject([{ timestamp: NOW }]); - expect(await storage.getStats()).toMatchObject({ usage: { input: 4 } }); - }); - - it("preserves historical stores across register-only commits", async () => { - const storage = new MemoryStorage({ now: () => NOW }); - await storage.commit({ - writes: [ - { kind: "entry", entry: messageEntry("root") }, - { kind: "usage", row: { id: "historical-usage", usage: usage(2, 3), adjustment: false } }, - ], - }); - const entriesBefore = await storage.scanEntries({}); - const usageBefore = await storage.scanUsage({}); - const statsBefore = await storage.getStats(); - - await expect( - storage.commit({ - writes: [ - { kind: "register", op: "set", namespace: "fact.name", key: "", value: "first" }, - { kind: "register", op: "set", namespace: "fact.name", key: "", value: "second" }, - ], - }), - ).resolves.toEqual({ firstSeq: 3, seqs: [3, 4], timestamp: NOW }); - - expect(await storage.scanEntries({})).toEqual(entriesBefore); - expect(await storage.scanUsage({})).toEqual(usageBefore); - expect(await storage.getStats()).toEqual(statsBefore); - expect(await storage.getRegister("fact.name", "")).toEqual({ - namespace: "fact.name", - key: "", - value: "second", - seq: 4, - }); - }); - - it("isolates every store when a mixed transaction fails", async () => { - const storage = new MemoryStorage({ now: () => NOW }); - await storage.commit({ - writes: [ - { kind: "entry", entry: messageEntry("root") }, - { kind: "usage", row: { id: "taken", usage: usage(1, 1), adjustment: false } }, - ], - }); - const statsBefore = await storage.getStats(); - - await expect( - storage.commit({ - writes: [ - { kind: "register", op: "set", namespace: "fact.name", key: "", value: "must roll back" }, - { kind: "entry", entry: messageEntry("transient", "root") }, - { kind: "usage", row: { id: "transient-usage", usage: usage(5, 8), adjustment: true } }, - { kind: "entry", entry: messageEntry("taken", "root") }, - ], - }), - ).rejects.toThrow("Duplicate entry or usage id: taken"); - - expect(await storage.getRegister("fact.name", "")).toBeUndefined(); - expect(await storage.getEntries(["root", "transient"])).toEqual( - new Map([["root", { ...messageEntry("root"), seq: 1, timestamp: NOW }]]), - ); - expect((await storage.scanUsage({})).map((row) => row.id)).toEqual(["taken"]); - expect(await storage.getStats()).toEqual(statsBefore); - await expect( - storage.commit({ writes: [{ kind: "register", op: "set", namespace: "fact.name", key: "", value: "after" }] }), - ).resolves.toEqual({ firstSeq: 3, seqs: [3], timestamp: NOW }); +const conformance = createStorageConformance(() => { + const storage = new MemoryStorage({ now: () => NOW }); + return Promise.resolve({ + storage, + [Symbol.asyncDispose]: () => storage.close(), + }); +}); + +describe("MemoryStorage 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()); + } + }); + } +}); + +describe("MemoryStorage", () => { + it("uses the injected clock once per transaction", async () => { + let timestamp = NOW; + const storage = new MemoryStorage({ now: () => timestamp++ }); + + const first = await storage.commit({ + writes: [ + { + kind: "entry", + entry: { + id: "first", + parentId: null, + type: "custom", + customType: "note", + }, + }, + { kind: "register", op: "set", namespace: "fact.name", key: "", value: "first" }, + ], + }); + const second = await storage.commit({ + writes: [{ kind: "register", op: "set", namespace: "fact.name", key: "", value: "second" }], + }); + + expect(first.timestamp).toBe(NOW); + expect(second.timestamp).toBe(NOW + 1); + expect((await storage.getEntries(["first"])).get("first")?.timestamp).toBe(first.timestamp); + await storage.close(); }); }); diff --git a/tsconfig.json b/tsconfig.json index 112b976ab..ea7a514f3 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -13,6 +13,9 @@ "@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/session/testing": [ + "./packages/agent/src/harness/session/testing/index.ts" + ], "@earendil-works/pi-agent-core/node": ["./packages/agent/src/node.ts"], "@earendil-works/pi-agent-core/*": ["./packages/agent/src/*"], "@earendil-works/pi-session-backend-sqlite-node": ["./packages/session-backends/sqlite-node/src/index.ts"],