Squashed 'third_party/pi/' content from commit d981de122
git-subtree-dir: third_party/pi git-subtree-split: d981de1229ef899957bbe968bc8dcda02a21f477
This commit is contained in:
696
packages/session-backends/sqlite-node/test/storage.test.ts
Normal file
696
packages/session-backends/sqlite-node/test/storage.test.ts
Normal file
@@ -0,0 +1,696 @@
|
||||
import * as storedValues from "@earendil-works/pi-agent-core";
|
||||
import * as sessionWrites from "@earendil-works/pi-agent-core";
|
||||
import { BACKGROUND_CONTEXT, prepareStorageCommit } from "@earendil-works/pi-agent-core";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import {
|
||||
createNodeSqliteFactory,
|
||||
type SqliteDatabase,
|
||||
type SqliteStatement,
|
||||
SqliteStorage,
|
||||
sql,
|
||||
} from "../src/index.ts";
|
||||
import { applyInitialSchema } from "../src/sqlite/migrations.ts";
|
||||
import { advanceNextSeq, readNextSeq } from "../src/sqlite/session/session-sequences.ts";
|
||||
import { listValueReadQuery } from "../src/sqlite/session/values.ts";
|
||||
|
||||
const SESSION_ID = "session";
|
||||
|
||||
class TransactionCountingDatabase implements SqliteDatabase {
|
||||
readonly source: SqliteDatabase;
|
||||
transactionCount = 0;
|
||||
|
||||
constructor(source: SqliteDatabase) {
|
||||
this.source = source;
|
||||
}
|
||||
|
||||
exec(query: string): void {
|
||||
this.source.exec(query);
|
||||
}
|
||||
|
||||
prepare(query: string): SqliteStatement {
|
||||
return this.source.prepare(query);
|
||||
}
|
||||
|
||||
transaction<T>(callback: () => T): T {
|
||||
this.transactionCount++;
|
||||
return this.source.transaction(callback);
|
||||
}
|
||||
|
||||
close(): void {
|
||||
this.source.close();
|
||||
}
|
||||
}
|
||||
|
||||
const ZERO_USAGE = {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
|
||||
async function withStorage<T>(run: (storage: SqliteStorage, db: SqliteDatabase) => Promise<T>): Promise<T> {
|
||||
const db = await createNodeSqliteFactory().open(":memory:");
|
||||
try {
|
||||
await applyInitialSchema(db);
|
||||
const storage = new SqliteStorage(db, { sessionId: SESSION_ID, now: () => 1_700_000_000_000 });
|
||||
return await run(storage, db);
|
||||
} finally {
|
||||
db.close();
|
||||
}
|
||||
}
|
||||
|
||||
function explainQueryPlan(db: SqliteDatabase, query: string, ...params: unknown[]): string[] {
|
||||
return db
|
||||
.prepare(`EXPLAIN QUERY PLAN ${query}`)
|
||||
.all<{ detail: string }>(...params)
|
||||
.map((row) => row.detail);
|
||||
}
|
||||
|
||||
function expectBranchPlan(plan: string[]): void {
|
||||
expect(plan.some((detail) => detail.includes("SEARCH b USING COVERING INDEX ix_be_seq"))).toBe(true);
|
||||
expect(plan.some((detail) => detail.includes("SEARCH e USING PRIMARY KEY"))).toBe(true);
|
||||
expect(plan.some((detail) => detail.includes("USE TEMP B-TREE"))).toBe(false);
|
||||
expect(plan.some((detail) => detail.includes("SCAN e"))).toBe(false);
|
||||
}
|
||||
|
||||
function insertCommitSessionRow(db: SqliteDatabase, nextSeq = 1): void {
|
||||
sql`INSERT INTO sessions
|
||||
(id, created_at, parent_session_id, storage_version, metadata, message_count, usage_payload, next_seq)
|
||||
VALUES (
|
||||
${SESSION_ID},
|
||||
${1},
|
||||
${null},
|
||||
${1},
|
||||
${null},
|
||||
${0},
|
||||
${JSON.stringify(ZERO_USAGE)},
|
||||
${nextSeq}
|
||||
)`.run(db);
|
||||
}
|
||||
|
||||
describe("SqliteStorage", () => {
|
||||
it("uses one write transaction for an ordinary commit", async () => {
|
||||
const db = new TransactionCountingDatabase(await createNodeSqliteFactory().open(":memory:"));
|
||||
try {
|
||||
await applyInitialSchema(db);
|
||||
insertCommitSessionRow(db);
|
||||
const storage = new SqliteStorage(db, { sessionId: SESSION_ID, now: () => 1_700_000_000_000 });
|
||||
|
||||
await storage.commit([storedValues.setValue(storedValues.sessionName, "name")], BACKGROUND_CONTEXT);
|
||||
|
||||
expect(db.transactionCount).toBe(1);
|
||||
await storage.close(BACKGROUND_CONTEXT);
|
||||
} finally {
|
||||
db.close();
|
||||
}
|
||||
});
|
||||
|
||||
it("commits root entries and append-to-tip entries into the branch index", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
insertCommitSessionRow(db);
|
||||
|
||||
expect(
|
||||
await storage.commit(
|
||||
[
|
||||
sessionWrites.insertEntry({
|
||||
id: "root",
|
||||
parentId: null,
|
||||
type: "message",
|
||||
message: { role: "user", content: "root", timestamp: 10 },
|
||||
}),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
).toEqual({
|
||||
firstSeq: 1,
|
||||
seqs: [1],
|
||||
timestamp: 1_700_000_000_000,
|
||||
stats: { messageCount: 1, usage: ZERO_USAGE },
|
||||
});
|
||||
expect(
|
||||
await storage.commit(
|
||||
[
|
||||
sessionWrites.insertEntry({
|
||||
id: "child",
|
||||
parentId: "root",
|
||||
type: "message",
|
||||
message: { role: "user", content: "child", timestamp: 11 },
|
||||
}),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
).toEqual({
|
||||
firstSeq: 2,
|
||||
seqs: [2],
|
||||
timestamp: 1_700_000_000_000,
|
||||
stats: { messageCount: 2, usage: ZERO_USAGE },
|
||||
});
|
||||
|
||||
expect(sql`SELECT branch_id, tip_entry_id, tip_seq FROM branch_meta`.all(db)).toEqual([
|
||||
{ branch_id: "root", tip_entry_id: "child", tip_seq: 2 },
|
||||
]);
|
||||
expect(
|
||||
sql`SELECT branch_id, entry_id, entry_seq, entry_type FROM branch_entries ORDER BY entry_seq`.all(db),
|
||||
).toEqual([
|
||||
{ branch_id: "root", entry_id: "root", entry_seq: 1, entry_type: "message" },
|
||||
{ branch_id: "root", entry_id: "child", entry_seq: 2, entry_type: "message" },
|
||||
]);
|
||||
expect((await storage.scanBranch({ start: "child" }, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([
|
||||
"child",
|
||||
"root",
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
it("commits divergent branch entries by materializing a new segment", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
insertCommitSessionRow(db);
|
||||
await storage.commit(
|
||||
[
|
||||
sessionWrites.insertEntry({
|
||||
id: "root",
|
||||
parentId: null,
|
||||
type: "message",
|
||||
message: { role: "user", content: "root", timestamp: 10 },
|
||||
}),
|
||||
sessionWrites.insertEntry({
|
||||
id: "left",
|
||||
parentId: "root",
|
||||
type: "message",
|
||||
message: { role: "user", content: "left", timestamp: 11 },
|
||||
}),
|
||||
sessionWrites.insertEntry({
|
||||
id: "right",
|
||||
parentId: "root",
|
||||
type: "message",
|
||||
message: { role: "user", content: "right", timestamp: 12 },
|
||||
}),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
expect(
|
||||
sql`SELECT branch_id, tip_entry_id, tip_seq, base_branch_id, base_seq FROM branch_meta ORDER BY branch_id`.all(
|
||||
db,
|
||||
),
|
||||
).toEqual([
|
||||
{ branch_id: "right", tip_entry_id: "right", tip_seq: 3, base_branch_id: null, base_seq: null },
|
||||
{ branch_id: "root", tip_entry_id: "left", tip_seq: 2, base_branch_id: null, base_seq: null },
|
||||
]);
|
||||
expect(
|
||||
sql`SELECT branch_id, entry_id, entry_seq, entry_type FROM branch_entries ORDER BY branch_id, entry_seq`.all(
|
||||
db,
|
||||
),
|
||||
).toEqual([
|
||||
{ branch_id: "right", entry_id: "root", entry_seq: 1, entry_type: "message" },
|
||||
{ branch_id: "right", entry_id: "right", entry_seq: 3, entry_type: "message" },
|
||||
{ branch_id: "root", entry_id: "root", entry_seq: 1, entry_type: "message" },
|
||||
{ branch_id: "root", entry_id: "left", entry_seq: 2, entry_type: "message" },
|
||||
]);
|
||||
expect((await storage.scanBranch({ start: "right" }, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([
|
||||
"right",
|
||||
"root",
|
||||
]);
|
||||
expect((await storage.scanBranch({ start: "left" }, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([
|
||||
"left",
|
||||
"root",
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
it("bases divergent branch segments at the newest compaction", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
insertCommitSessionRow(db);
|
||||
await storage.commit(
|
||||
[
|
||||
sessionWrites.insertEntry({
|
||||
id: "root",
|
||||
parentId: null,
|
||||
type: "message",
|
||||
message: { role: "user", content: "root", timestamp: 10 },
|
||||
}),
|
||||
sessionWrites.insertEntry({
|
||||
id: "compact",
|
||||
parentId: "root",
|
||||
type: "compaction",
|
||||
summary: "summary",
|
||||
retainedTail: [],
|
||||
tokensBefore: 1,
|
||||
fromHook: false,
|
||||
}),
|
||||
sessionWrites.insertEntry({
|
||||
id: "left",
|
||||
parentId: "compact",
|
||||
type: "message",
|
||||
message: { role: "user", content: "left", timestamp: 11 },
|
||||
}),
|
||||
sessionWrites.insertEntry({
|
||||
id: "leaf",
|
||||
parentId: "left",
|
||||
type: "message",
|
||||
message: { role: "user", content: "leaf", timestamp: 12 },
|
||||
}),
|
||||
sessionWrites.insertEntry({
|
||||
id: "right",
|
||||
parentId: "left",
|
||||
type: "message",
|
||||
message: { role: "user", content: "right", timestamp: 13 },
|
||||
}),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
expect(
|
||||
sql`SELECT branch_id, tip_entry_id, tip_seq, base_branch_id, base_seq FROM branch_meta WHERE branch_id = ${"right"}`.get(
|
||||
db,
|
||||
),
|
||||
).toEqual({
|
||||
branch_id: "right",
|
||||
tip_entry_id: "right",
|
||||
tip_seq: 5,
|
||||
base_branch_id: "root",
|
||||
base_seq: 2,
|
||||
});
|
||||
expect(
|
||||
sql`SELECT branch_id, entry_id, entry_seq, entry_type FROM branch_entries WHERE branch_id = ${"right"} ORDER BY entry_seq`.all(
|
||||
db,
|
||||
),
|
||||
).toEqual([
|
||||
{ branch_id: "right", entry_id: "left", entry_seq: 3, entry_type: "message" },
|
||||
{ branch_id: "right", entry_id: "right", entry_seq: 5, entry_type: "message" },
|
||||
]);
|
||||
expect((await storage.scanBranch({ start: "right" }, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([
|
||||
"right",
|
||||
"left",
|
||||
"compact",
|
||||
"root",
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
it("gets entries by requested id order", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
sql`INSERT INTO entries (session_id, id, parent_id, seq, type, custom_type, timestamp, payload)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"first"}, ${null}, ${1}, ${"custom"}, ${"note"}, ${10}, ${JSON.stringify({ data: { value: 1 } })}),
|
||||
(${SESSION_ID}, ${"second"}, ${"first"}, ${2}, ${"message"}, ${null}, ${11}, ${JSON.stringify({ message: { role: "user", content: "hi", timestamp: 11 } })})`.run(
|
||||
db,
|
||||
);
|
||||
|
||||
const entries = await storage.getEntries(["second", "missing", "first"], BACKGROUND_CONTEXT);
|
||||
|
||||
expect([...entries.keys()]).toEqual(["second", "first"]);
|
||||
expect(entries.get("second")).toMatchObject({ id: "second", type: "message", parentId: "first" });
|
||||
expect(entries.get("first")).toMatchObject({ id: "first", type: "custom", customType: "note" });
|
||||
});
|
||||
});
|
||||
|
||||
it("scans decoded entries with filters and sequence bounds", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
sql`INSERT INTO entries (session_id, id, parent_id, seq, type, custom_type, timestamp, payload)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"one"}, ${null}, ${1}, ${"custom"}, ${"note"}, ${10}, ${JSON.stringify({ data: 1 })}),
|
||||
(${SESSION_ID}, ${"two"}, ${"one"}, ${2}, ${"message"}, ${null}, ${11}, ${JSON.stringify({ message: { role: "user", content: "two", timestamp: 11 } })}),
|
||||
(${SESSION_ID}, ${"three"}, ${"two"}, ${3}, ${"custom"}, ${"note"}, ${12}, ${JSON.stringify({ data: 3 })})`.run(
|
||||
db,
|
||||
);
|
||||
|
||||
expect(
|
||||
(await storage.scanEntries({ order: "desc", type: "custom", fromSeq: 2 }, BACKGROUND_CONTEXT)).map(
|
||||
(entry) => entry.id,
|
||||
),
|
||||
).toEqual(["three"]);
|
||||
expect(
|
||||
(await storage.scanEntries({ order: "asc", customType: "note", limit: 2 }, BACKGROUND_CONTEXT)).map(
|
||||
(entry) => entry.id,
|
||||
),
|
||||
).toEqual(["one", "three"]);
|
||||
});
|
||||
});
|
||||
|
||||
it("scans branch entries through materialized branch segments", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
sql`INSERT INTO entries (session_id, id, parent_id, seq, type, custom_type, timestamp, payload)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"root"}, ${null}, ${1}, ${"message"}, ${null}, ${10}, ${JSON.stringify({ message: { role: "user", content: "root", timestamp: 10 } })}),
|
||||
(${SESSION_ID}, ${"compact"}, ${"root"}, ${2}, ${"compaction"}, ${null}, ${11}, ${JSON.stringify({ summary: "s", retainedTail: [], tokensBefore: 1, fromHook: false })}),
|
||||
(${SESSION_ID}, ${"old"}, ${"compact"}, ${3}, ${"message"}, ${null}, ${12}, ${JSON.stringify({ message: { role: "assistant", content: "old", timestamp: 12 } })}),
|
||||
(${SESSION_ID}, ${"custom"}, ${"old"}, ${4}, ${"custom"}, ${"note"}, ${13}, ${JSON.stringify({ data: 4 })}),
|
||||
(${SESSION_ID}, ${"leaf"}, ${"custom"}, ${5}, ${"message"}, ${null}, ${14}, ${JSON.stringify({ message: { role: "user", content: "leaf", timestamp: 14 } })})`.run(
|
||||
db,
|
||||
);
|
||||
sql`INSERT INTO branch_meta (session_id, branch_id, tip_entry_id, tip_seq, base_branch_id, base_seq)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"base"}, ${"compact"}, ${2}, ${null}, ${null}),
|
||||
(${SESSION_ID}, ${"new"}, ${"leaf"}, ${5}, ${"base"}, ${2})`.run(db);
|
||||
sql`INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq, entry_type)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"base"}, ${"root"}, ${1}, ${"message"}),
|
||||
(${SESSION_ID}, ${"base"}, ${"compact"}, ${2}, ${"compaction"}),
|
||||
(${SESSION_ID}, ${"new"}, ${"old"}, ${3}, ${"message"}),
|
||||
(${SESSION_ID}, ${"new"}, ${"custom"}, ${4}, ${"custom"}),
|
||||
(${SESSION_ID}, ${"new"}, ${"leaf"}, ${5}, ${"message"})`.run(db);
|
||||
|
||||
expect(
|
||||
(await storage.scanBranch({ start: "leaf", stopAtType: "compaction", limit: 2 }, BACKGROUND_CONTEXT)).map(
|
||||
(entry) => entry.id,
|
||||
),
|
||||
).toEqual(["leaf", "custom"]);
|
||||
expect(
|
||||
(
|
||||
await storage.scanBranch(
|
||||
{ start: "leaf", order: "oldestFirst", type: "message", cursor: { seq: 1 } },
|
||||
BACKGROUND_CONTEXT,
|
||||
)
|
||||
).map((entry) => entry.id),
|
||||
).toEqual(["old", "leaf"]);
|
||||
});
|
||||
});
|
||||
|
||||
it("resolves a branch segment that physically contains the start entry", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
sql`INSERT INTO entries (session_id, id, parent_id, seq, type, custom_type, timestamp, payload)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"root"}, ${null}, ${1}, ${"message"}, ${null}, ${10}, ${JSON.stringify({ message: { role: "user", content: "root", timestamp: 10 } })}),
|
||||
(${SESSION_ID}, ${"base-tip"}, ${"root"}, ${2}, ${"message"}, ${null}, ${11}, ${JSON.stringify({ message: { role: "assistant", content: "base", timestamp: 11 } })}),
|
||||
(${SESSION_ID}, ${"new-tip"}, ${"base-tip"}, ${3}, ${"message"}, ${null}, ${12}, ${JSON.stringify({ message: { role: "user", content: "new", timestamp: 12 } })})`.run(
|
||||
db,
|
||||
);
|
||||
sql`INSERT INTO branch_meta (session_id, branch_id, tip_entry_id, tip_seq, base_branch_id, base_seq)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"aaa-new"}, ${"new-tip"}, ${3}, ${"zzz-base"}, ${2}),
|
||||
(${SESSION_ID}, ${"zzz-base"}, ${"base-tip"}, ${2}, ${null}, ${null})`.run(db);
|
||||
sql`INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq, entry_type)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"aaa-new"}, ${"new-tip"}, ${3}, ${"message"}),
|
||||
(${SESSION_ID}, ${"zzz-base"}, ${"root"}, ${1}, ${"message"}),
|
||||
(${SESSION_ID}, ${"zzz-base"}, ${"base-tip"}, ${2}, ${"message"})`.run(db);
|
||||
|
||||
expect(
|
||||
(await storage.scanBranch({ start: "base-tip", order: "oldestFirst" }, BACKGROUND_CONTEXT)).map(
|
||||
(entry) => entry.id,
|
||||
),
|
||||
).toEqual(["root", "base-tip"]);
|
||||
});
|
||||
});
|
||||
|
||||
it("applies branch stop boundaries across base segments before filtering", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
sql`INSERT INTO entries (session_id, id, parent_id, seq, type, custom_type, timestamp, payload)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"root"}, ${null}, ${1}, ${"message"}, ${null}, ${10}, ${JSON.stringify({ message: { role: "user", content: "root", timestamp: 10 } })}),
|
||||
(${SESSION_ID}, ${"compact"}, ${"root"}, ${2}, ${"compaction"}, ${null}, ${11}, ${JSON.stringify({ summary: "s", retainedTail: [], tokensBefore: 1, fromHook: false })}),
|
||||
(${SESSION_ID}, ${"after"}, ${"compact"}, ${3}, ${"message"}, ${null}, ${12}, ${JSON.stringify({ message: { role: "assistant", content: "after", timestamp: 12 } })}),
|
||||
(${SESSION_ID}, ${"leaf"}, ${"after"}, ${4}, ${"custom"}, ${"note"}, ${13}, ${JSON.stringify({ data: 4 })})`.run(
|
||||
db,
|
||||
);
|
||||
sql`INSERT INTO branch_meta (session_id, branch_id, tip_entry_id, tip_seq, base_branch_id, base_seq)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"base"}, ${"compact"}, ${2}, ${null}, ${null}),
|
||||
(${SESSION_ID}, ${"new"}, ${"leaf"}, ${4}, ${"base"}, ${2})`.run(db);
|
||||
sql`INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq, entry_type)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"base"}, ${"root"}, ${1}, ${"message"}),
|
||||
(${SESSION_ID}, ${"base"}, ${"compact"}, ${2}, ${"compaction"}),
|
||||
(${SESSION_ID}, ${"new"}, ${"after"}, ${3}, ${"message"}),
|
||||
(${SESSION_ID}, ${"new"}, ${"leaf"}, ${4}, ${"custom"})`.run(db);
|
||||
|
||||
expect(
|
||||
(await storage.scanBranch({ start: "leaf", stopAtType: "compaction" }, BACKGROUND_CONTEXT)).map(
|
||||
(entry) => entry.id,
|
||||
),
|
||||
).toEqual(["leaf", "after", "compact"]);
|
||||
expect(
|
||||
(
|
||||
await storage.scanBranch(
|
||||
{ start: "leaf", stopAtType: "compaction", type: "message" },
|
||||
BACKGROUND_CONTEXT,
|
||||
)
|
||||
).map((entry) => entry.id),
|
||||
).toEqual(["after"]);
|
||||
});
|
||||
});
|
||||
|
||||
it("uses branch_entries as the outer scan for branch payload queries", async () => {
|
||||
await withStorage(async (_storage, db) => {
|
||||
const plan = explainQueryPlan(
|
||||
db,
|
||||
`SELECT e.id, e.parent_id, e.seq, e.type, e.custom_type, e.timestamp, e.payload
|
||||
FROM branch_entries b
|
||||
CROSS JOIN entries 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 b.entry_seq <= ?
|
||||
ORDER BY b.entry_seq DESC LIMIT ?`,
|
||||
SESSION_ID,
|
||||
"main",
|
||||
0,
|
||||
10,
|
||||
2,
|
||||
);
|
||||
|
||||
expectBranchPlan(plan);
|
||||
});
|
||||
});
|
||||
|
||||
it("scans branch structure without payloads", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
sql`INSERT INTO entries (session_id, id, parent_id, seq, type, custom_type, timestamp, payload)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"root"}, ${null}, ${1}, ${"message"}, ${null}, ${10}, ${JSON.stringify({ message: { role: "user", content: "root", timestamp: 10 } })}),
|
||||
(${SESSION_ID}, ${"custom"}, ${"root"}, ${2}, ${"custom"}, ${"note"}, ${11}, ${JSON.stringify({ data: 2 })})`.run(
|
||||
db,
|
||||
);
|
||||
sql`INSERT INTO branch_meta (session_id, branch_id, tip_entry_id, tip_seq, base_branch_id, base_seq)
|
||||
VALUES (${SESSION_ID}, ${"main"}, ${"custom"}, ${2}, ${null}, ${null})`.run(db);
|
||||
sql`INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq, entry_type)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"main"}, ${"root"}, ${1}, ${"message"}),
|
||||
(${SESSION_ID}, ${"main"}, ${"custom"}, ${2}, ${"custom"})`.run(db);
|
||||
|
||||
expect(await storage.scanBranchStructure({ start: "custom", customType: "note" }, BACKGROUND_CONTEXT)).toEqual(
|
||||
[
|
||||
{
|
||||
id: "custom",
|
||||
parentId: "root",
|
||||
seq: 2,
|
||||
timestamp: 11,
|
||||
type: "custom",
|
||||
customType: "note",
|
||||
},
|
||||
],
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
it("uses branch_entries as the outer scan for branch structure queries", async () => {
|
||||
await withStorage(async (_storage, db) => {
|
||||
const plan = explainQueryPlan(
|
||||
db,
|
||||
`SELECT e.id, e.parent_id, e.seq, e.type, e.custom_type, e.timestamp
|
||||
FROM branch_entries b
|
||||
CROSS JOIN entries 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 b.entry_seq <= ?
|
||||
ORDER BY b.entry_seq ASC LIMIT ?`,
|
||||
SESSION_ID,
|
||||
"main",
|
||||
0,
|
||||
10,
|
||||
2,
|
||||
);
|
||||
|
||||
expectBranchPlan(plan);
|
||||
});
|
||||
});
|
||||
|
||||
it("scans decoded usage rows with sequence bounds", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
const usage = {
|
||||
input: 1,
|
||||
output: 2,
|
||||
cacheRead: 3,
|
||||
cacheWrite: 4,
|
||||
totalTokens: 10,
|
||||
cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
|
||||
};
|
||||
sql`INSERT INTO usage_ledger (session_id, id, seq, entry_id, adjustment, usage, details)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"u1"}, ${1}, ${"e1"}, ${0}, ${JSON.stringify(usage)}, ${null}),
|
||||
(${SESSION_ID}, ${"u2"}, ${2}, ${null}, ${1}, ${JSON.stringify(usage)}, ${JSON.stringify({ reason: "adjust" })}),
|
||||
(${SESSION_ID}, ${"u3"}, ${3}, ${"e3"}, ${0}, ${JSON.stringify(usage)}, ${null})`.run(db);
|
||||
|
||||
expect(await storage.scanUsage({ fromSeq: 2, order: "asc", limit: 1 }, BACKGROUND_CONTEXT)).toEqual([
|
||||
{ id: "u2", seq: 2, usage, adjustment: true, details: { reason: "adjust" } },
|
||||
]);
|
||||
expect(
|
||||
(await storage.scanUsage({ toSeq: 2, order: "desc" }, BACKGROUND_CONTEXT)).map((row) => row.id),
|
||||
).toEqual(["u2", "u1"]);
|
||||
});
|
||||
});
|
||||
|
||||
it("prepares committed writes with assigned sequences and timestamp", () => {
|
||||
const prepared = prepareStorageCommit(
|
||||
[
|
||||
sessionWrites.insertEntry({
|
||||
id: "entry",
|
||||
parentId: null,
|
||||
type: "message",
|
||||
message: { role: "user", content: "hi", timestamp: 1 },
|
||||
}),
|
||||
sessionWrites.insertUsage({
|
||||
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,
|
||||
}),
|
||||
storedValues.setValue(storedValues.sessionName, "name"),
|
||||
storedValues.deleteValue(storedValues.entryLabel("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 },
|
||||
storedValues.setValue(storedValues.sessionName, "name"),
|
||||
storedValues.deleteValue(storedValues.entryLabel("entry")),
|
||||
]);
|
||||
});
|
||||
|
||||
it("reads and advances the next commit sequence", async () => {
|
||||
await withStorage(async (_storage, db) => {
|
||||
sql`INSERT INTO sessions
|
||||
(id, created_at, parent_session_id, storage_version, metadata, message_count, usage_payload, next_seq)
|
||||
VALUES (${SESSION_ID}, ${1}, ${null}, ${1}, ${null}, ${0}, ${JSON.stringify({})}, ${7})`.run(db);
|
||||
|
||||
expect(readNextSeq(db, SESSION_ID)).toBe(7);
|
||||
advanceNextSeq(db, SESSION_ID, 10);
|
||||
expect(readNextSeq(db, SESSION_ID)).toBe(10);
|
||||
});
|
||||
});
|
||||
|
||||
it("gets maintained session stats", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
const usage = {
|
||||
input: 1,
|
||||
output: 2,
|
||||
cacheRead: 3,
|
||||
cacheWrite: 4,
|
||||
totalTokens: 10,
|
||||
cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
|
||||
};
|
||||
sql`INSERT INTO sessions
|
||||
(id, created_at, parent_session_id, storage_version, metadata, message_count, usage_payload, next_seq)
|
||||
VALUES (${SESSION_ID}, ${1}, ${null}, ${1}, ${null}, ${2}, ${JSON.stringify(usage)}, ${3})`.run(db);
|
||||
|
||||
expect(await storage.getStats(BACKGROUND_CONTEXT)).toEqual({ messageCount: 2, usage });
|
||||
const next = await storage.commit(
|
||||
[storedValues.setValue(storedValues.sessionName, "after-history")],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
expect(next.stats).toEqual({ messageCount: 2, usage });
|
||||
});
|
||||
});
|
||||
|
||||
it("includes historical totals in the first commit after storage reopen", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
insertCommitSessionRow(db);
|
||||
const usage = {
|
||||
input: 1,
|
||||
output: 2,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 3,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
await storage.commit(
|
||||
[
|
||||
sessionWrites.insertEntry({
|
||||
id: "history",
|
||||
parentId: null,
|
||||
type: "message",
|
||||
message: { role: "user", content: "history", timestamp: 1 },
|
||||
}),
|
||||
sessionWrites.insertUsage({ id: "usage", usage, adjustment: false }),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
await storage.close(BACKGROUND_CONTEXT);
|
||||
|
||||
const reopened = new SqliteStorage(db, { sessionId: SESSION_ID, now: () => 1_700_000_000_000 });
|
||||
const result = await reopened.commit(
|
||||
[storedValues.setValue(storedValues.sessionName, "reopened")],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
expect(result.stats).toEqual({ messageCount: 1, usage });
|
||||
expect(result.stats).toEqual(await reopened.getStats(BACKGROUND_CONTEXT));
|
||||
await reopened.close(BACKGROUND_CONTEXT);
|
||||
});
|
||||
});
|
||||
|
||||
it("gets a decoded scalar value by bound address", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
const address = storedValues.value<{ ready: boolean }>("test.value", "state");
|
||||
sql`INSERT INTO scalar_values (session_id, namespace, key, seq, value)
|
||||
VALUES (${SESSION_ID}, ${address.namespace}, ${address.key}, ${1}, ${JSON.stringify({ ready: true })})`.run(
|
||||
db,
|
||||
);
|
||||
|
||||
expect(await storage.getValue(address, BACKGROUND_CONTEXT)).toEqual({
|
||||
address,
|
||||
seq: 1,
|
||||
value: { ready: true },
|
||||
});
|
||||
expect(
|
||||
await storage.getValue(storedValues.value<unknown>("test.value", "missing"), BACKGROUND_CONTEXT),
|
||||
).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
it("scans decoded scalar values by namespace and key prefix", async () => {
|
||||
await withStorage(async (storage, db) => {
|
||||
sql`INSERT INTO scalar_values (session_id, namespace, key, seq, value)
|
||||
VALUES
|
||||
(${SESSION_ID}, ${"test.value"}, ${"app:one"}, ${1}, ${JSON.stringify(1)}),
|
||||
(${SESSION_ID}, ${"test.value"}, ${"app:two"}, ${2}, ${JSON.stringify(2)}),
|
||||
(${SESSION_ID}, ${"test.value"}, ${"app:\ufffftail"}, ${3}, ${JSON.stringify(3)}),
|
||||
(${SESSION_ID}, ${"test.value"}, ${"app;other"}, ${4}, ${JSON.stringify(4)}),
|
||||
(${SESSION_ID}, ${"test.value"}, ${"other"}, ${5}, ${JSON.stringify(5)}),
|
||||
(${SESSION_ID}, ${"other.value"}, ${""}, ${6}, ${JSON.stringify("name")})`.run(db);
|
||||
|
||||
expect(
|
||||
await storage.scanValues(storedValues.value<unknown>("test.value", "app:"), BACKGROUND_CONTEXT),
|
||||
).toEqual([
|
||||
{ address: storedValues.value<unknown>("test.value", "app:one"), seq: 1, value: 1 },
|
||||
{ address: storedValues.value<unknown>("test.value", "app:two"), seq: 2, value: 2 },
|
||||
{ address: storedValues.value<unknown>("test.value", "app:\ufffftail"), seq: 3, value: 3 },
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
it("uses the list primary key for ascending and descending cursor pages", async () => {
|
||||
await withStorage(async (_storage, db) => {
|
||||
const address = storedValues.list<unknown>("test.list", "events");
|
||||
for (const options of [
|
||||
{ order: "asc" as const },
|
||||
{ order: "asc" as const, cursor: { seq: 4 } },
|
||||
{ order: "desc" as const },
|
||||
{ order: "desc" as const, cursor: { seq: 4 } },
|
||||
]) {
|
||||
const query = listValueReadQuery(SESSION_ID, address, options);
|
||||
const plan = explainQueryPlan(db, query.queryText, ...query.params);
|
||||
expect(plan.some((detail) => detail.includes("USING PRIMARY KEY"))).toBe(true);
|
||||
expect(plan.some((detail) => detail.includes("SCAN list_values"))).toBe(false);
|
||||
expect(plan.some((detail) => detail.includes("USE TEMP B-TREE"))).toBe(false);
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user