feat: sqlite session storage (#6594)
This PR: - Adds retainedTail to compaction entries in the new agent harness so we don't have to walk up the tree for the 2000 tokens before compaction, - Changes getPathToRoot to getPathToRootOrCompaction to only load until last compaction, as unnecessary to access all nodes where it is called, - Adds a SQLite storage backend, in a separate packages/session-backend-sqlite, with a migration system and schemas as per on-site discussions: sessions to match session header messages (except for metadata, which I couldn't understand what it's used for or where it gets written, so I omitted it), session_entries for shared entry types as columns plus payload as a json for what remains, session_sequences to represent the append-only, serialized nature of the jsonl files, branch_entries to attribute nodes to branches (relationship one-to-many), and session_materialized with the session info (see /session in TUI) to act as a "cache" or quick-access for costs, message count, token info, labels, session name, and model-thinking-level config (e.g. for fast resume). - This is compatible with the new agent harness Session abstraction.
This commit is contained in:
Cristina Poncela Cubeiro
committed by
GitHub
GitHub
parent
54fad505b9
commit
9e7582aa03
@@ -91,6 +91,7 @@ function createCompactionEntry(
|
||||
summary: string,
|
||||
firstKeptEntryId: string,
|
||||
parentId: string | null = null,
|
||||
retainedTail?: AgentMessage[],
|
||||
): CompactionEntry {
|
||||
return {
|
||||
type: "compaction",
|
||||
@@ -100,6 +101,7 @@ function createCompactionEntry(
|
||||
summary,
|
||||
firstKeptEntryId,
|
||||
tokensBefore: 1234,
|
||||
retainedTail,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -343,12 +345,38 @@ describe("harness compaction", () => {
|
||||
const a1 = createMessageEntry(createAssistantMessage("a"), u1.id);
|
||||
const u2 = createMessageEntry(createUserMessage("2"), a1.id);
|
||||
const a2 = createMessageEntry(createAssistantMessage("b"), u2.id);
|
||||
const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id, a2.id);
|
||||
const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id, a2.id, [
|
||||
createUserMessage("2"),
|
||||
createAssistantMessage("b"),
|
||||
]);
|
||||
const u3 = createMessageEntry(createUserMessage("3"), compaction.id);
|
||||
const a3 = createMessageEntry(createAssistantMessage("c"), u3.id);
|
||||
const loaded = buildSessionContext([u1, a1, u2, a2, compaction, u3, a3]);
|
||||
expect(loaded.messages).toHaveLength(5);
|
||||
expect(loaded.messages[0]?.role).toBe("compactionSummary");
|
||||
expect(loaded.messages.map((message) => message.role)).toEqual([
|
||||
"compactionSummary",
|
||||
"user",
|
||||
"assistant",
|
||||
"user",
|
||||
"assistant",
|
||||
]);
|
||||
});
|
||||
|
||||
it("falls back to firstKeptEntryId when a compaction has no retained tail", () => {
|
||||
const u1 = createMessageEntry(createUserMessage("1"));
|
||||
const a1 = createMessageEntry(createAssistantMessage("a"), u1.id);
|
||||
const u2 = createMessageEntry(createUserMessage("2"), a1.id);
|
||||
const a2 = createMessageEntry(createAssistantMessage("b"), u2.id);
|
||||
const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id, a2.id);
|
||||
const u3 = createMessageEntry(createUserMessage("3"), compaction.id);
|
||||
const loaded = buildSessionContext([u1, a1, u2, a2, compaction, u3]);
|
||||
expect(loaded.messages.map((message) => message.role)).toEqual([
|
||||
"compactionSummary",
|
||||
"user",
|
||||
"assistant",
|
||||
"user",
|
||||
]);
|
||||
});
|
||||
|
||||
it("tracks model and thinking level changes in built context", () => {
|
||||
@@ -374,6 +402,7 @@ describe("harness compaction", () => {
|
||||
expect(preparation).toBeDefined();
|
||||
expect(preparation?.previousSummary).toBe("First summary");
|
||||
expect(preparation?.firstKeptEntryId).toBeTruthy();
|
||||
expect(preparation?.retainedTail.length).toBeGreaterThan(0);
|
||||
expect(preparation?.tokensBefore).toBe(estimateContextTokens(buildSessionContext(pathEntries).messages).tokens);
|
||||
});
|
||||
|
||||
@@ -566,6 +595,7 @@ describe("harness compaction", () => {
|
||||
firstKeptEntryId: "entry-keep",
|
||||
messagesToSummarize: messages,
|
||||
turnPrefixMessages: messages,
|
||||
retainedTail: messages,
|
||||
isSplitTurn: true,
|
||||
tokensBefore: 600000,
|
||||
fileOps: { read: new Set(), written: new Set(), edited: new Set() },
|
||||
@@ -583,6 +613,7 @@ describe("harness compaction", () => {
|
||||
firstKeptEntryId: "entry-keep",
|
||||
messagesToSummarize: messages,
|
||||
turnPrefixMessages: [],
|
||||
retainedTail: messages,
|
||||
isSplitTurn: false,
|
||||
tokensBefore: 100,
|
||||
fileOps: { read: new Set(), written: new Set(), edited: new Set() },
|
||||
@@ -619,6 +650,7 @@ describe("harness compaction", () => {
|
||||
turnPrefixMessages: messages,
|
||||
isSplitTurn: true,
|
||||
tokensBefore: 100,
|
||||
retainedTail: messages,
|
||||
fileOps: { read: new Set(), written: new Set(), edited: new Set() },
|
||||
settings: { enabled: true, reserveTokens: 2000, keepRecentTokens: 20 },
|
||||
};
|
||||
@@ -642,6 +674,7 @@ describe("harness compaction", () => {
|
||||
firstKeptEntryId: "entry-keep",
|
||||
messagesToSummarize: [],
|
||||
turnPrefixMessages: messages,
|
||||
retainedTail: messages,
|
||||
isSplitTurn: true,
|
||||
tokensBefore: 100,
|
||||
fileOps: { read: new Set(), written: new Set(), edited: new Set() },
|
||||
@@ -659,6 +692,7 @@ describe("harness compaction", () => {
|
||||
firstKeptEntryId: "entry-keep",
|
||||
messagesToSummarize: [],
|
||||
turnPrefixMessages: messages,
|
||||
retainedTail: messages,
|
||||
isSplitTurn: true,
|
||||
tokensBefore: 100,
|
||||
fileOps: { read: new Set(), written: new Set(), edited: new Set() },
|
||||
@@ -697,6 +731,7 @@ describe("harness compaction", () => {
|
||||
expect(result.summary.length).toBeGreaterThan(0);
|
||||
expect(result.firstKeptEntryId).toBeTruthy();
|
||||
expect(result.usage?.totalTokens).toBeGreaterThan(0);
|
||||
expect(result.retainedTail?.length).toBeGreaterThan(0);
|
||||
expect(result.details).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -68,11 +68,20 @@ async function runSessionSuite(
|
||||
await session.appendMessage(createAssistantMessage("two"));
|
||||
const user2 = await session.appendMessage(createUserMessage("three"));
|
||||
await session.appendMessage(createAssistantMessage("four"));
|
||||
await session.appendCompaction("summary", user2, 1234);
|
||||
await session.appendCompaction("summary", user2, 1234, undefined, undefined, undefined, [
|
||||
createUserMessage("three"),
|
||||
createAssistantMessage("four"),
|
||||
]);
|
||||
await session.appendMessage(createUserMessage("five"));
|
||||
const context = await session.buildContext();
|
||||
expect(context.messages[0]?.role).toBe("compactionSummary");
|
||||
expect(context.messages).toHaveLength(4);
|
||||
expect(context.messages.map((message) => message.role)).toEqual([
|
||||
"compactionSummary",
|
||||
"user",
|
||||
"assistant",
|
||||
"user",
|
||||
]);
|
||||
});
|
||||
|
||||
it("supports moving with branch summary entries in context", async () => {
|
||||
|
||||
@@ -0,0 +1,441 @@
|
||||
import { mkdtempSync, writeFileSync } from "node:fs";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import {
|
||||
applyMigrations,
|
||||
createNodeSqliteFactory,
|
||||
type SqliteDatabase,
|
||||
type SqliteDatabaseFactory,
|
||||
type SqliteRunResult,
|
||||
type SqliteSessionMetadata,
|
||||
SqliteSessionRepo,
|
||||
SqliteSessionStorage,
|
||||
type SqliteStatement,
|
||||
} from "../../../storage/sqlite-node/src/index.ts";
|
||||
import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts";
|
||||
import { createAssistantMessage, createUserMessage } from "./session-test-utils.ts";
|
||||
|
||||
function createTempDir(): string {
|
||||
return mkdtempSync(join(tmpdir(), "pi-agent-sqlite-"));
|
||||
}
|
||||
|
||||
class ThrowingStatement implements SqliteStatement {
|
||||
private readonly onRun: () => Promise<SqliteRunResult>;
|
||||
|
||||
constructor(onRun: () => Promise<SqliteRunResult>) {
|
||||
this.onRun = onRun;
|
||||
}
|
||||
|
||||
async run(..._params: unknown[]): Promise<SqliteRunResult> {
|
||||
return this.onRun();
|
||||
}
|
||||
|
||||
async get<TRow extends object>(..._params: unknown[]): Promise<TRow | undefined> {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
async all<TRow extends object>(..._params: unknown[]): Promise<TRow[]> {
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
class CountingDatabase implements SqliteDatabase {
|
||||
closeCount = 0;
|
||||
private readonly statementFactory: (sql: string) => SqliteStatement;
|
||||
|
||||
constructor(statementFactory: (sql: string) => SqliteStatement) {
|
||||
this.statementFactory = statementFactory;
|
||||
}
|
||||
|
||||
async exec(_sql: string): Promise<void> {}
|
||||
|
||||
prepare(sql: string): SqliteStatement {
|
||||
return this.statementFactory(sql);
|
||||
}
|
||||
|
||||
async transaction<T>(fn: () => Promise<T>): Promise<T> {
|
||||
return fn();
|
||||
}
|
||||
|
||||
async close(): Promise<void> {
|
||||
this.closeCount += 1;
|
||||
}
|
||||
}
|
||||
|
||||
describe("SQLite migrations", () => {
|
||||
it("applies file-based migrations and records them", async () => {
|
||||
const root = createTempDir();
|
||||
const databasePath = join(root, "sessions.sqlite");
|
||||
const env = new NodeExecutionEnv({ cwd: root });
|
||||
const sqlite = createNodeSqliteFactory();
|
||||
const repo = new SqliteSessionRepo({ env, sqlite, databasePath });
|
||||
await repo.create({ cwd: root, id: "session-1" });
|
||||
|
||||
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"]);
|
||||
const tables = await db
|
||||
.prepare("SELECT name, sql FROM sqlite_master WHERE type = 'table' ORDER BY name")
|
||||
.all<{ name: string; sql: string | null }>();
|
||||
expect(tables.map((row) => row.name)).toEqual(
|
||||
expect.arrayContaining([
|
||||
"migrations",
|
||||
"sessions",
|
||||
"session_entries",
|
||||
"session_sequences",
|
||||
"branch_entries",
|
||||
"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");
|
||||
for (const tableName of [
|
||||
"sessions",
|
||||
"session_sequences",
|
||||
"branch_entries",
|
||||
"session_materialized",
|
||||
"entry_materialized",
|
||||
]) {
|
||||
const table = tables.find((row) => row.name === tableName);
|
||||
expect(table?.sql).toContain("WITHOUT ROWID");
|
||||
}
|
||||
} finally {
|
||||
await db.close();
|
||||
}
|
||||
});
|
||||
|
||||
it("persists session metadata through create, list, open, and fork", async () => {
|
||||
const root = createTempDir();
|
||||
const databasePath = join(root, "sessions.sqlite");
|
||||
const env = new NodeExecutionEnv({ cwd: root });
|
||||
const repo = new SqliteSessionRepo({ env, sqlite: createNodeSqliteFactory(), databasePath });
|
||||
const source = await repo.create({
|
||||
cwd: root,
|
||||
id: "session-1",
|
||||
metadata: { profile: "reviewer" },
|
||||
});
|
||||
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 fork = await repo.fork(sourceMetadata, { cwd: root, id: "session-2" });
|
||||
expect((await fork.getMetadata()).metadata).toEqual({ profile: "reviewer" });
|
||||
const overridden = await repo.fork(sourceMetadata, {
|
||||
cwd: root,
|
||||
id: "session-3",
|
||||
metadata: { profile: "writer" },
|
||||
});
|
||||
expect((await overridden.getMetadata()).metadata).toEqual({ profile: "writer" });
|
||||
});
|
||||
|
||||
it("materializes active leaf id in sessions transactionally", async () => {
|
||||
const root = createTempDir();
|
||||
const databasePath = join(root, "sessions.sqlite");
|
||||
const env = new NodeExecutionEnv({ cwd: root });
|
||||
const sqlite = createNodeSqliteFactory();
|
||||
const repo = new SqliteSessionRepo({ 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"));
|
||||
await session.getStorage().setLeafId(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);
|
||||
} finally {
|
||||
await db.close();
|
||||
}
|
||||
|
||||
const reopened = await repo.open(await session.getMetadata());
|
||||
expect(await reopened.getLeafId()).toBe(rootId);
|
||||
expect(childId).not.toBe(rootId);
|
||||
});
|
||||
|
||||
it("materializes a new branch when appending from a parent with an existing child", async () => {
|
||||
const root = createTempDir();
|
||||
const databasePath = join(root, "sessions.sqlite");
|
||||
const env = new NodeExecutionEnv({ cwd: root });
|
||||
const sqlite = createNodeSqliteFactory();
|
||||
const repo = new SqliteSessionRepo({ env, sqlite, databasePath });
|
||||
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.getStorage().setLeafId(rootId);
|
||||
const secondChildId = await session.appendMessage(createAssistantMessage("second child"));
|
||||
|
||||
const db = await sqlite.open(databasePath);
|
||||
try {
|
||||
const branchRows = await db
|
||||
.prepare(
|
||||
"SELECT branch_id, entry_id, entry_seq FROM branch_entries WHERE session_id = ? ORDER BY branch_id, entry_seq",
|
||||
)
|
||||
.all<{ branch_id: string; entry_id: string; entry_seq: number }>("session-1");
|
||||
const branchIds = [...new Set(branchRows.map((row) => row.branch_id))];
|
||||
expect(branchIds).toHaveLength(3);
|
||||
expect(branchRows.filter((row) => row.entry_id === rootId)).toHaveLength(3);
|
||||
expect(branchRows.filter((row) => row.entry_id === firstChildId)).toHaveLength(1);
|
||||
expect(branchRows.filter((row) => row.entry_id === secondChildId)).toHaveLength(1);
|
||||
} finally {
|
||||
await db.close();
|
||||
}
|
||||
});
|
||||
|
||||
it("reopens using branch materialization and session summary state", async () => {
|
||||
const root = createTempDir();
|
||||
const databasePath = join(root, "sessions.sqlite");
|
||||
const env = new NodeExecutionEnv({ cwd: root });
|
||||
const repo = new SqliteSessionRepo({ env, sqlite: createNodeSqliteFactory(), databasePath });
|
||||
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.getStorage().setLeafId(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({
|
||||
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 SqliteSessionRepo({ env, sqlite: createNodeSqliteFactory(), databasePath });
|
||||
const session = await repo.create({ cwd: root, id: "session-1" });
|
||||
await session.appendMessage(createUserMessage("one"));
|
||||
await session.appendMessage(createAssistantMessage("two"));
|
||||
await session.appendMessage(createUserMessage("three"));
|
||||
|
||||
expect((await session.getEntries({ limit: 2 })).map((entry) => entry.type)).toEqual(["message", "message"]);
|
||||
expect((await session.getEntries({ afterEntrySeq: 2, limit: 2 })).map((entry) => entry.type)).toEqual([
|
||||
"message",
|
||||
"message",
|
||||
]);
|
||||
});
|
||||
|
||||
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 () => {
|
||||
throw new Error("insert failed");
|
||||
});
|
||||
}
|
||||
return new ThrowingStatement(async () => ({ changes: 1 }));
|
||||
});
|
||||
const sqlite: SqliteDatabaseFactory = {
|
||||
open: async () => db,
|
||||
};
|
||||
const env = new NodeExecutionEnv({ cwd: root });
|
||||
const repo = new SqliteSessionRepo({ env, sqlite, databasePath: join(root, "sessions.sqlite") });
|
||||
|
||||
await expect(repo.create({ cwd: root, id: "session-1" })).rejects.toThrow("insert failed");
|
||||
expect(db.closeCount).toBe(1);
|
||||
});
|
||||
|
||||
it("closes the database when open fails after openDatabase succeeds", async () => {
|
||||
const root = createTempDir();
|
||||
const db = new CountingDatabase((sql) => {
|
||||
if (sql.includes("FROM sessions WHERE id = ?")) {
|
||||
return new ThrowingStatement(async () => ({ changes: 0 }));
|
||||
}
|
||||
return new ThrowingStatement(async () => ({ changes: 1 }));
|
||||
});
|
||||
const sqlite: SqliteDatabaseFactory = {
|
||||
open: async () => db,
|
||||
};
|
||||
const env = new NodeExecutionEnv({ cwd: root });
|
||||
const repo = new SqliteSessionRepo({ env, sqlite, databasePath: join(root, "sessions.sqlite") });
|
||||
const metadata: SqliteSessionMetadata = {
|
||||
id: "missing",
|
||||
createdAt: new Date().toISOString(),
|
||||
cwd: root,
|
||||
path: join(root, "sessions.sqlite"),
|
||||
};
|
||||
writeFileSync(metadata.path, "");
|
||||
|
||||
await expect(repo.open(metadata)).rejects.toThrow("Session not found: missing");
|
||||
expect(db.closeCount).toBe(1);
|
||||
});
|
||||
|
||||
it("closes the source storage after fork reads its entries", async () => {
|
||||
const root = createTempDir();
|
||||
const databasePath = join(root, "sessions.sqlite");
|
||||
const env = new NodeExecutionEnv({ cwd: root });
|
||||
const repo = new SqliteSessionRepo({ env, sqlite: createNodeSqliteFactory(), databasePath });
|
||||
let cleanupCount = 0;
|
||||
const sourceStorage = {
|
||||
async getEntries() {
|
||||
return [];
|
||||
},
|
||||
async getPathToRootOrCompaction() {
|
||||
return [];
|
||||
},
|
||||
async cleanup() {
|
||||
cleanupCount += 1;
|
||||
},
|
||||
} as const;
|
||||
const originalOpen = repo.open.bind(repo);
|
||||
repo.open = async () =>
|
||||
({
|
||||
getStorage() {
|
||||
return sourceStorage;
|
||||
},
|
||||
}) as never;
|
||||
|
||||
try {
|
||||
await repo.fork(
|
||||
{
|
||||
id: "session-1",
|
||||
createdAt: new Date().toISOString(),
|
||||
cwd: root,
|
||||
path: databasePath,
|
||||
},
|
||||
{ cwd: root, id: "session-2" },
|
||||
);
|
||||
} finally {
|
||||
repo.open = originalOpen;
|
||||
}
|
||||
|
||||
expect(cleanupCount).toBe(1);
|
||||
});
|
||||
|
||||
it("restores in-memory state when appendEntry fails after mutating caches", async () => {
|
||||
const root = createTempDir();
|
||||
const databasePath = join(root, "sessions.sqlite");
|
||||
const sqlite = createNodeSqliteFactory();
|
||||
const db = await sqlite.open(databasePath);
|
||||
await applyMigrations(db);
|
||||
const storage = await SqliteSessionStorage.create(db, databasePath, {
|
||||
cwd: root,
|
||||
sessionId: "session-1",
|
||||
});
|
||||
const originalPrepare = db.prepare.bind(db);
|
||||
db.prepare = (sql: string) => {
|
||||
if (sql.startsWith("UPDATE sessions SET active_leaf_id = ?")) {
|
||||
return new ThrowingStatement(async () => {
|
||||
throw new Error("active leaf update failed");
|
||||
});
|
||||
}
|
||||
return originalPrepare(sql);
|
||||
};
|
||||
|
||||
await expect(
|
||||
storage.appendEntry({
|
||||
type: "message",
|
||||
id: "root",
|
||||
parentId: null,
|
||||
timestamp: new Date().toISOString(),
|
||||
message: createUserMessage("root"),
|
||||
}),
|
||||
).rejects.toMatchObject({ code: "storage" });
|
||||
expect(await storage.getLeafId()).toBeNull();
|
||||
expect(await storage.getEntry("root")).toBeUndefined();
|
||||
expect(await storage.getEntries()).toEqual([]);
|
||||
await db.close();
|
||||
});
|
||||
|
||||
it("materializes session summary fields transactionally", async () => {
|
||||
const root = createTempDir();
|
||||
const databasePath = join(root, "sessions.sqlite");
|
||||
const env = new NodeExecutionEnv({ cwd: root });
|
||||
const sqlite = createNodeSqliteFactory();
|
||||
const repo = new SqliteSessionRepo({ 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",
|
||||
model: "claude-sonnet-4-5",
|
||||
usage: {
|
||||
input: 100,
|
||||
output: 25,
|
||||
cacheRead: 40,
|
||||
cacheWrite: 10,
|
||||
totalTokens: 175,
|
||||
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, {
|
||||
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 },
|
||||
});
|
||||
await session.moveTo(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 },
|
||||
},
|
||||
});
|
||||
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);
|
||||
} finally {
|
||||
await db.close();
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,21 @@
|
||||
import { join } from "node:path";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { createNodeSqliteFactory } from "../../../storage/sqlite-node/src/index.ts";
|
||||
import { createTempDir } from "./session-test-utils.ts";
|
||||
|
||||
describe("sqlite-node adapter", () => {
|
||||
it("supports node:sqlite-style named parameters", async () => {
|
||||
const root = createTempDir();
|
||||
const databasePath = join(root, "adapter.sqlite");
|
||||
const sqlite = createNodeSqliteFactory();
|
||||
const db = await sqlite.open(databasePath);
|
||||
try {
|
||||
await db.exec("CREATE TABLE items (id INTEGER PRIMARY KEY, text TEXT NOT NULL)");
|
||||
await db.prepare("INSERT INTO items (id, text) VALUES ($id, $text)").run({ $id: 1, $text: "hello" });
|
||||
const row = await db.prepare("SELECT text FROM items WHERE id = $id").get<{ text: string }>({ $id: 1 });
|
||||
expect(row).toEqual({ text: "hello" });
|
||||
} finally {
|
||||
await db.close();
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -4,7 +4,13 @@ import { describe, expect, it } from "vitest";
|
||||
import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts";
|
||||
import { JsonlSessionStorage, loadJsonlSessionMetadata } from "../../src/harness/session/jsonl-storage.ts";
|
||||
import { InMemorySessionStorage } from "../../src/harness/session/memory-storage.ts";
|
||||
import { type MessageEntry, ok, type SessionMetadata } from "../../src/harness/types.ts";
|
||||
import {
|
||||
type BranchSummaryEntry,
|
||||
type CompactionEntry,
|
||||
type MessageEntry,
|
||||
ok,
|
||||
type SessionMetadata,
|
||||
} from "../../src/harness/types.ts";
|
||||
import { createAssistantMessage, createTempDir, createUserMessage } from "./session-test-utils.ts";
|
||||
|
||||
describe("InMemorySessionStorage", () => {
|
||||
@@ -80,7 +86,74 @@ describe("InMemorySessionStorage", () => {
|
||||
expect(await storage.getLabel("entry-1")).toBeUndefined();
|
||||
});
|
||||
|
||||
it("walks paths to root", async () => {
|
||||
it("includes summary-entry usage in session stats", async () => {
|
||||
const assistant: MessageEntry = {
|
||||
type: "message",
|
||||
id: "assistant",
|
||||
parentId: null,
|
||||
timestamp: "2026-01-01T00:00:00.000Z",
|
||||
message: {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text: "reply" }],
|
||||
api: "anthropic-messages",
|
||||
provider: "anthropic",
|
||||
model: "claude-sonnet-4-5",
|
||||
usage: {
|
||||
input: 10,
|
||||
output: 20,
|
||||
cacheRead: 30,
|
||||
cacheWrite: 40,
|
||||
totalTokens: 100,
|
||||
cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: 0,
|
||||
},
|
||||
};
|
||||
const compaction: CompactionEntry = {
|
||||
type: "compaction",
|
||||
id: "compaction",
|
||||
parentId: "assistant",
|
||||
timestamp: "2026-01-01T00:00:01.000Z",
|
||||
summary: "summary",
|
||||
firstKeptEntryId: "assistant",
|
||||
tokensBefore: 1234,
|
||||
usage: {
|
||||
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 branchSummary: BranchSummaryEntry = {
|
||||
type: "branch_summary",
|
||||
id: "branch-summary",
|
||||
parentId: "compaction",
|
||||
timestamp: "2026-01-01T00:00:02.000Z",
|
||||
fromId: "assistant",
|
||||
summary: "branch",
|
||||
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 },
|
||||
},
|
||||
};
|
||||
const storage = new InMemorySessionStorage({ entries: [assistant, compaction, branchSummary] });
|
||||
expect(await storage.getSessionStats()).toEqual({
|
||||
messageCount: 1,
|
||||
cachedTokens: 40,
|
||||
uncachedTokens: 68,
|
||||
totalTokens: 136,
|
||||
costTotal: 1.36,
|
||||
});
|
||||
});
|
||||
|
||||
it("walks paths to root or retained-tail compaction", async () => {
|
||||
const root: MessageEntry = {
|
||||
type: "message",
|
||||
id: "root",
|
||||
@@ -94,9 +167,29 @@ describe("InMemorySessionStorage", () => {
|
||||
parentId: "root",
|
||||
message: createAssistantMessage("child"),
|
||||
};
|
||||
const storage = new InMemorySessionStorage({ entries: [root, child] });
|
||||
expect((await storage.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
|
||||
expect(await storage.getPathToRoot(null)).toEqual([]);
|
||||
const compaction: CompactionEntry = {
|
||||
type: "compaction",
|
||||
id: "compaction",
|
||||
parentId: "child",
|
||||
timestamp: "2026-01-01T00:00:01.000Z",
|
||||
summary: "summary",
|
||||
firstKeptEntryId: "child",
|
||||
tokensBefore: 1234,
|
||||
retainedTail: [createAssistantMessage("child")],
|
||||
};
|
||||
const afterCompaction: MessageEntry = {
|
||||
...root,
|
||||
id: "after-compaction",
|
||||
parentId: "compaction",
|
||||
message: createUserMessage("after"),
|
||||
};
|
||||
const storage = new InMemorySessionStorage({ entries: [root, child, compaction, afterCompaction] });
|
||||
expect((await storage.getPathToRootOrCompaction("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
|
||||
expect((await storage.getPathToRootOrCompaction("after-compaction")).map((entry) => entry.id)).toEqual([
|
||||
"compaction",
|
||||
"after-compaction",
|
||||
]);
|
||||
expect(await storage.getPathToRootOrCompaction(null)).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -255,7 +348,7 @@ describe("JsonlSessionStorage", () => {
|
||||
const reloaded = await JsonlSessionStorage.open(env, filePath);
|
||||
expect(await reloaded.getLeafId()).toBe("root");
|
||||
expect((await reloaded.getEntries()).at(-1)).toMatchObject({ type: "leaf", targetId: "root" });
|
||||
expect((await loaded.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
|
||||
expect((await loaded.getPathToRootOrCompaction("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
|
||||
});
|
||||
|
||||
it("finds entries by type", async () => {
|
||||
@@ -309,6 +402,76 @@ describe("JsonlSessionStorage", () => {
|
||||
expect(await loaded.getLabel("entry-1")).toBeUndefined();
|
||||
});
|
||||
|
||||
it("includes summary-entry usage in session stats", async () => {
|
||||
const dir = createTempDir();
|
||||
const env = new NodeExecutionEnv({ cwd: dir });
|
||||
const filePath = join(dir, "session.jsonl");
|
||||
const storage = await JsonlSessionStorage.create(env, filePath, { cwd: dir, sessionId: "session-1" });
|
||||
await storage.appendEntry({
|
||||
type: "message",
|
||||
id: "assistant",
|
||||
parentId: null,
|
||||
timestamp: "2026-01-01T00:00:00.000Z",
|
||||
message: {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text: "reply" }],
|
||||
api: "anthropic-messages",
|
||||
provider: "anthropic",
|
||||
model: "claude-sonnet-4-5",
|
||||
usage: {
|
||||
input: 10,
|
||||
output: 20,
|
||||
cacheRead: 30,
|
||||
cacheWrite: 40,
|
||||
totalTokens: 100,
|
||||
cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: 0,
|
||||
},
|
||||
});
|
||||
await storage.appendEntry({
|
||||
type: "compaction",
|
||||
id: "compaction",
|
||||
parentId: "assistant",
|
||||
timestamp: "2026-01-01T00:00:01.000Z",
|
||||
summary: "summary",
|
||||
firstKeptEntryId: "assistant",
|
||||
tokensBefore: 1234,
|
||||
usage: {
|
||||
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 },
|
||||
},
|
||||
});
|
||||
await storage.appendEntry({
|
||||
type: "branch_summary",
|
||||
id: "branch-summary",
|
||||
parentId: "compaction",
|
||||
timestamp: "2026-01-01T00:00:02.000Z",
|
||||
fromId: "assistant",
|
||||
summary: "branch",
|
||||
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 },
|
||||
},
|
||||
});
|
||||
expect(await storage.getSessionStats()).toEqual({
|
||||
messageCount: 1,
|
||||
cachedTokens: 40,
|
||||
uncachedTokens: 68,
|
||||
totalTokens: 136,
|
||||
costTotal: 1.36,
|
||||
});
|
||||
});
|
||||
|
||||
it("reads session metadata through the line-reading filesystem operation", async () => {
|
||||
const dir = createTempDir();
|
||||
const filePath = join(dir, "session.jsonl");
|
||||
|
||||
Reference in New Issue
Block a user