sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

refactor: update sqlite for lanes (#7591)

Browse source 
* feat: align sqlite storage with session lanes

* fix: delete lane from entry, action

* fix: errors and types

* fix: tests

* fix: usage records instead of session materialized

* chore: 1 mig

* fix

* refactor(agent): move harness experimental changes to split branch

* refactor and cleanup

* fix: use assertJsonSerializable for session metadata

* fix(sqlite): reject corrupt lane leafs on open

* fix: rebuild cache

* refactor: branch cache

* feat: session stats

* fix: enforce leases

* feat(agent): enforce fenced SQLite writer leases

* fix(agent): keep SQLite transactions synchronous

* fix(agent): align SQLite storage with conformance suite

---------

Co-authored-by: Christian Klotz <hello@christianklotz.co.uk>
This commit is contained in:
Cristina Poncela Cubeiro 2026-08-04 15:56:56 +02:00 • committed by GitHub
parent 04133eb01b
commit f119b01cb1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
37 changed files with 2536 additions and 2044 deletions
Download patch file Download diff file Expand all files Collapse all files

4
packages/agent/docs/harness-v2.md
View file

@ -1649,7 +1649,7 @@ lane_moves (session_id, seq, lane, leaf_id) -- history; getLog parity
facts (session_id, seq, kind, key, value) -- name, labels; latest by seq
branch_entries (session_id, branch_id, entry_id, entry_seq, entry_type, custom_type)
branch_tips (session_id, branch_id, tip_id) -- PRIMARY KEY (session_id, tip_id)
leases (session_id, owner, heartbeat) -- writer claim
leases (session_id, owner_id, fence, expires_at_ms) -- writer claim
-- indexes
records: (session_id, lane, type, seq), (session_id, lane, type, op_kind, seq)
@ -1657,6 +1657,8 @@ branch_entries: (session_id, branch_id, entry_type, entry_seq)
(session_id, entry_id) -- reverse lookup: entry → branches
```
`leases` enforces one writer per session with expiring, fenced claims. Storage renews the claim inside every write transaction and while idle. Repository-owned cleanup releases only its matching owner and fence.
`branch_entries` and `branch_tips` are a private read cache. No interface exposes them; no other backend has them; rebuilding them from parent pointers is an explicit repair operation, never a runtime fallback.
Two invariants carry the whole design:

2
packages/agent/src/harness/experimental/session/session.ts
View file

@ -37,7 +37,7 @@ function assertValidCursor(afterSeq: number | undefined): void {
}
}
function assertJsonSerializable(value: unknown): void {
export function assertJsonSerializable(value: unknown): void {
const active = new WeakSet<object>();
const stack: JsonValidationFrame[] = [{ value }];
while (stack.length > 0) {

51
packages/agent/test/harness/branch-query.test.ts
View file

@ -171,7 +171,7 @@ describe("bounded session branch queries", () => {
const db = await sqlite.open(databasePath);
try {
await db
.prepare("UPDATE session_entries SET payload = ? WHERE session_id = ? AND id = ?")
.prepare("UPDATE entries SET payload = ? WHERE session_id = ? AND id = ?")
.run("not json", "bounded-sqlite", middleId);
const branch = await db
.prepare("SELECT branch_id FROM branch_entries WHERE session_id = ? AND entry_id = ?")
@ -187,21 +187,6 @@ describe("bounded session branch queries", () => {
expect((await session.findEntriesOnBranch({ start: tailId, stopAtId: tailId })).map((entry) => entry.id)).toEqual(
[tailId],
);
const inspection = await sqlite.open(databasePath);
try {
const repaired = await inspection
.prepare(
`SELECT entry_id FROM branch_entries
WHERE session_id = ? AND branch_id = (
SELECT branch_id FROM branch_entries WHERE session_id = ? AND entry_id = ? LIMIT 1
)
ORDER BY entry_seq`,
)
.all<{ entry_id: string }>("bounded-sqlite", "bounded-sqlite", tailId);
expect(repaired.map((row) => row.entry_id)).toEqual([rootId, tailId]);
} finally {
await inspection.close();
}
expect(
(
await session.findEntriesOnBranch({
@ -214,7 +199,7 @@ describe("bounded session branch queries", () => {
).toEqual([rootId]);
await expect(session.findEntriesOnBranch({ start: tailId, limit: 2 })).rejects.toMatchObject({
code: "invalid_entry",
message: expect.stringContaining(`failed to decode entry ${middleId}`),
message: expect.stringContaining(`Entry ${middleId} not found`),
});
});
@ -236,7 +221,7 @@ describe("bounded session branch queries", () => {
const db = await sqlite.open(databasePath);
try {
await db
.prepare("UPDATE session_entries SET payload = ? WHERE session_id = ? AND id = ?")
.prepare("UPDATE entries SET payload = ? WHERE session_id = ? AND id = ?")
.run("{}", "invalid-filtered-sqlite", customId);
} finally {
await db.close();
@ -249,7 +234,7 @@ describe("bounded session branch queries", () => {
const invalidJsonDb = await sqlite.open(databasePath);
try {
await invalidJsonDb
.prepare("UPDATE session_entries SET payload = ? WHERE session_id = ? AND id = ?")
.prepare("UPDATE entries SET payload = ? WHERE session_id = ? AND id = ?")
.run("not json", "invalid-filtered-sqlite", customId);
} finally {
await invalidJsonDb.close();
@ -277,7 +262,7 @@ describe("bounded session branch queries", () => {
const db = await sqlite.open(databasePath);
try {
await db
.prepare("UPDATE session_entries SET parent_id = ? WHERE session_id = ? AND id = ?")
.prepare("UPDATE entries SET parent_id = ? WHERE session_id = ? AND id = ?")
.run("missing-parent", "bounded-corrupt-sqlite", childId);
} finally {
await db.close();
@ -289,17 +274,17 @@ describe("bounded session branch queries", () => {
(await session.findEntriesOnBranch({ start: childId, stopAtType: "message" })).map((entry) => entry.id),
).toEqual([childId]);
await expect(session.findEntriesOnBranch({ start: childId })).rejects.toMatchObject({
code: "invalid_session",
code: "invalid_entry",
message: expect.stringContaining("Entry missing-parent not found"),
});
const cycleDb = await sqlite.open(databasePath);
try {
await cycleDb
.prepare("UPDATE session_entries SET parent_id = ? WHERE session_id = ? AND id = ?")
.prepare("UPDATE entries SET parent_id = ? WHERE session_id = ? AND id = ?")
.run(rootId, "bounded-corrupt-sqlite", childId);
await cycleDb
.prepare("UPDATE session_entries SET parent_id = ? WHERE session_id = ? AND id = ?")
.prepare("UPDATE entries SET parent_id = ? WHERE session_id = ? AND id = ?")
.run(childId, "bounded-corrupt-sqlite", rootId);
} finally {
await cycleDb.close();
@ -311,24 +296,8 @@ describe("bounded session branch queries", () => {
(await session.findEntriesOnBranch({ start: childId, stopAtType: "message" })).map((entry) => entry.id),
).toEqual([childId]);
await expect(session.findEntriesOnBranch({ start: childId })).rejects.toMatchObject({
code: "invalid_session",
message: expect.stringContaining(`cycle in parent chain at entry ${childId}`),
code: "invalid_entry",
message: expect.stringContaining(`Entry ${childId} not found`),
});
});
it("provides identical SQLite query semantics", async () => {
const root = createTempDir();
const repo = new SqliteSessionRepository({
env: new NodeExecutionEnv({ cwd: root }),
sqlite: createNodeSqliteFactory(),
databasePath: join(root, "sessions.sqlite"),
});
ownedRepositories.push(repo);
const session = await repo.create({ id: "sqlite", cwd: root });
const expected = await verifyBranchQueries(session);
const reopened = await repo.open(await session.getMetadata());
expect(
(await reopened.findEntriesOnBranch({ start: expected.tail, order: "oldestFirst" })).map((entry) => entry.id),
).toEqual(expected.fullPath);
});
});

58
packages/agent/test/harness/experimental/session/sqlite.test.ts Normal file
View file

@ -0,0 +1,58 @@
import { join } from "node:path";
import type { SessionMetadata, SessionRepo } from "@earendil-works/pi-agent-core/experimental";
import { describe, it } from "vitest";
import {
createNodeSqliteFactory,
type SqliteSessionMetadata,
SqliteSessionRepository,
} from "../../../../../storage/sqlite-node/src/index.ts";
import { NodeExecutionEnv } from "../../../../src/harness/env/nodejs.ts";
import {
createSessionBackendConformance,
type SessionBackendFixture,
} from "../../../../src/harness/experimental/session/testing/index.ts";
import { createTempDir } from "../../session-test-utils.ts";
function requireSqliteMetadata(metadata: SessionMetadata): SqliteSessionMetadata {
const cwd = "cwd" in metadata ? metadata.cwd : undefined;
if (typeof cwd !== "string") {
throw new Error(`Expected SQLite metadata for session ${metadata.id}`);
}
const path = "path" in metadata ? metadata.path : undefined;
if (typeof path !== "string") {
throw new Error(`Expected SQLite metadata for session ${metadata.id}`);
}
return { ...metadata, cwd, path };
}
const conformance = createSessionBackendConformance(async () => {
const root = createTempDir();
const sqliteRepository = new SqliteSessionRepository({
env: new NodeExecutionEnv({ cwd: root }),
sqlite: createNodeSqliteFactory(),
databasePath: join(root, "sessions.sqlite"),
});
const repository: SessionRepo = {
create: (options = {}) => sqliteRepository.create({ ...options, cwd: root }),
open: (metadata) => sqliteRepository.open(requireSqliteMetadata(metadata)),
list: () => sqliteRepository.list(),
delete: (metadata) => sqliteRepository.delete(requireSqliteMetadata(metadata)),
fork: (source, options = {}) => sqliteRepository.fork(requireSqliteMetadata(source), { ...options, cwd: root }),
};
return {
repository,
async [Symbol.asyncDispose]() {
await sqliteRepository.close();
},
} satisfies SessionBackendFixture;
});
describe("SqliteSessionRepository 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());
}
});
}
});

128
packages/agent/test/harness/session-test-utils.ts
View file

@ -2,7 +2,18 @@ import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type {
BranchSummaryEntry,
CompactionEntry,
Session as CoreSession,
Entry,
MessageEntry,
ModelChangeEntry,
ThinkingLevelChangeEntry,
} from "@earendil-works/pi-agent-core/experimental";
import type { Usage } from "@earendil-works/pi-ai";
import { afterEach } from "vitest";
import type { SqliteSessionMetadata } from "../../../storage/sqlite-node/src/index.ts";
import { InMemorySessionRepository } from "../../src/harness/session/memory-repo.ts";
import type { Session } from "../../src/harness/session/session.ts";
@ -38,6 +49,123 @@ export function createAssistantMessage(text: string): AgentMessage {
};
}
export type SqliteTestSession = CoreSession<SqliteSessionMetadata>;
export type SqliteTestMessage = MessageEntry["message"];
export async function appendSqliteCompaction(
session: SqliteTestSession,
summary: string,
_firstKeptEntryId: string | undefined,
tokensBefore: number,
details?: unknown,
_fromHook?: boolean,
usage?: Usage,
retainedTail: SqliteTestMessage[] = [],
): Promise<string> {
const provisioned = {
type: "compaction",
id: session.idGenerator.next(),
summary,
retainedTail,
tokensBefore,
...(details === undefined ? {} : { details }),
...(usage === undefined ? {} : { usage }),
} satisfies Omit<CompactionEntry, "parentId" | "seq" | "timestamp">;
const entry = await session.appendEntry(provisioned, "main");
return entry.id;
}
export async function moveSqliteMainLane(
session: SqliteTestSession,
entryId: string | null,
summary?: { summary: string; details?: unknown; usage?: Usage; fromHook?: boolean },
): Promise<string | undefined> {
await session.moveLane("main", entryId);
if (!summary) return undefined;
const provisioned = {
type: "branch_summary",
id: session.idGenerator.next(),
fromId: entryId ?? "root",
summary: summary.summary,
...(summary.details === undefined ? {} : { details: summary.details }),
...(summary.usage === undefined ? {} : { usage: summary.usage }),
} satisfies Omit<BranchSummaryEntry, "parentId" | "seq" | "timestamp">;
const entry = await session.appendEntry(provisioned, "main");
return entry.id;
}
export async function getSqliteBranch(session: SqliteTestSession, fromId?: string | null): Promise<Entry[]> {
const start = fromId === undefined ? await session.getLeafId() : fromId;
if (start === null) return [];
const newestWindow = await session.findEntriesOnBranch({ start, stopAtType: "compaction" });
return newestWindow.reverse();
}
export async function getSqliteEntries(
session: SqliteTestSession,
options?: { afterEntrySeq?: number; limit?: number },
): Promise<Entry[]> {
return session.findEntries({
order: "oldestFirst",
limit: options?.limit,
cursor: options?.afterEntrySeq === undefined ? undefined : { afterSeq: options.afterEntrySeq },
});
}
export async function appendSqliteSessionName(session: SqliteTestSession, name: string): Promise<void> {
await session.setName(name.replace(/[\r\n]+/g, " ").trim());
}
export async function appendSqliteLabel(
session: SqliteTestSession,
targetId: string,
label: string | undefined,
): Promise<void> {
await session.setLabel(targetId, label);
}
export async function buildSqliteContext(session: SqliteTestSession): Promise<{ messages: SqliteTestMessage[] }> {
const entries = await getSqliteBranch(session);
const messages = entries.flatMap((entry): SqliteTestMessage[] => {
if (entry.type === "message") return [entry.message];
if (entry.type === "compaction") return entry.retainedTail;
return [];
});
return { messages };
}
export async function appendSqliteThinkingLevelChange(
session: SqliteTestSession,
thinkingLevel: string,
): Promise<string> {
const entry = await session.appendEntry(
{
type: "thinking_level_change",
id: session.idGenerator.next(),
thinkingLevel,
} satisfies Omit<ThinkingLevelChangeEntry, "parentId" | "seq" | "timestamp">,
"main",
);
return entry.id;
}
export async function appendSqliteModelChange(
session: SqliteTestSession,
provider: string,
modelId: string,
): Promise<string> {
const entry = await session.appendEntry(
{
type: "model_change",
id: session.idGenerator.next(),
provider,
modelId,
} satisfies Omit<ModelChangeEntry, "parentId" | "seq" | "timestamp">,
"main",
);
return entry.id;
}
const tempDirs: string[] = [];
export function createTempDir(): string {

36
packages/agent/test/harness/sqlite-adapter.test.ts Normal file
View file

@ -0,0 +1,36 @@
import { describe, expect, it } from "vitest";
import { createNodeSqliteFactory } from "../../../storage/sqlite-node/src/index.ts";
describe("node:sqlite adapter", () => {
it("runs transaction callbacks synchronously", async () => {
const db = await createNodeSqliteFactory().open(":memory:");
try {
db.exec("CREATE TABLE values_table (value INTEGER NOT NULL)");
const result = db.transaction(() => {
db.prepare("INSERT INTO values_table (value) VALUES (?)").run(42);
return "committed";
});
expect(result).toBe("committed");
expect(db.prepare("SELECT value FROM values_table").get()).toEqual({ value: 42 });
} finally {
db.close();
}
});
it("rejects asynchronous transaction callbacks", async () => {
const db = await createNodeSqliteFactory().open(":memory:");
try {
db.exec("CREATE TABLE values_table (value INTEGER NOT NULL)");
const asynchronous = async () => {
db.prepare("INSERT INTO values_table (value) VALUES (?)").run(42);
await Promise.resolve();
};
expect(() => db.transaction(asynchronous)).toThrow("SQLite transaction callbacks must be synchronous");
await Promise.resolve();
expect(db.prepare("SELECT value FROM values_table").all()).toEqual([]);
} finally {
db.close();
}
});
});

139
packages/agent/test/harness/sqlite-branch-cache.test.ts
View file

@ -4,7 +4,13 @@ import { join } from "node:path";
import { describe, expect, it } from "vitest";
import { createNodeSqliteFactory, SqliteSessionRepository } from "../../../storage/sqlite-node/src/index.ts";
import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts";
import { createAssistantMessage, createUserMessage } from "./session-test-utils.ts";
import {
appendSqliteCompaction,
createAssistantMessage,
createUserMessage,
getSqliteBranch,
moveSqliteMainLane,
} from "./session-test-utils.ts";
function createTempDir(): string {
return mkdtempSync(join(tmpdir(), "pi-agent-sqlite-branch-cache-"));
@ -20,9 +26,9 @@ describe("SQLite branch cache", () => {
const session = await repo.create({ cwd: root, id: "session-1" });
const rootId = await session.appendMessage(createUserMessage("root"));
const keptId = await session.appendMessage(createUserMessage("kept"));
const compactionId = await session.appendCompaction("summary", keptId, 100);
const compactionId = await appendSqliteCompaction(session, "summary", keptId, 100);
await session.appendMessage(createAssistantMessage("first child"));
await session.moveTo(compactionId);
await moveSqliteMainLane(session, compactionId);
const branchedId = await session.appendMessage(createAssistantMessage("branched child"));
const db = await sqlite.open(databasePath);
@ -49,19 +55,19 @@ describe("SQLite branch cache", () => {
const session = await repo.create({ cwd: root, id: "session-1" });
const oldId = await session.appendMessage(createUserMessage("old"));
const keptId = await session.appendMessage(createUserMessage("kept"));
const compactionId = await session.appendCompaction("summary", keptId, 100);
const compactionId = await appendSqliteCompaction(session, "summary", keptId, 100);
const leafId = await session.appendMessage(createAssistantMessage("new"));
const db = await sqlite.open(databasePath);
try {
await db
.prepare("UPDATE session_entries SET payload = ? WHERE session_id = ? AND id = ?")
.prepare("UPDATE entries SET payload = ? WHERE session_id = ? AND id = ?")
.run("not json", "session-1", oldId);
} finally {
await db.close();
}
expect((await session.getBranch()).map((entry) => entry.id)).toEqual([keptId, compactionId, leafId]);
expect((await getSqliteBranch(session)).map((entry) => entry.id)).toEqual([compactionId, leafId]);
});
it("preserves nested compaction boundaries when reading the cache", async () => {
@ -71,7 +77,8 @@ describe("SQLite branch cache", () => {
const repo = new SqliteSessionRepository({ env, sqlite: createNodeSqliteFactory(), databasePath });
const session = await repo.create({ cwd: root, id: "session-1" });
const rootId = await session.appendMessage(createUserMessage("root"));
const firstCompactionId = await session.appendCompaction(
const firstCompactionId = await appendSqliteCompaction(
session,
"first summary",
rootId,
100,
@ -81,18 +88,15 @@ describe("SQLite branch cache", () => {
[],
);
const middleId = await session.appendMessage(createUserMessage("middle"));
const secondCompactionId = await session.appendCompaction("second summary", rootId, 200);
const secondCompactionId = await appendSqliteCompaction(session, "second summary", rootId, 200);
const leafId = await session.appendMessage(createAssistantMessage("new"));
expect((await session.getBranch()).map((entry) => entry.id)).toEqual([
firstCompactionId,
middleId,
secondCompactionId,
leafId,
]);
expect(firstCompactionId).not.toBe(secondCompactionId);
expect(middleId).not.toBe(leafId);
expect((await getSqliteBranch(session)).map((entry) => entry.id)).toEqual([secondCompactionId, leafId]);
});
it("repairs a missing branch cache from canonical parent links", async () => {
it("fails loudly when the private branch cache is missing", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const env = new NodeExecutionEnv({ cwd: root });
@ -110,25 +114,11 @@ describe("SQLite branch cache", () => {
await db.close();
}
expect((await session.getBranch()).map((entry) => entry.id)).toEqual([rootId, childId]);
const inspection = await sqlite.open(databasePath);
try {
const rows = await inspection
.prepare("SELECT branch_id, entry_id FROM branch_entries WHERE session_id = ? ORDER BY entry_seq")
.all<{ branch_id: string; entry_id: string }>("session-1");
expect(rows.map((row) => row.entry_id)).toEqual([rootId, childId]);
const tip = await inspection
.prepare("SELECT branch_id, tip_id FROM branch_tips WHERE session_id = ?")
.get<{ branch_id: string; tip_id: string }>("session-1");
expect(rows).toHaveLength(2);
expect(tip).toEqual({ branch_id: rows[0]?.branch_id, tip_id: childId });
} finally {
await inspection.close();
}
expect(rootId).not.toBe(childId);
await expect(getSqliteBranch(session)).rejects.toMatchObject({ code: "invalid_entry" });
});
it("rolls back an interrupted branch cache repair", async () => {
it("does not repair the private branch cache during normal reads", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const env = new NodeExecutionEnv({ cwd: root });
@ -141,35 +131,49 @@ describe("SQLite branch cache", () => {
try {
await db.prepare("DELETE FROM branch_tips WHERE session_id = ?").run("session-1");
await db.prepare("DELETE FROM branch_entries WHERE session_id = ?").run("session-1");
await db.exec(`
CREATE TRIGGER fail_branch_cache_repair
BEFORE INSERT ON branch_tips
BEGIN
SELECT RAISE(ABORT, 'repair failed');
END;
`);
} finally {
await db.close();
}
await expect(session.getBranch()).rejects.toMatchObject({
code: "storage",
message: expect.stringContaining("Failed to rebuild SQLite branch cache"),
});
await expect(getSqliteBranch(session)).rejects.toMatchObject({ code: "invalid_entry" });
const inspection = await sqlite.open(databasePath);
try {
expect(
await inspection.prepare("SELECT entry_id FROM branch_entries WHERE session_id = ?").all("session-1"),
).toEqual([]);
await inspection.exec("DROP TRIGGER fail_branch_cache_repair");
} finally {
await inspection.close();
}
expect(await session.getBranch()).toHaveLength(1);
});
it("repairs a missing source branch cache while forking transactionally", async () => {
it("repairs the private branch cache explicitly", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const env = new NodeExecutionEnv({ cwd: root });
const sqlite = createNodeSqliteFactory();
const repo = new SqliteSessionRepository({ 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"));
const metadata = await session.getMetadata();
const db = await sqlite.open(databasePath);
try {
await db.prepare("DELETE FROM branch_tips WHERE session_id = ?").run("session-1");
await db.prepare("DELETE FROM branch_entries WHERE session_id = ?").run("session-1");
} finally {
await db.close();
}
await expect(getSqliteBranch(session)).rejects.toMatchObject({ code: "invalid_entry" });
await repo.repairBranchCache(metadata);
expect((await getSqliteBranch(session)).map((entry) => entry.id)).toEqual([rootId, childId]);
});
it("fails when forking from a source with a missing branch cache", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const env = new NodeExecutionEnv({ cwd: root });
@ -187,16 +191,18 @@ describe("SQLite branch cache", () => {
await db.close();
}
const fork = await repo.fork(await source.getMetadata(), {
cwd: root,
id: "fork",
entryId: childId,
position: "at",
});
expect((await fork.getEntries()).map((entry) => entry.id)).toEqual([rootId, childId]);
expect(rootId).not.toBe(childId);
await expect(
repo.fork(await source.getMetadata(), {
cwd: root,
id: "fork",
entryId: childId,
position: "at",
}),
).rejects.toMatchObject({ code: "invalid_fork_target" });
});
it("repairs a stale branch cache from canonical parent links", async () => {
it("fails when the private branch cache is stale", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const env = new NodeExecutionEnv({ cwd: root });
@ -210,31 +216,16 @@ describe("SQLite branch cache", () => {
const db = await sqlite.open(databasePath);
try {
await db
.prepare("UPDATE session_entries SET parent_id = ? WHERE session_id = ? AND id = ?")
.prepare("UPDATE entries SET parent_id = ? WHERE session_id = ? AND id = ?")
.run(rootId, "session-1", leafId);
} finally {
await db.close();
}
expect(
(await session.findEntriesOnBranch({ start: leafId, order: "oldestFirst" })).map((entry) => entry.id),
).toEqual([rootId, leafId]);
const inspection = await sqlite.open(databasePath);
try {
const rows = await inspection
.prepare(
`SELECT entry_id FROM branch_entries
WHERE session_id = ? AND branch_id = (
SELECT branch_id FROM branch_entries WHERE session_id = ? AND entry_id = ? LIMIT 1
)
ORDER BY entry_seq`,
)
.all<{ entry_id: string }>("session-1", "session-1", leafId);
expect(rows.map((row) => row.entry_id)).toEqual([rootId, leafId]);
expect(rows.map((row) => row.entry_id)).not.toContain(staleId);
} finally {
await inspection.close();
}
expect(staleId).not.toBe(leafId);
await expect(session.findEntriesOnBranch({ start: leafId, order: "oldestFirst" })).rejects.toMatchObject({
code: "invalid_entry",
});
});
it("deletes branch entries and tips with the session", async () => {

136
packages/agent/test/harness/sqlite-leases.test.ts Normal file
View file

@ -0,0 +1,136 @@
import { join } from "node:path";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
createNodeSqliteFactory,
type SqliteSessionMetadata,
SqliteSessionRepository,
} from "../../../storage/sqlite-node/src/index.ts";
import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts";
import { createTempDir, createUserMessage } from "./session-test-utils.ts";
const repositories: SqliteSessionRepository[] = [];
function createRepository(root: string, databasePath: string, lease?: { ttlMs: number; heartbeatIntervalMs: number }) {
const repository = new SqliteSessionRepository({
env: new NodeExecutionEnv({ cwd: root }),
sqlite: createNodeSqliteFactory(),
databasePath,
writerLease: lease,
});
repositories.push(repository);
return repository;
}
afterEach(async () => {
vi.useRealTimers();
for (const repository of repositories.splice(0)) await repository.close();
});
describe("SQLite session writer leases", () => {
it("rejects a second writer until the first session releases its claim", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const firstRepository = createRepository(root, databasePath);
const secondRepository = createRepository(root, databasePath);
const first = await firstRepository.create({ cwd: root, id: "session-1" });
const metadata = await first.getMetadata();
await expect(secondRepository.open(metadata)).rejects.toMatchObject({
code: "storage",
message: expect.stringContaining("already has an active writer"),
});
await firstRepository.close();
const second = await secondRepository.open(metadata);
await expect(second.appendMessage(createUserMessage("new owner"))).resolves.toBeTypeOf("string");
});
it("fences a stale owner after an expired lease is acquired by another writer", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const lease = { ttlMs: 120_000, heartbeatIntervalMs: 60_000 };
const firstRepository = createRepository(root, databasePath, lease);
const secondRepository = createRepository(root, databasePath, lease);
const first = await firstRepository.create({ cwd: root, id: "session-1" });
const metadata = await first.getMetadata();
const sqlite = createNodeSqliteFactory();
const db = await sqlite.open(databasePath);
try {
await db.prepare("UPDATE leases SET expires_at_ms = 0 WHERE session_id = ?").run(metadata.id);
} finally {
await db.close();
}
const second = await secondRepository.open(metadata);
await expect(first.appendMessage(createUserMessage("stale owner"))).rejects.toMatchObject({
code: "storage",
message: expect.stringContaining("writer lease was lost"),
});
expect(await second.findEntries()).toEqual([]);
const inspection = await sqlite.open(databasePath);
let currentLease: { owner_id: string; fence: number } | undefined;
try {
currentLease = await inspection
.prepare("SELECT owner_id, fence FROM leases WHERE session_id = ?")
.get<{ owner_id: string; fence: number }>(metadata.id);
expect(currentLease?.fence).toBe(2);
} finally {
await inspection.close();
}
await firstRepository.close();
const afterStaleClose = await sqlite.open(databasePath);
try {
expect(
await afterStaleClose.prepare("SELECT owner_id, fence FROM leases WHERE session_id = ?").get(metadata.id),
).toEqual(currentLease);
} finally {
await afterStaleClose.close();
}
await expect(second.appendMessage(createUserMessage("current owner"))).resolves.toBeTypeOf("string");
});
it("serializes lease-checked writes for sessions sharing one database connection", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const repository = createRepository(root, databasePath);
const first = await repository.create({ cwd: root, id: "session-1" });
const second = await repository.create({ cwd: root, id: "session-2" });
await expect(
Promise.all([
first.appendMessage(createUserMessage("first")),
second.appendMessage(createUserMessage("second")),
]),
).resolves.toHaveLength(2);
});
it("renews an idle writer lease with a heartbeat", async () => {
vi.useFakeTimers();
vi.setSystemTime(new Date("2026-01-01T00:00:00.000Z"));
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const repository = createRepository(root, databasePath, { ttlMs: 30_000, heartbeatIntervalMs: 10_000 });
const session = await repository.create({ cwd: root, id: "session-1" });
const metadata = (await session.getMetadata()) as SqliteSessionMetadata;
const sqlite = createNodeSqliteFactory();
const readExpiry = async (): Promise<number | undefined> => {
const db = await sqlite.open(databasePath);
try {
return (
await db
.prepare("SELECT expires_at_ms FROM leases WHERE session_id = ?")
.get<{ expires_at_ms: number }>(metadata.id)
)?.expires_at_ms;
} finally {
await db.close();
}
};
const initialExpiry = await readExpiry();
await vi.advanceTimersByTimeAsync(10_000);
expect(await readExpiry()).toBe((initialExpiry ?? 0) + 10_000);
});
});

387
packages/agent/test/harness/sqlite-migrations.test.ts
View file

@ -3,9 +3,7 @@ import { tmpdir } from "node:os";
import { join } from "node:path";
import { describe, expect, it } from "vitest";
import {
applyMigrations,
createNodeSqliteFactory,
loadMigrations,
type SqliteDatabase,
type SqliteDatabaseFactory,
type SqliteRunResult,
@ -13,30 +11,38 @@ import {
SqliteSessionRepository,
type SqliteStatement,
} from "../../../storage/sqlite-node/src/index.ts";
import { SqliteSessionConnection } from "../../../storage/sqlite-node/src/sqlite/storage/index.ts";
import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts";
import { createAssistantMessage, createUserMessage } from "./session-test-utils.ts";
import {
appendSqliteCompaction,
appendSqliteLabel,
appendSqliteSessionName,
buildSqliteContext,
createAssistantMessage,
createUserMessage,
getSqliteEntries,
moveSqliteMainLane,
} from "./session-test-utils.ts";
function createTempDir(): string {
return mkdtempSync(join(tmpdir(), "pi-agent-sqlite-"));
}
class ThrowingStatement implements SqliteStatement {
private readonly onRun: () => Promise<SqliteRunResult>;
private readonly onRun: () => SqliteRunResult;
constructor(onRun: () => Promise<SqliteRunResult>) {
constructor(onRun: () => SqliteRunResult) {
this.onRun = onRun;
}
async run(..._params: unknown[]): Promise<SqliteRunResult> {
run(..._params: unknown[]): SqliteRunResult {
return this.onRun();
}
async get<TRow extends object>(..._params: unknown[]): Promise<TRow | undefined> {
get<TRow extends object>(..._params: unknown[]): TRow | undefined {
return undefined;
}
async all<TRow extends object>(..._params: unknown[]): Promise<TRow[]> {
all<TRow extends object>(..._params: unknown[]): TRow[] {
return [];
}
}
@ -49,17 +55,17 @@ class CountingDatabase implements SqliteDatabase {
this.statementFactory = statementFactory;
}
async exec(_sql: string): Promise<void> {}
exec(_sql: string): void {}
prepare(sql: string): SqliteStatement {
return this.statementFactory(sql);
}
async transaction<T>(fn: () => Promise<T>): Promise<T> {
transaction<T>(fn: () => T): T {
return fn();
}
async close(): Promise<void> {
close(): void {
this.closeCount += 1;
}
}
@ -80,9 +86,9 @@ function createCloseCountingSqliteFactory(): {
exec: (sql) => db.exec(sql),
prepare: (sql) => db.prepare(sql),
transaction: (fn) => db.transaction(fn),
async close() {
close() {
counts.closes += 1;
await db.close();
db.close();
},
};
},
@ -102,7 +108,7 @@ describe("SQLite migrations", () => {
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", "002_branch_tips.sql"]);
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 }>();
@ -110,16 +116,17 @@ describe("SQLite migrations", () => {
expect.arrayContaining([
"migrations",
"sessions",
"session_entries",
"entries",
"session_sequences",
"branch_entries",
"branch_tips",
"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");
expect(sessionColumns.map((column) => column.name)).not.toContain("leaf_id");
expect(tables.map((row) => row.name)).toEqual(
expect.arrayContaining(["lanes", "records", "lane_moves", "facts", "leases", "session_stats"]),
);
const branchIndexes = await db
.prepare("SELECT name FROM sqlite_master WHERE type = 'index' AND tbl_name = 'branch_entries'")
.all<{ name: string }>();
@ -128,10 +135,13 @@ describe("SQLite migrations", () => {
for (const tableName of [
"sessions",
"session_sequences",
"session_stats",
"branch_entries",
"branch_tips",
"session_materialized",
"entry_materialized",
"lanes",
"records",
"lane_moves",
"facts",
]) {
const table = tables.find((row) => row.name === tableName);
expect(table?.sql).toContain("WITHOUT ROWID");
@ -141,37 +151,6 @@ describe("SQLite migrations", () => {
}
});
it("clears legacy branch projections when adding explicit tips", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const sqlite = createNodeSqliteFactory();
const db = await sqlite.open(databasePath);
try {
const initial = (await loadMigrations()).find((migration) => migration.id === "001_initial.sql");
if (!initial) throw new Error("Missing initial SQLite migration");
await db.exec(initial.sql);
await db.exec("CREATE TABLE migrations (id TEXT PRIMARY KEY, applied_at TEXT NOT NULL)");
await db
.prepare("INSERT INTO migrations (id, applied_at) VALUES (?, ?)")
.run(initial.id, "2026-01-01T00:00:00.000Z");
await db
.prepare("INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) VALUES (?, ?, ?, ?)")
.run("session-1", "legacy-branch", "entry-1", 1);
await applyMigrations(db);
expect(await db.prepare("SELECT entry_id FROM branch_entries").all<{ entry_id: string }>()).toEqual([]);
expect(
await db.prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'branch_tips'").get(),
).toBeDefined();
expect(
(await db.prepare("SELECT id FROM migrations ORDER BY id").all<{ id: string }>()).map((row) => row.id),
).toEqual(["001_initial.sql", "002_branch_tips.sql"]);
} finally {
await db.close();
}
});
it("persists session metadata through create, list, open, and fork", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
@ -185,7 +164,8 @@ describe("SQLite migrations", () => {
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 reopened = await repo.open(sourceMetadata);
expect((await reopened.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, {
@ -209,8 +189,8 @@ describe("SQLite migrations", () => {
const db = await sqlite.open(databasePath);
try {
await db.exec(`
CREATE TRIGGER fail_fork_entry BEFORE INSERT ON session_entries
WHEN new.session_id = 'fork' AND new.entry_seq = 2
CREATE TRIGGER fail_fork_entry BEFORE INSERT ON entries
WHEN new.session_id = 'fork' AND new.seq = 2
BEGIN
SELECT RAISE(ABORT, 'fail fork');
END;
@ -228,14 +208,14 @@ END;
await inspection.prepare("SELECT id FROM sessions WHERE id = ?").get<{ id: string }>("fork"),
).toBeUndefined();
expect(
await inspection.prepare("SELECT id FROM session_entries WHERE session_id = ?").all<{ id: string }>("fork"),
await inspection.prepare("SELECT id FROM entries WHERE session_id = ?").all<{ id: string }>("fork"),
).toEqual([]);
} finally {
await inspection.close();
}
});
it("materializes active leaf id in sessions transactionally", async () => {
it("materializes the main lane leaf transactionally", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const env = new NodeExecutionEnv({ cwd: root });
@ -244,29 +224,18 @@ END;
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.moveTo(rootId);
await moveSqliteMainLane(session, 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);
if (!latestBranchRow) throw new Error("Missing latest branch row");
const branchTip = await db
.prepare("SELECT branch_id, tip_id FROM branch_tips WHERE session_id = ? AND branch_id = ?")
.get<{ branch_id: string; tip_id: string }>("session-1", latestBranchRow.branch_id);
expect(branchTip?.tip_id).toBe(latestSessionEntry?.id);
.prepare("SELECT leaf_id FROM lanes WHERE session_id = ? AND lane = ?")
.get<{ leaf_id: string | null }>("session-1", "main");
expect(row?.leaf_id).toBe(rootId);
const latestLaneMove = await db
.prepare("SELECT lane, leaf_id FROM lane_moves WHERE session_id = ? ORDER BY seq DESC LIMIT 1")
.get<{ lane: string; leaf_id: string | null }>("session-1");
expect(latestLaneMove).toEqual({ lane: "main", leaf_id: rootId });
} finally {
await db.close();
}
@ -285,7 +254,7 @@ END;
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.moveTo(rootId);
await moveSqliteMainLane(session, rootId);
const secondChildId = await session.appendMessage(createAssistantMessage("second child"));
const db = await sqlite.open(databasePath);
@ -318,43 +287,30 @@ END;
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.moveTo(rootId);
await appendSqliteSessionName(session, " Reopened Session ");
await moveSqliteMainLane(session, 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({
expect(await reopened.getName()).toBe("Reopened Session");
expect((await buildSqliteContext(reopened)).messages.map((message) => message.role)).toEqual([
"user",
"assistant",
]);
expect((await buildSqliteContext(reopened)).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 SqliteSessionRepository({ env, sqlite: createNodeSqliteFactory(), databasePath });
const session = await repo.create({ cwd: root, id: "session-1" });
const ids = [
await session.appendMessage(createUserMessage("one")),
await session.appendMessage(createAssistantMessage("two")),
await session.appendMessage(createUserMessage("three")),
];
expect((await session.getEntries({ limit: 2 })).map((entry) => entry.id)).toEqual(ids.slice(0, 2));
expect((await session.getEntries({ afterEntrySeq: 1, limit: 2 })).map((entry) => entry.id)).toEqual(ids.slice(1));
});
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 () => {
return new ThrowingStatement(() => {
throw new Error("insert failed");
});
}
return new ThrowingStatement(async () => ({ changes: 1 }));
return new ThrowingStatement(() => ({ changes: 1 }));
});
const sqlite: SqliteDatabaseFactory = {
open: async () => db,
@ -372,9 +328,9 @@ END;
const root = createTempDir();
const db = new CountingDatabase((sql) => {
if (sql.includes("FROM sessions WHERE id = ?")) {
return new ThrowingStatement(async () => ({ changes: 0 }));
return new ThrowingStatement(() => ({ changes: 0 }));
}
return new ThrowingStatement(async () => ({ changes: 1 }));
return new ThrowingStatement(() => ({ changes: 1 }));
});
const sqlite: SqliteDatabaseFactory = {
open: async () => db,
@ -383,7 +339,7 @@ END;
const repo = new SqliteSessionRepository({ env, sqlite, databasePath: join(root, "sessions.sqlite") });
const metadata: SqliteSessionMetadata = {
id: "missing",
createdAt: new Date().toISOString(),
createdAt: Date.now(),
cwd: root,
path: join(root, "sessions.sqlite"),
};
@ -404,7 +360,7 @@ END;
const session = await repo.create({ cwd: root, id: "session-1" });
for (let i = 0; i < 10; i++) await session.appendMessage(createUserMessage(`message ${i}`));
await session.getEntries();
await getSqliteEntries(session);
expect(counts).toEqual({ opens: 1, closes: 0 });
await repo[Symbol.asyncDispose]();
expect(counts).toEqual({ opens: 1, closes: 1 });
@ -427,7 +383,7 @@ END;
expect(counts).toEqual({ opens: 1, closes: 1 });
});
it("rejects a missing active leaf when opened", async () => {
it("rejects a missing lane leaf when listing lanes and opening", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const env = new NodeExecutionEnv({ cwd: root });
@ -438,14 +394,20 @@ END;
const db = await sqlite.open(databasePath);
try {
await db.prepare("UPDATE sessions SET active_leaf_id = ? WHERE id = ?").run("missing", metadata.id);
await db
.prepare("UPDATE lanes SET leaf_id = ? WHERE session_id = ? AND lane = ?")
.run("missing", metadata.id, "main");
} finally {
await db.close();
}
await expect(session.getLanes()).rejects.toMatchObject({
code: "storage",
message: expect.stringContaining("Lane main points at missing entry missing"),
});
await expect(repo.open(metadata)).rejects.toMatchObject({
code: "invalid_session",
message: "Entry missing not found",
code: "storage",
message: expect.stringContaining("Lane main points at missing entry missing"),
});
});
@ -462,64 +424,46 @@ END;
const db = await sqlite.open(databasePath);
try {
await db
.prepare("UPDATE session_entries SET payload = ? WHERE session_id = ? AND id = ?")
.prepare("UPDATE entries SET payload = ? WHERE session_id = ? AND id = ?")
.run("not json", metadata.id, entryId);
} finally {
await db.close();
}
const reopened = await repo.open(metadata);
await expect(reopened.getEntries()).rejects.toMatchObject({ code: "invalid_entry" });
await expect(getSqliteEntries(reopened)).rejects.toMatchObject({ code: "invalid_entry" });
});
it("does not publish connection state when an append transaction fails", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const env = new NodeExecutionEnv({ cwd: root });
const sqlite = createNodeSqliteFactory();
const repo = new SqliteSessionRepository({ env, sqlite, databasePath });
const session = await repo.create({ cwd: root, id: "session-1" });
const db = await sqlite.open(databasePath);
await applyMigrations(db);
const storage = await SqliteSessionConnection.create(db, databasePath, {
cwd: root,
sessionId: "session-1",
});
await db.exec(`
CREATE TEMP TRIGGER fail_branch_tip_insert
BEFORE INSERT ON branch_tips
BEGIN
SELECT RAISE(ABORT, 'branch insert failed');
END;
`);
const rootEntry = {
type: "message" as const,
id: "root",
parentId: null,
timestamp: new Date().toISOString(),
message: createUserMessage("root"),
};
try {
await expect(storage.appendEntry(rootEntry)).rejects.toMatchObject({ code: "storage" });
} finally {
await db.exec(`
CREATE TRIGGER fail_branch_tip_insert
BEFORE INSERT ON branch_tips
BEGIN
SELECT RAISE(ABORT, 'branch insert failed');
END;
`);
await expect(session.appendMessage(createUserMessage("root"))).rejects.toThrow("branch insert failed");
const lane = await db
.prepare("SELECT leaf_id FROM lanes WHERE session_id = ? AND lane = ?")
.get<{ leaf_id: string | null }>("session-1", "main");
expect(lane?.leaf_id).toBeNull();
expect(await db.prepare("SELECT id FROM entries WHERE session_id = ?").all("session-1")).toEqual([]);
expect(await session.getStats()).toMatchObject({ messageCount: 0 });
await db.exec("DROP TRIGGER fail_branch_tip_insert");
} finally {
await db.close();
}
const sessionRow = await db
.prepare("SELECT active_leaf_id FROM sessions WHERE id = ?")
.get<{ active_leaf_id: string | null }>("session-1");
expect(sessionRow?.active_leaf_id).toBeNull();
expect(await storage.readEntries()).toEqual([]);
await expect(
storage.appendEntry({
type: "leaf",
id: "leaf",
parentId: null,
timestamp: new Date().toISOString(),
targetId: rootEntry.id,
}),
).rejects.toMatchObject({ code: "not_found" });
expect(await storage.readEntries()).toEqual([]);
await storage.appendEntry(rootEntry);
expect(await storage.readEntries()).toEqual([rootEntry]);
await db.close();
const entryId = await session.appendMessage(createUserMessage("root"));
expect((await getSqliteEntries(session)).map((entry) => entry.id)).toEqual([entryId]);
expect(await session.getStats()).toMatchObject({ messageCount: 1 });
});
it("materializes session summary fields transactionally", async () => {
@ -530,8 +474,6 @@ END;
const repo = new SqliteSessionRepository({ 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",
@ -545,62 +487,117 @@ END;
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, {
const assistantId = await session.appendMessage(assistant);
await session.appendRecord({
type: "usage",
id: "assistant-usage",
lane: "main",
cause: "assistant",
runId: "run",
entryId: assistantId,
attempt: 1,
stopReason: "stop",
usage: assistant.usage,
});
const compactionUsage = {
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 compactionId = await appendSqliteCompaction(
session,
"summary",
userId,
200,
undefined,
false,
compactionUsage,
);
await session.appendRecord({
type: "usage",
id: "compaction-usage",
lane: "main",
cause: "compaction",
runId: "run",
entryId: compactionId,
attempt: 1,
stopReason: "stop",
usage: compactionUsage,
});
await session.moveTo(userId, {
const branchUsage = {
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 branchSummaryId = await moveSqliteMainLane(session, 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 },
},
usage: branchUsage,
});
if (!branchSummaryId) throw new Error("Expected branch summary");
await session.appendRecord({
type: "usage",
id: "branch-summary-usage",
lane: "main",
cause: "branch_summary",
runId: "run",
entryId: branchSummaryId,
attempt: 1,
stopReason: "stop",
usage: branchUsage,
});
await appendSqliteSessionName(session, " My Session ");
await appendSqliteLabel(session, userId, "checkpoint");
expect(await session.getStats()).toMatchObject({
messageCount: 2,
cachedTokens: 50,
uncachedTokens: 128,
totalTokens: 211,
costTotal: 0.73,
});
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);
await db
.prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'session_materialized'")
.get(),
).toBeUndefined();
expect(
await db
.prepare("SELECT type, COUNT(*) AS count FROM records WHERE session_id = ? AND type = ? GROUP BY type")
.get<{ type: string; count: number }>("session-1", "usage"),
).toEqual({ type: "usage", count: 3 });
expect(
await db
.prepare(
`SELECT message_count, cached_tokens, uncached_tokens, total_tokens, cost_total
FROM session_stats
WHERE session_id = ?`,
)
.get("session-1"),
).toEqual({
message_count: 2,
cached_tokens: 50,
uncached_tokens: 128,
total_tokens: 211,
cost_total: 0.73,
});
const nameFact = await db
.prepare("SELECT value FROM facts WHERE session_id = ? AND kind = 'name' ORDER BY seq DESC LIMIT 1")
.get<{ value: string }>("session-1");
expect(JSON.parse(nameFact?.value ?? "null")).toBe("My Session");
const labelFact = await db
.prepare("SELECT key, value FROM facts WHERE session_id = ? AND kind = 'label' ORDER BY seq DESC LIMIT 1")
.get<{ key: string; value: string }>("session-1");
expect(labelFact?.key).toBe(userId);
expect(JSON.parse(labelFact?.value ?? "null")).toBe("checkpoint");
} finally {
await db.close();
}

50
packages/agent/test/harness/sqlite-node.test.ts
View file

@ -9,8 +9,8 @@ import {
import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts";
import { JsonlSessionRepository } from "../../src/harness/session/jsonl-repo.ts";
import { createScanningSessionSearch } from "../../src/harness/session/search.ts";
import type { SessionSearch, SessionSearchHit, SessionSearchOptions } from "../../src/harness/types.ts";
import { createTempDir, createUserMessage } from "./session-test-utils.ts";
import type { SessionSearchOptions } from "../../src/harness/types.ts";
import { createTempDir, createUserMessage, getSqliteEntries } from "./session-test-utils.ts";
const ownedRepositories: AsyncDisposable[] = [];
@ -46,6 +46,42 @@ describe("JsonlSessionBackend with scanning search", () => {
});
});
describe("SqliteSessionRepository writer leases", () => {
it("shares one storage queue for repeated opens in the same repository", async () => {
const root = createTempDir();
const env = new NodeExecutionEnv({ cwd: root });
const sqlite = createNodeSqliteFactory();
const databasePath = join(root, "sessions.sqlite");
const { repository: repo } = createSqliteFixture({ env, sqlite, databasePath });
const session = await repo.create({ cwd: root, id: "session" });
const reopened = await repo.open(await session.getMetadata());
const [first, second] = await Promise.all([
session.appendMessage(createUserMessage("first")),
reopened.appendMessage(createUserMessage("second")),
]);
expect((await getSqliteEntries(session)).map((entry) => entry.id)).toEqual([first, second]);
});
it("rejects a second repository while a session lease is active", async () => {
const root = createTempDir();
const env = new NodeExecutionEnv({ cwd: root });
const sqlite = createNodeSqliteFactory();
const databasePath = join(root, "sessions.sqlite");
const { repository: firstRepo } = createSqliteFixture({ env, sqlite, databasePath });
const session = await firstRepo.create({ cwd: root, id: "session" });
const metadata = await session.getMetadata();
const secondRepo = new SqliteSessionRepository({ env, sqlite, databasePath });
ownedRepositories.push(secondRepo);
await expect(secondRepo.open(metadata)).rejects.toMatchObject({ code: "storage" });
await firstRepo.close();
await expect(secondRepo.open(metadata)).resolves.toBeDefined();
});
});
describe("SqliteSessionBackend with explicit SQLite FTS5 search", () => {
it("uses SQLite FTS5 when composed with its search implementation", async () => {
const root = createTempDir();
@ -118,7 +154,7 @@ describe("SqliteSessionBackend with explicit SQLite FTS5 search", () => {
}
await expect(session.appendMessage(createUserMessage("must roll back"))).rejects.toThrow();
await expect(session.getEntries()).resolves.toEqual([]);
await expect(getSqliteEntries(session)).resolves.toEqual([]);
});
it("rolls back canonical deletion when co-located FTS cleanup fails", async () => {
@ -141,7 +177,7 @@ describe("SqliteSessionBackend with explicit SQLite FTS5 search", () => {
await expect(repo.delete(metadata)).rejects.toThrow();
const reopened = await repo.open(metadata);
await expect(reopened.getEntries()).resolves.toHaveLength(1);
await expect(getSqliteEntries(reopened)).resolves.toHaveLength(1);
});
it("initializes canonical storage when searched before the first session is created", async () => {
@ -168,8 +204,10 @@ describe("SqliteSessionRepository with custom search", () => {
it("uses an independently supplied search implementation", async () => {
const root = createTempDir();
const searches: SessionSearchOptions[] = [];
const search: SessionSearch<SqliteSessionMetadata> = {
async search(options): Promise<SessionSearchHit<SqliteSessionMetadata>[]> {
const search: {
search(options: SessionSearchOptions): Promise<{ metadata: SqliteSessionMetadata; entryId: string }[]>;
} = {
async search(options) {
searches.push(options);
return [];
},

2
packages/agent/vitest.config.ts
View file

@ -4,6 +4,7 @@ import { defineConfig } from "vitest/config";
const aiSrcIndex = fileURLToPath(new URL("../ai/src/index.ts", import.meta.url));
const aiSrcCompat = fileURLToPath(new URL("../ai/src/compat.ts", import.meta.url));
const agentSrcIndex = fileURLToPath(new URL("./src/index.ts", import.meta.url));
const agentSrcExperimental = fileURLToPath(new URL("./src/experimental.ts", import.meta.url));
export default defineConfig({
test: {
@ -16,6 +17,7 @@ export default defineConfig({
resolve: {
alias: [
{ find: /^@earendil-works\/pi-agent-core$/, replacement: agentSrcIndex },
{ find: /^@earendil-works\/pi-agent-core\/experimental$/, replacement: agentSrcExperimental },
{ find: /^@earendil-works\/pi-ai$/, replacement: aiSrcIndex },
{ find: /^@earendil-works\/pi-ai\/compat$/, replacement: aiSrcCompat },
],

2
packages/agent/vitest.harness.config.ts
View file

@ -4,6 +4,7 @@ import { defineConfig } from "vitest/config";
const aiSrcIndex = fileURLToPath(new URL("../ai/src/index.ts", import.meta.url));
const aiSrcCompat = fileURLToPath(new URL("../ai/src/compat.ts", import.meta.url));
const agentSrcIndex = fileURLToPath(new URL("../agent/src/index.ts", import.meta.url));
const agentSrcExperimental = fileURLToPath(new URL("../agent/src/experimental.ts", import.meta.url));
export default defineConfig({
test: {
@ -23,6 +24,7 @@ export default defineConfig({
resolve: {
alias: [
{ find: /^@earendil-works\/pi-agent-core$/, replacement: agentSrcIndex },
{ find: /^@earendil-works\/pi-agent-core\/experimental$/, replacement: agentSrcExperimental },
{ find: /^@earendil-works\/pi-ai$/, replacement: aiSrcIndex },
{ find: /^@earendil-works\/pi-ai\/compat$/, replacement: aiSrcCompat },
],

23
packages/storage/sqlite-node/src/index.ts
View file

@ -8,6 +8,10 @@ function isNamedParameters(value: unknown): value is Record<string, SQLInputValu
return true;
}
function isAsyncResult(value: unknown): boolean {
return value !== null && (typeof value === "object" || typeof value === "function") && "then" in value;
}
class NodeSqliteStatement implements SqliteStatement {
private readonly statement: ReturnType<DatabaseSync["prepare"]>;
@ -15,7 +19,7 @@ class NodeSqliteStatement implements SqliteStatement {
this.statement = statement;
}
async run(...params: unknown[]): Promise<SqliteRunResult> {
run(...params: unknown[]): SqliteRunResult {
const [first, ...rest] = params;
const result = isNamedParameters(first)
? this.statement.run(first, ...(rest as SQLInputValue[]))
@ -26,7 +30,7 @@ class NodeSqliteStatement implements SqliteStatement {
};
}
async get<TRow extends object>(...params: unknown[]): Promise<TRow | undefined> {
get<TRow extends object>(...params: unknown[]): TRow | undefined {
const [first, ...rest] = params;
return (
isNamedParameters(first)
@ -35,7 +39,7 @@ class NodeSqliteStatement implements SqliteStatement {
) as TRow | undefined;
}
async all<TRow extends object>(...params: unknown[]): Promise<TRow[]> {
all<TRow extends object>(...params: unknown[]): TRow[] {
const [first, ...rest] = params;
return (
isNamedParameters(first)
@ -52,7 +56,7 @@ class NodeSqliteDatabase implements SqliteDatabase {
this.db = db;
}
async exec(sql: string): Promise<void> {
exec(sql: string): void {
this.db.exec(sql);
}
@ -60,10 +64,13 @@ class NodeSqliteDatabase implements SqliteDatabase {
return new NodeSqliteStatement(this.db.prepare(sql));
}
async transaction<T>(fn: () => Promise<T>): Promise<T> {
this.db.exec("BEGIN");
transaction<T>(fn: () => T): T {
this.db.exec("BEGIN IMMEDIATE");
try {
const result = await fn();
const result = fn();
if (isAsyncResult(result)) {
throw new TypeError("SQLite transaction callbacks must be synchronous");
}
this.db.exec("COMMIT");
return result;
} catch (error) {
@ -76,7 +83,7 @@ class NodeSqliteDatabase implements SqliteDatabase {
}
}
async close(): Promise<void> {
close(): void {
this.db.close();
}
}

105
packages/storage/sqlite-node/src/sqlite/branch-cache.ts Normal file
View file

@ -0,0 +1,105 @@
import { SessionError } from "@earendil-works/pi-agent-core/experimental";
import { uuidv7 } from "@earendil-works/pi-ai";
import {
copyBranchEntriesThroughSeq,
deleteBranchEntries,
insertBranchEntriesForPath,
insertBranchEntry,
readBranchContainingEntry,
} from "./storage/branch-entries.ts";
import { deleteBranchTips, insertBranchTip, readBranchTipBranchId, updateBranchTip } from "./storage/branch-tips.ts";
import type { SqliteDatabase } from "./types.ts";
export function deleteBranchCache(db: SqliteDatabase, sessionId: string) {
deleteBranchTips(db, sessionId);
deleteBranchEntries(db, sessionId);
}
export function rebuildBranchCache(db: SqliteDatabase, sessionId: string) {
const tips = db
.prepare(
`SELECT leaf.id
FROM entries AS leaf
WHERE leaf.session_id = ?
AND NOT EXISTS (
SELECT 1 FROM entries AS child WHERE child.session_id = leaf.session_id AND child.parent_id = leaf.id
)
ORDER BY leaf.seq`,
)
.all<{ id: string }>(sessionId);
deleteBranchCache(db, sessionId);
for (const tip of tips) buildCachedBranch(db, sessionId, tip.id);
}
export function buildCachedBranch(db: SqliteDatabase, sessionId: string, leafId: string) {
db.exec("SAVEPOINT build_branch_cache");
try {
const branchId = uuidv7();
insertBranchEntriesForPath(db, sessionId, branchId, leafId);
insertBranchTip(db, sessionId, leafId, branchId);
db.exec("RELEASE SAVEPOINT build_branch_cache");
} catch (error) {
try {
db.exec("ROLLBACK TO SAVEPOINT build_branch_cache");
db.exec("RELEASE SAVEPOINT build_branch_cache");
} catch {
// Preserve the original build failure.
}
if (error instanceof SessionError) throw error;
throw new SessionError(
"storage",
`Failed to build SQLite branch cache at entry ${leafId}`,
error instanceof Error ? error : undefined,
);
}
}
function extendBranch(
db: SqliteDatabase,
sessionId: string,
branchId: string,
parentId: string,
entryId: string,
entrySeq: number,
entryType: string,
customType: string | null,
) {
insertBranchEntry(db, sessionId, branchId, entryId, entrySeq, entryType, customType);
if (!updateBranchTip(db, sessionId, branchId, parentId, entryId)) {
throw new SessionError("invalid_entry", `Branch tip ${parentId} changed during append`);
}
}
export function appendEntryToBranchCache(
db: SqliteDatabase,
sessionId: string,
entryId: string,
entrySeq: number,
entryType: string,
customType: string | null,
parentId: string | null,
) {
if (parentId === null) {
const branchId = uuidv7();
insertBranchEntry(db, sessionId, branchId, entryId, entrySeq, entryType, customType);
insertBranchTip(db, sessionId, entryId, branchId);
return;
}
const tipBranchId = readBranchTipBranchId(db, sessionId, parentId);
if (tipBranchId !== undefined) {
extendBranch(db, sessionId, tipBranchId, parentId, entryId, entrySeq, entryType, customType);
return;
}
const source = readBranchContainingEntry(db, sessionId, parentId);
if (!source) {
throw new SessionError("invalid_entry", `Branch cache has no branch containing parent entry ${parentId}`);
}
const branchId = uuidv7();
copyBranchEntriesThroughSeq(db, sessionId, branchId, source.branchId, source.entrySeq);
insertBranchEntry(db, sessionId, branchId, entryId, entrySeq, entryType, customType);
insertBranchTip(db, sessionId, entryId, branchId);
}

11
packages/storage/sqlite-node/src/sqlite/index.ts
View file

@ -1,7 +1,16 @@
export * from "./migrations.ts";
export {
type SqliteSessionCreateOptions,
type SqliteSessionListOptions,
type SqliteSessionMetadata,
SqliteSessionRepository,
type SqliteSessionRepositoryOptions,
type SqliteWriterLeaseOptions,
} from "./repo.ts";
export * from "./search-backend.ts";
export * from "./types.ts";
export type {
SqliteDatabase,
SqliteDatabaseFactory,
SqliteRunResult,
SqliteStatement,
} from "./types.ts";

24
packages/storage/sqlite-node/src/sqlite/migrations.ts
View file

@ -19,16 +19,11 @@ export async function loadMigrations(): Promise<SqliteMigration[]> {
order: 1,
sql: await loadMigrationSql("./migrations/001_initial.sql"),
},
{
id: "002_branch_tips.sql",
order: 2,
sql: await loadMigrationSql("./migrations/002_branch_tips.sql"),
},
];
}
async function ensureMigrationsTable(db: SqliteDatabase): Promise<void> {
await db.exec(`
function ensureMigrationsTable(db: SqliteDatabase): void {
db.exec(`
CREATE TABLE IF NOT EXISTS migrations (
id TEXT PRIMARY KEY,
applied_at TEXT NOT NULL
@ -37,18 +32,19 @@ CREATE TABLE IF NOT EXISTS migrations (
}
export async function applyMigrations(db: SqliteDatabase): Promise<void> {
await ensureMigrationsTable(db);
ensureMigrationsTable(db);
const migrations = await loadMigrations();
const appliedRows = await db.prepare("SELECT id FROM migrations ORDER BY applied_at, id").all<{ id: string }>();
const appliedRows = db.prepare("SELECT id FROM migrations ORDER BY applied_at, id").all<{ id: string }>();
const applied = new Set(appliedRows.map((row) => row.id));
for (const migration of migrations) {
if (applied.has(migration.id)) continue;
await db.transaction(async () => {
await db.exec(migration.sql);
await db
.prepare("INSERT INTO migrations (id, applied_at) VALUES (?, ?)")
.run(migration.id, new Date().toISOString());
db.transaction(() => {
db.exec(migration.sql);
db.prepare("INSERT INTO migrations (id, applied_at) VALUES (?, ?)").run(
migration.id,
new Date().toISOString(),
);
});
applied.add(migration.id);
}

105
packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql
View file

@ -3,57 +3,122 @@ CREATE TABLE IF NOT EXISTS sessions (
created_at TEXT NOT NULL,
cwd TEXT NOT NULL,
parent_session_id TEXT NULL,
metadata TEXT NULL,
active_leaf_id TEXT NULL
metadata TEXT NULL
) WITHOUT ROWID;
CREATE INDEX IF NOT EXISTS idx_sessions_created_at ON sessions(created_at DESC);
CREATE INDEX IF NOT EXISTS idx_sessions_cwd ON sessions(cwd);
CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_session_id);
CREATE TABLE IF NOT EXISTS session_entries (
CREATE TABLE IF NOT EXISTS entries (
session_id TEXT NOT NULL,
seq INTEGER NOT NULL,
id TEXT NOT NULL,
entry_seq INTEGER NOT NULL,
parent_id TEXT NULL,
type TEXT NOT NULL,
timestamp TEXT NOT NULL,
payload TEXT NOT NULL,
PRIMARY KEY (session_id, id)
PRIMARY KEY (session_id, id),
UNIQUE (session_id, seq)
);
CREATE UNIQUE INDEX IF NOT EXISTS idx_session_entries_session_seq ON session_entries(session_id, entry_seq);
CREATE INDEX IF NOT EXISTS idx_session_entries_session_parent ON session_entries(session_id, parent_id);
CREATE INDEX IF NOT EXISTS idx_session_entries_session_type ON session_entries(session_id, type);
CREATE INDEX IF NOT EXISTS idx_entries_session_seq ON entries(session_id, seq);
CREATE INDEX IF NOT EXISTS idx_entries_session_parent ON entries(session_id, parent_id);
CREATE INDEX IF NOT EXISTS idx_entries_session_type_seq ON entries(session_id, type, seq);
CREATE TABLE IF NOT EXISTS session_sequences (
session_id TEXT PRIMARY KEY,
next_seq INTEGER NOT NULL
) WITHOUT ROWID;
CREATE TABLE IF NOT EXISTS session_stats (
session_id TEXT PRIMARY KEY,
message_count INTEGER NOT NULL,
cached_tokens REAL NOT NULL,
uncached_tokens REAL NOT NULL,
total_tokens REAL NOT NULL,
cost_total REAL NOT NULL
) WITHOUT ROWID;
-- Derived branch cache. Parent links in entries remain canonical; this cache
-- exists only to make branch scans cheap.
CREATE TABLE IF NOT EXISTS branch_entries (
session_id TEXT NOT NULL,
branch_id TEXT NOT NULL,
entry_id TEXT NOT NULL,
entry_seq INTEGER NOT NULL,
entry_type TEXT NULL,
custom_type TEXT NULL,
PRIMARY KEY (session_id, branch_id, entry_id)
) WITHOUT ROWID;
CREATE INDEX IF NOT EXISTS idx_branch_entries_session_branch ON branch_entries(session_id, branch_id);
CREATE INDEX IF NOT EXISTS idx_branch_entries_session_branch_seq ON branch_entries(session_id, branch_id, entry_seq);
CREATE INDEX IF NOT EXISTS idx_branch_entries_session_entry ON branch_entries(session_id, entry_id);
CREATE INDEX IF NOT EXISTS idx_branch_entries_session_branch_type_seq ON branch_entries(session_id, branch_id, entry_type, entry_seq);
CREATE INDEX IF NOT EXISTS idx_branch_entries_session_branch_custom_seq ON branch_entries(session_id, branch_id, custom_type, entry_seq);
CREATE TABLE IF NOT EXISTS session_materialized (
session_id TEXT PRIMARY KEY,
payload TEXT NOT NULL
) WITHOUT ROWID;
CREATE TABLE IF NOT EXISTS entry_materialized (
CREATE TABLE IF NOT EXISTS lanes (
session_id TEXT NOT NULL,
entry_seq INTEGER NOT NULL,
type TEXT NOT NULL,
payload TEXT NOT NULL,
PRIMARY KEY (session_id, entry_seq, type)
lane TEXT NOT NULL,
leaf_id TEXT NULL,
PRIMARY KEY (session_id, lane)
) WITHOUT ROWID;
CREATE INDEX IF NOT EXISTS idx_entry_materialized_session_type_seq ON entry_materialized(session_id, type, entry_seq);
CREATE INDEX IF NOT EXISTS idx_lanes_session_leaf ON lanes(session_id, leaf_id);
CREATE TABLE IF NOT EXISTS records (
session_id TEXT NOT NULL,
seq INTEGER NOT NULL,
id TEXT NOT NULL,
lane TEXT NOT NULL,
run_id TEXT NULL,
type TEXT NOT NULL,
op_kind TEXT NULL,
timestamp TEXT NOT NULL,
payload TEXT NOT NULL,
PRIMARY KEY (session_id, id),
UNIQUE (session_id, seq)
) WITHOUT ROWID;
CREATE INDEX IF NOT EXISTS idx_records_session_seq ON records(session_id, seq);
CREATE INDEX IF NOT EXISTS idx_records_session_lane_type_seq ON records(session_id, lane, type, seq);
CREATE INDEX IF NOT EXISTS idx_records_session_lane_type_op_kind_seq ON records(session_id, lane, type, op_kind, seq);
CREATE INDEX IF NOT EXISTS idx_records_session_run_id_seq ON records(session_id, run_id, seq);
CREATE TABLE IF NOT EXISTS lane_moves (
session_id TEXT NOT NULL,
seq INTEGER NOT NULL,
lane TEXT NOT NULL,
leaf_id TEXT NULL,
PRIMARY KEY (session_id, seq)
) WITHOUT ROWID;
CREATE INDEX IF NOT EXISTS idx_lane_moves_session_lane_seq ON lane_moves(session_id, lane, seq);
CREATE TABLE IF NOT EXISTS facts (
session_id TEXT NOT NULL,
seq INTEGER NOT NULL,
kind TEXT NOT NULL,
key TEXT NULL,
value TEXT NULL,
PRIMARY KEY (session_id, seq)
) WITHOUT ROWID;
CREATE INDEX IF NOT EXISTS idx_facts_session_kind_key_seq ON facts(session_id, kind, key, seq);
CREATE TABLE IF NOT EXISTS branch_tips (
session_id TEXT NOT NULL,
tip_id TEXT NOT NULL,
branch_id TEXT NOT NULL,
PRIMARY KEY (session_id, tip_id),
UNIQUE (session_id, branch_id)
) WITHOUT ROWID;
-- Per-session writer claim. The fence prevents an expired owner from writing
-- after a new owner takes over the session.
CREATE TABLE IF NOT EXISTS leases (
session_id TEXT PRIMARY KEY,
owner_id TEXT NOT NULL,
fence INTEGER NOT NULL,
expires_at_ms INTEGER NOT NULL
) WITHOUT ROWID;

12
packages/storage/sqlite-node/src/sqlite/migrations/002_branch_tips.sql
View file

@ -1,12 +0,0 @@
CREATE TABLE IF NOT EXISTS branch_tips (
session_id TEXT NOT NULL,
tip_id TEXT NOT NULL,
branch_id TEXT NOT NULL,
PRIMARY KEY (session_id, tip_id),
UNIQUE (session_id, branch_id)
) WITHOUT ROWID;
DELETE FROM branch_tips;
DELETE FROM branch_entries;
DROP INDEX IF EXISTS idx_branch_entries_session_branch;

1116
packages/storage/sqlite-node/src/sqlite/repo.ts
View file

File diff suppressed because it is too large Load diff

42
packages/storage/sqlite-node/src/sqlite/search-backend.ts
View file

@ -17,10 +17,10 @@ function getParentPath(path: string): string {
return normalized.slice(0, lastSlash);
}
async function configureSqliteDatabase(db: SqliteDatabase): Promise<void> {
await db.exec("PRAGMA journal_mode=WAL");
await db.exec("PRAGMA synchronous=FULL");
await db.exec("PRAGMA busy_timeout=5000");
function configureSqliteDatabase(db: SqliteDatabase): void {
db.exec("PRAGMA journal_mode=WAL");
db.exec("PRAGMA synchronous=FULL");
db.exec("PRAGMA busy_timeout=5000");
}
export interface SqliteSessionSearchOptions {
@ -29,33 +29,33 @@ export interface SqliteSessionSearchOptions {
databasePath: string;
}
async function tableExists(db: SqliteDatabase, name: string): Promise<boolean> {
return !!(await db
function tableExists(db: SqliteDatabase, name: string): boolean {
return !!db
.prepare("SELECT 1 AS found FROM sqlite_master WHERE type = 'table' AND name = ? LIMIT 1")
.get<{ found: number }>(name));
.get<{ found: number }>(name);
}
async function ensureSearchSchema(db: SqliteDatabase): Promise<void> {
const ftsExists = await tableExists(db, "session_search_fts");
await db.exec(`
function ensureSearchSchema(db: SqliteDatabase): void {
const ftsExists = tableExists(db, "session_search_fts");
db.exec(`
CREATE VIRTUAL TABLE IF NOT EXISTS session_search_fts USING fts5(
payload,
content = 'session_entries',
content = 'entries',
content_rowid = 'rowid',
tokenize = 'trigram remove_diacritics 1'
);
CREATE TRIGGER IF NOT EXISTS session_search_fts_ai AFTER INSERT ON session_entries BEGIN
CREATE TRIGGER IF NOT EXISTS session_search_fts_ai AFTER INSERT ON entries BEGIN
INSERT INTO session_search_fts(rowid, payload) VALUES (new.rowid, new.payload);
END;
CREATE TRIGGER IF NOT EXISTS session_search_fts_ad AFTER DELETE ON session_entries BEGIN
CREATE TRIGGER IF NOT EXISTS session_search_fts_ad AFTER DELETE ON entries BEGIN
INSERT INTO session_search_fts(session_search_fts, rowid, payload) VALUES('delete', old.rowid, old.payload);
END;
CREATE TRIGGER IF NOT EXISTS session_search_fts_au AFTER UPDATE OF payload ON session_entries BEGIN
CREATE TRIGGER IF NOT EXISTS session_search_fts_au AFTER UPDATE OF payload ON entries BEGIN
INSERT INTO session_search_fts(session_search_fts, rowid, payload) VALUES('delete', old.rowid, old.payload);
INSERT INTO session_search_fts(rowid, payload) VALUES (new.rowid, new.payload);
END;
`);
if (!ftsExists) await db.exec("INSERT INTO session_search_fts(session_search_fts) VALUES('rebuild')");
if (!ftsExists) db.exec("INSERT INTO session_search_fts(session_search_fts) VALUES('rebuild')");
}
/** SQLite FTS search over a co-located canonical session database. */
@ -86,12 +86,12 @@ class SqliteSessionSearch implements SessionSearch<SqliteSessionMetadata> {
);
const db = await this.options.sqlite.open(path);
try {
await configureSqliteDatabase(db);
configureSqliteDatabase(db);
await applyMigrations(db);
await ensureSearchSchema(db);
ensureSearchSchema(db);
return db;
} catch (error) {
await db.close();
db.close();
throw error;
}
}
@ -102,9 +102,9 @@ class SqliteSessionSearch implements SessionSearch<SqliteSessionMetadata> {
const db = await this.openDatabase();
try {
const query = `"${text.replaceAll('"', '""')}"`;
const rows = await db
const rows = db
.prepare(
"SELECT s.id, s.created_at, s.metadata, s.cwd, s.parent_session_id, s.active_leaf_id, se.id AS entry_id, se.timestamp, bm25(session_search_fts) AS score FROM session_search_fts JOIN session_entries se ON se.rowid = session_search_fts.rowid JOIN sessions s ON s.id = se.session_id WHERE session_search_fts MATCH ? AND (? IS NULL OR s.cwd = ?) ORDER BY score",
"SELECT s.id, s.created_at, s.metadata, s.cwd, s.parent_session_id, se.id AS entry_id, se.timestamp, bm25(session_search_fts) AS score FROM session_search_fts JOIN entries se ON se.rowid = session_search_fts.rowid JOIN sessions s ON s.id = se.session_id WHERE session_search_fts MATCH ? AND (? IS NULL OR s.cwd = ?) ORDER BY score",
)
.all<SessionRow & { entry_id: string; timestamp: string; score: number }>(
query,
@ -119,7 +119,7 @@ class SqliteSessionSearch implements SessionSearch<SqliteSessionMetadata> {
score: row.score,
}));
} finally {
await db.close();
db.close();
}
}
}

326
packages/storage/sqlite-node/src/sqlite/storage/branch-cache.ts
View file

@ -1,326 +0,0 @@
import { uuidv7 } from "@earendil-works/pi-ai";
import type { SqliteDatabase } from "../types.ts";
import type { SessionEntryRow } from "./session-entries.ts";
import { invalidSession } from "./shared.ts";
/** Derived root-to-tip paths. Canonical parent links remain authoritative. */
export interface CachedBranch {
branchId: string;
leafSeq: number;
}
export interface CachedBranchQuery {
stopAtType?: SessionEntryRow["type"];
stopAtId?: string;
order?: "newestFirst" | "oldestFirst";
}
export async function readCachedBranch(
db: SqliteDatabase,
sessionId: string,
leafId: string,
): Promise<CachedBranch | undefined> {
const membership = await db
.prepare(
"SELECT branch_id, entry_seq FROM branch_entries WHERE session_id = ? AND entry_id = ? ORDER BY branch_id LIMIT 1",
)
.get<{ branch_id: string; entry_seq: number }>(sessionId, leafId);
if (!membership) return undefined;
return { branchId: membership.branch_id, leafSeq: membership.entry_seq };
}
export async function isCachedBranchValid(
db: SqliteDatabase,
sessionId: string,
branch: CachedBranch,
leafId: string,
startSeq = 0,
): Promise<boolean> {
const result = await db
.prepare(
`WITH path AS (
SELECT
b.entry_id,
b.entry_seq,
e.id AS stored_entry_id,
e.parent_id,
LAG(b.entry_id) OVER (ORDER BY b.entry_seq) AS previous_entry_id
FROM branch_entries AS b
LEFT JOIN session_entries AS 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 BETWEEN ? AND ?
)
SELECT
COUNT(*) AS row_count,
COALESCE(SUM(
stored_entry_id IS NULL OR
(? = 0 AND previous_entry_id IS NULL AND parent_id IS NOT NULL) OR
(previous_entry_id IS NOT NULL AND parent_id IS NOT previous_entry_id)
), 0) AS invalid_count,
COALESCE(MAX(entry_seq = ? AND entry_id = ?), 0) AS contains_leaf
FROM path`,
)
.get<{ row_count: number; invalid_count: number; contains_leaf: number }>(
sessionId,
branch.branchId,
startSeq,
branch.leafSeq,
startSeq,
branch.leafSeq,
leafId,
);
return result?.row_count !== 0 && result?.invalid_count === 0 && result.contains_leaf === 1;
}
export async function readNewestCachedStopSeq(
db: SqliteDatabase,
sessionId: string,
branch: CachedBranch,
stopAtType: SessionEntryRow["type"] | undefined,
stopAtId: string | undefined,
): Promise<number | undefined> {
const predicates: string[] = [];
const params: unknown[] = [sessionId, branch.branchId, branch.leafSeq];
if (stopAtType !== undefined) {
predicates.push("e.type = ?");
params.push(stopAtType);
}
if (stopAtId !== undefined) {
predicates.push("b.entry_id = ?");
params.push(stopAtId);
}
if (predicates.length === 0) return undefined;
const row = await db
.prepare(
`SELECT MAX(b.entry_seq) AS entry_seq
FROM branch_entries AS b
JOIN session_entries AS 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 (${predicates.join(" OR ")})`,
)
.get<{ entry_seq: number | null }>(...params);
return row?.entry_seq ?? undefined;
}
export async function readCachedBranchRows(
db: SqliteDatabase,
sessionId: string,
branch: CachedBranch,
startSeq: number,
): Promise<SessionEntryRow[]> {
return db
.prepare(
`SELECT e.session_id, e.id, e.entry_seq, e.parent_id, e.type, e.timestamp, e.payload
FROM branch_entries AS b
JOIN session_entries AS 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 BETWEEN ? AND ?
ORDER BY b.entry_seq`,
)
.all<SessionEntryRow>(sessionId, branch.branchId, startSeq, branch.leafSeq);
}
export async function queryCachedBranchRows(
db: SqliteDatabase,
sessionId: string,
branch: CachedBranch,
query: CachedBranchQuery,
): Promise<SessionEntryRow[]> {
const oldestFirst = query.order === "oldestFirst";
const boundaryParams: unknown[] = [sessionId, branch.branchId, branch.leafSeq];
const stopPredicates: string[] = [];
if (query.stopAtType !== undefined) {
stopPredicates.push("stop_entry.type = ?");
boundaryParams.push(query.stopAtType);
}
if (query.stopAtId !== undefined) {
stopPredicates.push("stop.entry_id = ?");
boundaryParams.push(query.stopAtId);
}
const boundary = stopPredicates.length
? `WITH boundary AS (
SELECT ${oldestFirst ? "MIN" : "MAX"}(stop.entry_seq) AS entry_seq
FROM branch_entries AS stop
JOIN session_entries AS stop_entry
ON stop_entry.session_id = stop.session_id AND stop_entry.id = stop.entry_id
WHERE stop.session_id = ? AND stop.branch_id = ? AND stop.entry_seq <= ?
AND (${stopPredicates.join(" OR ")})
)`
: "";
const range = stopPredicates.length
? `AND b.entry_seq ${oldestFirst ? "<=" : ">="} COALESCE(
(SELECT entry_seq FROM boundary), ${oldestFirst ? branch.leafSeq : 0}
)`
: "";
const sql = `${boundary}
SELECT e.session_id, e.id, e.entry_seq, e.parent_id, e.type, e.timestamp, e.payload
FROM branch_entries AS b
JOIN session_entries AS 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 <= ?
${range}
ORDER BY b.entry_seq ${oldestFirst ? "ASC" : "DESC"}`;
const params = [...(stopPredicates.length === 0 ? [] : boundaryParams), sessionId, branch.branchId, branch.leafSeq];
return db.prepare(sql).all<SessionEntryRow>(...params);
}
export async function readCachedEntryRowsByType(
db: SqliteDatabase,
sessionId: string,
branch: CachedBranch,
type: SessionEntryRow["type"],
): Promise<SessionEntryRow[]> {
// Drive the join from the usually sparse entry type. Ordering from branch_entries
// makes SQLite scan the complete cached path before filtering by type.
return db
.prepare(
`SELECT e.session_id, e.id, e.entry_seq, e.parent_id, e.type, e.timestamp, e.payload
FROM session_entries AS e INDEXED BY idx_session_entries_session_type
CROSS JOIN branch_entries AS b
WHERE e.session_id = ? AND e.type = ?
AND b.session_id = e.session_id AND b.entry_id = e.id
AND b.branch_id = ? AND b.entry_seq <= ?
ORDER BY e.entry_seq DESC`,
)
.all<SessionEntryRow>(sessionId, type, branch.branchId, branch.leafSeq);
}
export async function readCachedEntrySeq(
db: SqliteDatabase,
sessionId: string,
branchId: string,
entryId: string,
): Promise<number | undefined> {
const row = await db
.prepare("SELECT entry_seq FROM branch_entries WHERE session_id = ? AND branch_id = ? AND entry_id = ?")
.get<{ entry_seq: number }>(sessionId, branchId, entryId);
return row?.entry_seq;
}
export async function rebuildCachedBranch(
db: SqliteDatabase,
sessionId: string,
leafId: string,
branchIdToReplace?: string,
): Promise<void> {
await db.exec("SAVEPOINT rebuild_branch_cache");
try {
const tip = await db
.prepare("SELECT branch_id FROM branch_tips WHERE session_id = ? AND tip_id = ?")
.get<{ branch_id: string }>(sessionId, leafId);
const branchIds = new Set([branchIdToReplace, tip?.branch_id].filter((id): id is string => id !== undefined));
for (const branchId of branchIds) {
await db.prepare("DELETE FROM branch_tips WHERE session_id = ? AND branch_id = ?").run(sessionId, branchId);
await db.prepare("DELETE FROM branch_entries WHERE session_id = ? AND branch_id = ?").run(sessionId, branchId);
}
const branchId = uuidv7();
await db
.prepare(
`WITH RECURSIVE path(id, entry_seq, parent_id) AS (
SELECT id, entry_seq, parent_id
FROM session_entries
WHERE session_id = ? AND id = ?
UNION ALL
SELECT parent.id, parent.entry_seq, parent.parent_id
FROM session_entries AS parent
JOIN path AS child ON child.parent_id = parent.id
WHERE parent.session_id = ?
)
INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq)
SELECT ?, ?, id, entry_seq FROM path`,
)
.run(sessionId, leafId, sessionId, sessionId, branchId);
await db
.prepare("INSERT INTO branch_tips (session_id, tip_id, branch_id) VALUES (?, ?, ?)")
.run(sessionId, leafId, branchId);
await db.exec("RELEASE SAVEPOINT rebuild_branch_cache");
} catch (error) {
try {
await db.exec("ROLLBACK TO SAVEPOINT rebuild_branch_cache");
await db.exec("RELEASE SAVEPOINT rebuild_branch_cache");
} catch {
// Preserve the original repair failure.
}
throw error;
}
}
async function extendBranch(
db: SqliteDatabase,
sessionId: string,
branchId: string,
parentId: string,
entryId: string,
entrySeq: number,
): Promise<void> {
await db
.prepare("INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) VALUES (?, ?, ?, ?)")
.run(sessionId, branchId, entryId, entrySeq);
const result = await db
.prepare("UPDATE branch_tips SET tip_id = ? WHERE session_id = ? AND branch_id = ? AND tip_id = ?")
.run(entryId, sessionId, branchId, parentId);
if (result.changes !== 1) throw invalidSession(`branch tip ${parentId} changed during append`);
}
export async function appendEntryToBranchCache(
db: SqliteDatabase,
sessionId: string,
entryId: string,
entrySeq: number,
parentId: string | null,
repairParent: (parentId: string) => Promise<void>,
): Promise<void> {
if (parentId === null) {
const branchId = uuidv7();
await db
.prepare("INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) VALUES (?, ?, ?, ?)")
.run(sessionId, branchId, entryId, entrySeq);
await db
.prepare("INSERT INTO branch_tips (session_id, tip_id, branch_id) VALUES (?, ?, ?)")
.run(sessionId, entryId, branchId);
return;
}
let tip = await db
.prepare("SELECT branch_id FROM branch_tips WHERE session_id = ? AND tip_id = ?")
.get<{ branch_id: string }>(sessionId, parentId);
if (tip) {
await extendBranch(db, sessionId, tip.branch_id, parentId, entryId, entrySeq);
return;
}
const source = await db
.prepare(
`SELECT b.branch_id, b.entry_seq
FROM branch_entries AS b
WHERE b.session_id = ? AND b.entry_id = ?
ORDER BY b.branch_id
LIMIT 1`,
)
.get<{ branch_id: string; entry_seq: number }>(sessionId, parentId);
if (!source) {
await repairParent(parentId);
tip = await db
.prepare("SELECT branch_id FROM branch_tips WHERE session_id = ? AND tip_id = ?")
.get<{ branch_id: string }>(sessionId, parentId);
if (!tip) throw invalidSession(`branch cache repair did not create tip ${parentId}`);
await extendBranch(db, sessionId, tip.branch_id, parentId, entryId, entrySeq);
return;
}
const branchId = uuidv7();
await db
.prepare(
`INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq)
SELECT session_id, ?, entry_id, entry_seq
FROM branch_entries
WHERE session_id = ? AND branch_id = ? AND entry_seq <= ?`,
)
.run(branchId, sessionId, source.branch_id, source.entry_seq);
await db
.prepare("INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) VALUES (?, ?, ?, ?)")
.run(sessionId, branchId, entryId, entrySeq);
await db
.prepare("INSERT INTO branch_tips (session_id, tip_id, branch_id) VALUES (?, ?, ?)")
.run(sessionId, entryId, branchId);
}

146
packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts Normal file
View file

@ -0,0 +1,146 @@
import type { Entry } from "@earendil-works/pi-agent-core/experimental";
import type { SqliteDatabase } from "../types.ts";
/** Derived root-to-tip branch cache membership. Canonical parent links remain in entries. */
export interface CachedBranch {
branchId: string;
leafSeq: number;
}
export interface CachedBranchEntryRow {
session_id: string;
id: string;
entry_seq: number;
parent_id: string | null;
type: Entry["type"];
timestamp: string;
payload: string;
}
export interface CachedBranchQuery {
stopAtType?: Entry["type"];
stopAtId?: string;
order?: "newestFirst" | "oldestFirst";
}
export function readCachedBranch(db: SqliteDatabase, sessionId: string, leafId: string) {
const membership = db
.prepare(
"SELECT branch_id, entry_seq FROM branch_entries WHERE session_id = ? AND entry_id = ? ORDER BY branch_id LIMIT 1",
)
.get<{ branch_id: string; entry_seq: number }>(sessionId, leafId);
if (!membership) return undefined;
return { branchId: membership.branch_id, leafSeq: membership.entry_seq };
}
export function queryCachedBranchRows(
db: SqliteDatabase,
sessionId: string,
branch: CachedBranch,
query: CachedBranchQuery,
) {
const oldestFirst = query.order === "oldestFirst";
const boundaryParams: unknown[] = [sessionId, branch.branchId, branch.leafSeq];
const stopPredicates: string[] = [];
if (query.stopAtType !== undefined) {
stopPredicates.push("stop_entry.type = ?");
boundaryParams.push(query.stopAtType);
}
if (query.stopAtId !== undefined) {
stopPredicates.push("stop.entry_id = ?");
boundaryParams.push(query.stopAtId);
}
const boundary = stopPredicates.length
? `WITH boundary AS (
SELECT ${oldestFirst ? "MIN" : "MAX"}(stop.entry_seq) AS entry_seq
FROM branch_entries AS stop
JOIN entries AS stop_entry
ON stop_entry.session_id = stop.session_id AND stop_entry.id = stop.entry_id
WHERE stop.session_id = ? AND stop.branch_id = ? AND stop.entry_seq <= ?
AND (${stopPredicates.join(" OR ")})
)`
: "";
const range = stopPredicates.length
? `AND b.entry_seq ${oldestFirst ? "<=" : ">="} COALESCE(
(SELECT entry_seq FROM boundary), ${oldestFirst ? branch.leafSeq : 0}
)`
: "";
const sql = `${boundary}
SELECT e.session_id, e.id, e.seq AS entry_seq, e.parent_id, e.type, e.timestamp, e.payload
FROM branch_entries AS b
JOIN entries AS 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 <= ?
${range}
ORDER BY b.entry_seq ${oldestFirst ? "ASC" : "DESC"}`;
const params = [...(stopPredicates.length === 0 ? [] : boundaryParams), sessionId, branch.branchId, branch.leafSeq];
return db.prepare(sql).all<CachedBranchEntryRow>(...params);
}
export function deleteBranchEntries(db: SqliteDatabase, sessionId: string) {
db.prepare("DELETE FROM branch_entries WHERE session_id = ?").run(sessionId);
}
export function insertBranchEntry(
db: SqliteDatabase,
sessionId: string,
branchId: string,
entryId: string,
entrySeq: number,
entryType: string,
customType: string | null,
) {
db.prepare(
`INSERT INTO branch_entries
(session_id, branch_id, entry_id, entry_seq, entry_type, custom_type)
VALUES (?, ?, ?, ?, ?, ?)`,
).run(sessionId, branchId, entryId, entrySeq, entryType, customType);
}
export function insertBranchEntriesForPath(db: SqliteDatabase, sessionId: string, branchId: string, leafId: string) {
db.prepare(
`WITH RECURSIVE path(id, entry_seq, parent_id, type, custom_type) AS (
SELECT id, seq, parent_id, type,
CASE WHEN type = 'custom' THEN json_extract(payload, '$.customType') ELSE NULL END
FROM entries
WHERE session_id = ? AND id = ?
UNION ALL
SELECT parent.id, parent.seq, parent.parent_id, parent.type,
CASE WHEN parent.type = 'custom' THEN json_extract(parent.payload, '$.customType') ELSE NULL END
FROM entries AS parent
JOIN path AS child ON child.parent_id = parent.id
WHERE parent.session_id = ?
)
INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq, entry_type, custom_type)
SELECT ?, ?, id, entry_seq, type, custom_type FROM path`,
).run(sessionId, leafId, sessionId, sessionId, branchId);
}
export function readBranchContainingEntry(db: SqliteDatabase, sessionId: string, entryId: string) {
const row = db
.prepare(
`SELECT b.branch_id, b.entry_seq
FROM branch_entries AS b
WHERE b.session_id = ? AND b.entry_id = ?
ORDER BY b.branch_id
LIMIT 1`,
)
.get<{ branch_id: string; entry_seq: number }>(sessionId, entryId);
return row === undefined ? undefined : { branchId: row.branch_id, entrySeq: row.entry_seq };
}
export function copyBranchEntriesThroughSeq(
db: SqliteDatabase,
sessionId: string,
targetBranchId: string,
sourceBranchId: string,
throughSeq: number,
) {
db.prepare(
`INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq, entry_type, custom_type)
SELECT session_id, ?, entry_id, entry_seq, entry_type, custom_type
FROM branch_entries
WHERE session_id = ? AND branch_id = ? AND entry_seq <= ?`,
).run(targetBranchId, sessionId, sourceBranchId, throughSeq);
}

40
packages/storage/sqlite-node/src/sqlite/storage/branch-tips.ts Normal file
View file

@ -0,0 +1,40 @@
import type { SqliteDatabase } from "../types.ts";
export function readBranchTipIds(db: SqliteDatabase, sessionId: string) {
return db
.prepare("SELECT tip_id FROM branch_tips WHERE session_id = ? ORDER BY tip_id")
.all<{ tip_id: string }>(sessionId)
.map((row) => row.tip_id);
}
export function readBranchTipBranchId(db: SqliteDatabase, sessionId: string, tipId: string) {
const tip = db
.prepare("SELECT branch_id FROM branch_tips WHERE session_id = ? AND tip_id = ?")
.get<{ branch_id: string }>(sessionId, tipId);
return tip?.branch_id;
}
export function insertBranchTip(db: SqliteDatabase, sessionId: string, tipId: string, branchId: string) {
db.prepare("INSERT INTO branch_tips (session_id, tip_id, branch_id) VALUES (?, ?, ?)").run(
sessionId,
tipId,
branchId,
);
}
export function updateBranchTip(
db: SqliteDatabase,
sessionId: string,
branchId: string,
oldTipId: string,
newTipId: string,
) {
const result = db
.prepare("UPDATE branch_tips SET tip_id = ? WHERE session_id = ? AND branch_id = ? AND tip_id = ?")
.run(newTipId, sessionId, branchId, oldTipId);
return result.changes === 1;
}
export function deleteBranchTips(db: SqliteDatabase, sessionId: string) {
db.prepare("DELETE FROM branch_tips WHERE session_id = ?").run(sessionId);
}

75
packages/storage/sqlite-node/src/sqlite/storage/entries.ts Normal file
View file

@ -0,0 +1,75 @@
import type { Entry, EntryOrder } from "@earendil-works/pi-agent-core/experimental";
import type { SqliteDatabase } from "../types.ts";
export interface EntryRow {
session_id: string;
seq: number;
id: string;
parent_id: string | null;
type: Entry["type"];
timestamp: string;
payload: string;
}
export interface NewEntryRow {
seq: number;
id: string;
parentId: string | null;
type: Entry["type"];
timestamp: string;
payload: string;
}
export function entryPayload(entry: Entry): Record<string, unknown> {
const { type: _type, id: _id, seq: _seq, parentId: _parentId, timestamp: _timestamp, ...payload } = entry;
return payload;
}
function orderedSql(order: EntryOrder | undefined): string {
return order === "oldestFirst" ? "ASC" : "DESC";
}
export function insertEntryRow(db: SqliteDatabase, sessionId: string, entry: NewEntryRow) {
db.prepare(
"INSERT INTO entries (session_id, id, seq, parent_id, type, timestamp, payload) VALUES (?, ?, ?, ?, ?, ?, ?)",
).run(sessionId, entry.id, entry.seq, entry.parentId, entry.type, entry.timestamp, entry.payload);
}
export function readEntryRow(db: SqliteDatabase, sessionId: string, entryId: string) {
return db
.prepare(
"SELECT session_id, seq, id, parent_id, type, timestamp, payload FROM entries WHERE session_id = ? AND id = ?",
)
.get<EntryRow>(sessionId, entryId);
}
export function readEntryRows(
db: SqliteDatabase,
sessionId: string,
options: { afterSeq?: number; order?: EntryOrder } = {},
) {
const predicates = ["session_id = ?"];
const params: unknown[] = [sessionId];
if (options.afterSeq !== undefined) {
predicates.push("seq > ?");
params.push(options.afterSeq);
}
return db
.prepare(
`SELECT session_id, seq, id, parent_id, type, timestamp, payload
FROM entries
WHERE ${predicates.join(" AND ")}
ORDER BY seq ${orderedSql(options.order)}`,
)
.all<EntryRow>(...params);
}
export function idExistsInEntries(db: SqliteDatabase, sessionId: string, id: string) {
return !!db
.prepare("SELECT 1 AS found FROM entries WHERE session_id = ? AND id = ? LIMIT 1")
.get<{ found: number }>(sessionId, id);
}
export function deleteEntryRows(db: SqliteDatabase, sessionId: string) {
db.prepare("DELETE FROM entries WHERE session_id = ?").run(sessionId);
}

73
packages/storage/sqlite-node/src/sqlite/storage/facts.ts Normal file
View file

@ -0,0 +1,73 @@
import type { SqliteDatabase } from "../types.ts";
export interface FactRow {
session_id: string;
seq: number;
kind: string;
key: string | null;
value: string | null;
}
export function appendFact(
db: SqliteDatabase,
sessionId: string,
seq: number,
kind: string,
key: string | null,
value: string | null,
) {
db.prepare("INSERT INTO facts (session_id, seq, kind, key, value) VALUES (?, ?, ?, ?, ?)").run(
sessionId,
seq,
kind,
key,
value,
);
}
export function readLatestFact(db: SqliteDatabase, sessionId: string, kind: string, key: string | null) {
return db
.prepare(
`SELECT session_id, seq, kind, key, value
FROM facts
WHERE session_id = ? AND kind = ? AND key IS ?
ORDER BY seq DESC
LIMIT 1`,
)
.get<FactRow>(sessionId, kind, key);
}
export function readLatestLabelFacts(db: SqliteDatabase, sessionId: string) {
return db
.prepare(
`SELECT key, value FROM (
SELECT key, value, ROW_NUMBER() OVER (PARTITION BY key ORDER BY seq DESC) AS rank
FROM facts
WHERE session_id = ? AND kind = 'label'
)
WHERE rank = 1 AND value IS NOT NULL
ORDER BY key`,
)
.all<{ key: string; value: string }>(sessionId);
}
export function readFactRows(db: SqliteDatabase, sessionId: string, options: { afterSeq?: number } = {}) {
const predicates = ["session_id = ?"];
const params: unknown[] = [sessionId];
if (options.afterSeq !== undefined) {
predicates.push("seq > ?");
params.push(options.afterSeq);
}
return db
.prepare(
`SELECT session_id, seq, kind, key, value
FROM facts
WHERE ${predicates.join(" AND ")}
ORDER BY seq`,
)
.all<FactRow>(...params);
}
export function deleteFactRows(db: SqliteDatabase, sessionId: string) {
db.prepare("DELETE FROM facts WHERE session_id = ?").run(sessionId);
}

459
packages/storage/sqlite-node/src/sqlite/storage/index.ts
View file

@ -1,459 +0,0 @@
import type {
SessionBranchQuery,
SessionEntryCursorOptions,
SessionStats,
SessionTreeEntry,
} from "@earendil-works/pi-agent-core";
import { SessionError, toError } from "@earendil-works/pi-agent-core";
import type { SqliteDatabase, SqliteSessionMetadata } from "../types.ts";
import {
appendEntryToBranchCache,
type CachedBranch,
isCachedBranchValid,
queryCachedBranchRows,
readCachedBranch,
readCachedBranchRows,
readCachedEntryRowsByType,
readCachedEntrySeq,
readNewestCachedStopSeq,
rebuildCachedBranch,
} from "./branch-cache.ts";
import { decodeEntry, encodeEntry, type SessionEntryRow } from "./session-entries.ts";
import {
applyEntryToMaterializedState,
createEmptyMaterializedState,
type EntryMaterializedRow,
entryMaterializedValues,
materializedStateFromRows,
materializedStateValues,
type SessionMaterializedRow,
type SessionMaterializedState,
serializeSummary,
} from "./session-materialized.ts";
import { advanceSequence, getNextSequence } from "./session-sequences.ts";
import { rowToMetadata, type SessionRow } from "./sessions.ts";
import { invalidEntry, invalidSession, leafIdAfterEntry } from "./shared.ts";
function decodeEntryRows(entryRows: SessionEntryRow[]): SessionTreeEntry[] {
const entries: SessionTreeEntry[] = [];
for (const entryRow of entryRows) {
try {
const entry = decodeEntry(entryRow);
entries.push(entry);
} catch (error) {
throw invalidEntry(`failed to decode entry ${entryRow.id}`, toError(error));
}
}
return entries;
}
async function loadSqliteSession(
db: SqliteDatabase,
sessionId: string,
): Promise<{
row: SessionRow;
materializedState: SessionMaterializedState;
}> {
const row = await db
.prepare("SELECT id, created_at, metadata, cwd, parent_session_id, active_leaf_id FROM sessions WHERE id = ?")
.get<SessionRow>(sessionId);
if (!row) throw new SessionError("not_found", `Session not found: ${sessionId}`);
const materializedRow = await db
.prepare("SELECT session_id, payload FROM session_materialized WHERE session_id = ?")
.get<SessionMaterializedRow>(sessionId);
if (!materializedRow) throw invalidSession(`missing materialized row for session ${sessionId}`);
const entryMaterializedRows = await db
.prepare(
"SELECT session_id, entry_seq, type, payload FROM entry_materialized WHERE session_id = ? ORDER BY entry_seq, type",
)
.all<EntryMaterializedRow>(sessionId);
return {
row,
materializedState: materializedStateFromRows(materializedRow, entryMaterializedRows),
};
}
export class SqliteSessionConnection {
private readonly db: SqliteDatabase;
readonly metadata: SqliteSessionMetadata;
private byId: Map<string, SessionTreeEntry>;
private materializedState: SessionMaterializedState;
async findEntriesOnBranch(query: SessionBranchQuery & { start: string | null }): Promise<SessionTreeEntry[]> {
if (query.limit !== undefined && (!Number.isInteger(query.limit) || query.limit <= 0)) {
throw new RangeError("Session branch query limit must be a positive integer");
}
if (query.start === null) return [];
const startId = query.start;
let cached = await readCachedBranch(this.db, this.metadata.id, startId);
const validationStartSeq =
cached && query.order !== "oldestFirst"
? await readNewestCachedStopSeq(this.db, this.metadata.id, cached, query.stopAtType, query.stopAtId)
: undefined;
if (!cached || !(await isCachedBranchValid(this.db, this.metadata.id, cached, startId, validationStartSeq))) {
if (query.order !== "oldestFirst" && (query.stopAtId !== undefined || query.stopAtType !== undefined)) {
return this.findEntriesOnCanonicalBranch({ ...query, start: startId });
}
cached = await this.repairBranchCacheForQuery(startId, cached?.branchId);
}
const decoded = decodeEntryRows(await queryCachedBranchRows(this.db, this.metadata.id, cached, query));
const filtered = decoded.filter(
(entry) =>
(query.type === undefined || entry.type === query.type) &&
(query.customType === undefined || (entry.type === "custom" && entry.customType === query.customType)),
);
const entries = query.limit === undefined ? filtered : filtered.slice(0, query.limit);
for (const entry of entries) this.byId.set(entry.id, entry);
return entries;
}
private async findEntriesOnCanonicalBranch(
query: SessionBranchQuery & { start: string },
): Promise<SessionTreeEntry[]> {
const rows: SessionEntryRow[] = [];
const visited = new Set<string>();
let currentId: string | null = query.start;
while (currentId !== null) {
if (visited.has(currentId)) throw invalidSession(`cycle in parent chain at entry ${currentId}`);
visited.add(currentId);
const row: SessionEntryRow | undefined = await this.db
.prepare(
"SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND id = ?",
)
.get<SessionEntryRow>(this.metadata.id, currentId);
if (!row) {
if (currentId === query.start) throw new SessionError("not_found", `Entry ${query.start} not found`);
throw invalidSession(`Entry ${currentId} not found`);
}
rows.push(row);
if (row.id === query.stopAtId || row.type === query.stopAtType) break;
currentId = row.parent_id;
}
const entries = decodeEntryRows(rows).filter(
(entry) =>
(query.type === undefined || entry.type === query.type) &&
(query.customType === undefined || (entry.type === "custom" && entry.customType === query.customType)),
);
const limited = query.limit === undefined ? entries : entries.slice(0, query.limit);
for (const entry of limited) this.byId.set(entry.id, entry);
return limited;
}
private async repairBranchCacheForQuery(leafId: string, branchIdToReplace?: string): Promise<CachedBranch> {
const visited = new Set<string>();
let currentId: string | null = leafId;
while (currentId !== null) {
if (visited.has(currentId)) throw invalidSession(`cycle in parent chain at entry ${currentId}`);
visited.add(currentId);
const row: { parent_id: string | null } | undefined = await this.db
.prepare("SELECT parent_id FROM session_entries WHERE session_id = ? AND id = ?")
.get<{ parent_id: string | null }>(this.metadata.id, currentId);
if (!row) {
if (currentId === leafId) throw new SessionError("not_found", `Entry ${leafId} not found`);
throw invalidSession(`Entry ${currentId} not found`);
}
currentId = row.parent_id;
}
try {
await rebuildCachedBranch(this.db, this.metadata.id, leafId, branchIdToReplace);
} catch (error) {
if (error instanceof SessionError) throw error;
throw new SessionError("storage", `Failed to rebuild SQLite branch cache at entry ${leafId}`, toError(error));
}
const cached = await readCachedBranch(this.db, this.metadata.id, leafId);
if (!cached || !(await isCachedBranchValid(this.db, this.metadata.id, cached, leafId))) {
throw invalidSession(`branch cache repair did not produce a valid path to entry ${leafId}`);
}
return cached;
}
async readPathToRootOrCompaction(leafId: string | null): Promise<SessionTreeEntry[]> {
if (leafId === null) return [];
const cached = await readCachedBranch(this.db, this.metadata.id, leafId);
if (cached) {
const compactionRows = await readCachedEntryRowsByType(this.db, this.metadata.id, cached, "compaction");
let startSeq = 0;
let expectedStartId: string | null = null;
let pendingStop: { id: string; seq: number } | undefined;
for (const compactionRow of compactionRows) {
if (pendingStop && pendingStop.seq >= compactionRow.entry_seq) {
startSeq = pendingStop.seq;
expectedStartId = pendingStop.id;
break;
}
const compaction = decodeEntryRows([compactionRow])[0]!;
if (compaction.type !== "compaction") throw invalidSession(`entry ${compaction.id} is not a compaction`);
if (compaction.retainedTail) {
startSeq = compactionRow.entry_seq;
expectedStartId = compaction.id;
pendingStop = undefined;
break;
} else if (compaction.firstKeptEntryId !== undefined) {
const firstKeptSeq = await readCachedEntrySeq(
this.db,
this.metadata.id,
cached.branchId,
compaction.firstKeptEntryId,
);
pendingStop =
firstKeptSeq !== undefined && firstKeptSeq < compactionRow.entry_seq
? { id: compaction.firstKeptEntryId, seq: firstKeptSeq }
: undefined;
} else {
pendingStop = undefined;
}
}
if (startSeq === 0 && pendingStop) {
startSeq = pendingStop.seq;
expectedStartId = pendingStop.id;
}
const entries = decodeEntryRows(await readCachedBranchRows(this.db, this.metadata.id, cached, startSeq));
if (this.isValidCachedPath(entries, leafId, expectedStartId)) {
for (const entry of entries) this.byId.set(entry.id, entry);
return entries;
}
}
const entries = await this.repairBranchCache(leafId, cached?.branchId);
return this.trimPathToRootOrCompaction(entries);
}
private async repairBranchCache(leafId: string, branchIdToReplace?: string): Promise<SessionTreeEntry[]> {
const entries = await this.readCanonicalPathToRoot(leafId);
try {
await rebuildCachedBranch(this.db, this.metadata.id, leafId, branchIdToReplace);
} catch (error) {
if (error instanceof SessionError) throw error;
throw new SessionError("storage", `Failed to rebuild SQLite branch cache at entry ${leafId}`, toError(error));
}
return entries;
}
private async readCanonicalPathToRoot(leafId: string): Promise<SessionTreeEntry[]> {
const path: SessionTreeEntry[] = [];
let current = await this.readEntry(leafId);
if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`);
const visited = new Set<string>();
while (current) {
if (visited.has(current.id)) throw invalidSession(`cycle in parent chain at entry ${current.id}`);
visited.add(current.id);
path.push(current);
if (!current.parentId) break;
const parent = await this.readEntry(current.parentId);
if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`);
current = parent;
}
return path.reverse();
}
private isValidCachedPath(
entries: readonly SessionTreeEntry[],
leafId: string,
expectedStartId: string | null,
): boolean {
if (entries.length === 0 || entries.at(-1)!.id !== leafId) return false;
if (expectedStartId === null ? entries[0]!.parentId !== null : entries[0]!.id !== expectedStartId) return false;
for (let index = 1; index < entries.length; index++) {
if (entries[index]!.parentId !== entries[index - 1]!.id) return false;
}
return true;
}
private trimPathToRootOrCompaction(entries: readonly SessionTreeEntry[]): SessionTreeEntry[] {
const path: SessionTreeEntry[] = [];
let stopAtEntryId: string | null = null;
for (let index = entries.length - 1; index >= 0; index--) {
const entry = entries[index]!;
path.push(entry);
if (stopAtEntryId !== null && entry.id === stopAtEntryId) break;
if (entry.type === "compaction") {
if (entry.retainedTail) break;
stopAtEntryId = entry.firstKeptEntryId ?? null;
}
}
return path.reverse();
}
private constructor(
db: SqliteDatabase,
metadata: SqliteSessionMetadata,
materializedState: SessionMaterializedState,
) {
this.db = db;
this.metadata = metadata;
this.byId = new Map<string, SessionTreeEntry>();
this.materializedState = materializedState;
}
static async open(db: SqliteDatabase, metadata: SqliteSessionMetadata): Promise<SqliteSessionConnection> {
const loaded = await loadSqliteSession(db, metadata.id);
return new SqliteSessionConnection(db, rowToMetadata(loaded.row, metadata.path), loaded.materializedState);
}
static async create(
db: SqliteDatabase,
path: string,
options: {
cwd: string;
sessionId: string;
parentSessionId?: string;
metadata?: Record<string, unknown>;
},
): Promise<SqliteSessionConnection> {
const createdAt = new Date().toISOString();
await db
.prepare(
"INSERT INTO sessions (id, created_at, metadata, cwd, parent_session_id, active_leaf_id) VALUES (?, ?, ?, ?, ?, ?)",
)
.run(
options.sessionId,
createdAt,
options.metadata === undefined ? null : JSON.stringify(options.metadata),
options.cwd,
options.parentSessionId ?? null,
null,
);
await db.prepare("INSERT INTO session_sequences (session_id, next_seq) VALUES (?, ?)").run(options.sessionId, 1);
await db
.prepare("INSERT INTO session_materialized (session_id, payload) VALUES (?, ?)")
.run(...materializedStateValues(options.sessionId, createEmptyMaterializedState()));
return new SqliteSessionConnection(
db,
{
id: options.sessionId,
createdAt,
cwd: options.cwd,
path,
parentSessionId: options.parentSessionId,
metadata: options.metadata,
},
createEmptyMaterializedState(),
);
}
async getLabel(id: string): Promise<string | undefined> {
return this.materializedState.labelsById.get(id);
}
async getName(): Promise<string | undefined> {
return this.materializedState.name;
}
async getStats(): Promise<SessionStats> {
const { messageCount, cachedTokens, uncachedTokens, totalTokens, costTotal } = this.materializedState;
return { messageCount, cachedTokens, uncachedTokens, totalTokens, costTotal };
}
async readHead(): Promise<{ leafId: string | null }> {
const row = await this.db
.prepare(
`SELECT
s.active_leaf_id,
(s.active_leaf_id IS NULL OR EXISTS (
SELECT 1 FROM session_entries AS e WHERE e.session_id = s.id AND e.id = s.active_leaf_id
)) AS active_leaf_exists
FROM sessions AS s
WHERE s.id = ?`,
)
.get<{ active_leaf_id: string | null; active_leaf_exists: number }>(this.metadata.id);
if (!row) throw new SessionError("not_found", `Session not found: ${this.metadata.id}`);
if (row.active_leaf_exists === 0) {
throw new SessionError("invalid_session", `Entry ${row.active_leaf_id} not found`);
}
return { leafId: row.active_leaf_id };
}
async appendEntry(entry: SessionTreeEntry, options: { transaction?: boolean } = {}): Promise<void> {
if (entry.type === "leaf" && entry.targetId !== null && !(await this.readEntry(entry.targetId))) {
throw new SessionError("not_found", `Entry ${entry.targetId} not found`);
}
const encoded = encodeEntry(entry);
const nextMaterializedState: SessionMaterializedState = {
...this.materializedState,
labelsById: new Map(this.materializedState.labelsById),
modelThinkingConfigs: [...this.materializedState.modelThinkingConfigs],
currentModel: this.materializedState.currentModel ? { ...this.materializedState.currentModel } : null,
};
const nextLeafId = leafIdAfterEntry(entry);
try {
applyEntryToMaterializedState(nextMaterializedState, entry);
const write = async () => {
const nextSeq = await getNextSequence(this.db, this.metadata.id);
await this.db
.prepare(
"INSERT INTO session_entries (session_id, id, entry_seq, parent_id, type, timestamp, payload) VALUES (?, ?, ?, ?, ?, ?, ?)",
)
.run(this.metadata.id, entry.id, nextSeq, entry.parentId, entry.type, entry.timestamp, encoded.payload);
await advanceSequence(this.db, this.metadata.id, nextSeq);
await this.db
.prepare("UPDATE session_materialized SET payload = ? WHERE session_id = ?")
.run(serializeSummary(nextMaterializedState), this.metadata.id);
for (const materializedEntry of entryMaterializedValues(entry)) {
await this.db
.prepare("INSERT INTO entry_materialized (session_id, entry_seq, type, payload) VALUES (?, ?, ?, ?)")
.run(this.metadata.id, nextSeq, materializedEntry.type, materializedEntry.payload);
}
await this.db
.prepare("UPDATE sessions SET active_leaf_id = ? WHERE id = ?")
.run(nextLeafId, this.metadata.id);
await appendEntryToBranchCache(
this.db,
this.metadata.id,
entry.id,
nextSeq,
entry.parentId,
async (parentId) => {
await this.repairBranchCache(parentId);
},
);
};
if (options.transaction === false) await write();
else await this.db.transaction(write);
this.materializedState = nextMaterializedState;
this.byId.set(entry.id, entry);
} catch (error) {
if (error instanceof SessionError) throw error;
throw new SessionError("storage", `Failed to append SQLite session entry ${entry.id}`, toError(error));
}
}
async readEntry(id: string): Promise<SessionTreeEntry | undefined> {
const cached = this.byId.get(id);
if (cached) return cached;
const row = await this.db
.prepare(
"SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND id = ?",
)
.get<SessionEntryRow>(this.metadata.id, id);
if (!row) return undefined;
try {
const entry = decodeEntry(row);
this.byId.set(entry.id, entry);
return entry;
} catch (error) {
throw invalidEntry(`failed to decode entry ${row.id}`, toError(error));
}
}
async readEntries(options?: SessionEntryCursorOptions): Promise<SessionTreeEntry[]> {
const afterEntrySeq = options?.afterEntrySeq ?? 0;
const rows =
options?.limit === undefined
? await this.db
.prepare(
"SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND entry_seq > ? ORDER BY entry_seq",
)
.all<SessionEntryRow>(this.metadata.id, afterEntrySeq)
: await this.db
.prepare(
"SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND entry_seq > ? ORDER BY entry_seq LIMIT ?",
)
.all<SessionEntryRow>(this.metadata.id, afterEntrySeq, options.limit);
const entries = decodeEntryRows(rows);
for (const entry of entries) {
this.byId.set(entry.id, entry);
}
return entries;
}
}

116
packages/storage/sqlite-node/src/sqlite/storage/lanes.ts Normal file
View file

@ -0,0 +1,116 @@
import { SessionError } from "@earendil-works/pi-agent-core/experimental";
import type { SqliteDatabase } from "../types.ts";
export interface LaneRow {
session_id: string;
lane: string;
leaf_id: string | null;
}
export interface LaneMoveRow {
session_id: string;
seq: number;
lane: string;
leaf_id: string | null;
}
export function createInitialLane(db: SqliteDatabase, sessionId: string, lane = "main", leafId: string | null = null) {
db.prepare("INSERT INTO lanes (session_id, lane, leaf_id) VALUES (?, ?, ?)").run(sessionId, lane, leafId);
}
export function readLanes(db: SqliteDatabase, sessionId: string) {
const rows = db
.prepare(
`SELECT
l.session_id,
l.lane,
l.leaf_id,
(l.leaf_id IS NULL OR EXISTS (
SELECT 1 FROM entries AS e WHERE e.session_id = l.session_id AND e.id = l.leaf_id
)) AS leaf_exists
FROM lanes AS l
WHERE l.session_id = ?
ORDER BY l.lane`,
)
.all<LaneRow & { leaf_exists: number }>(sessionId);
for (const row of rows) {
if (row.leaf_exists === 0) {
throw new SessionError("storage", `Lane ${row.lane} points at missing entry ${row.leaf_id}`);
}
}
return rows.map(({ session_id, lane, leaf_id }) => ({ session_id, lane, leaf_id }));
}
export function readLane(db: SqliteDatabase, sessionId: string, lane: string) {
return db
.prepare("SELECT session_id, lane, leaf_id FROM lanes WHERE session_id = ? AND lane = ?")
.get<LaneRow>(sessionId, lane);
}
export function readLaneHead(db: SqliteDatabase, sessionId: string, lane: string) {
const row = db
.prepare(
`SELECT
l.leaf_id,
(l.leaf_id IS NULL OR EXISTS (
SELECT 1 FROM entries AS e WHERE e.session_id = l.session_id AND e.id = l.leaf_id
)) AS leaf_exists
FROM lanes AS l
WHERE l.session_id = ? AND l.lane = ?`,
)
.get<{ leaf_id: string | null; leaf_exists: number }>(sessionId, lane);
if (!row) throw new SessionError("invalid_lane", `Lane not found: ${lane}`);
if (row.leaf_exists === 0) throw new SessionError("storage", `Entry ${row.leaf_id} not found`);
return { leafId: row.leaf_id };
}
export function createLane(db: SqliteDatabase, sessionId: string, seq: number, lane: string, leafId: string | null) {
db.prepare("INSERT INTO lanes (session_id, lane, leaf_id) VALUES (?, ?, ?)").run(sessionId, lane, leafId);
appendLaneMove(db, sessionId, seq, lane, leafId);
}
export function moveLane(db: SqliteDatabase, sessionId: string, seq: number, lane: string, leafId: string | null) {
const result = db
.prepare("UPDATE lanes SET leaf_id = ? WHERE session_id = ? AND lane = ?")
.run(leafId, sessionId, lane);
if (result.changes !== 1) throw new SessionError("invalid_lane", `Lane not found: ${lane}`);
appendLaneMove(db, sessionId, seq, lane, leafId);
}
export function setLaneLeaf(db: SqliteDatabase, sessionId: string, lane: string, leafId: string | null) {
const result = db
.prepare("UPDATE lanes SET leaf_id = ? WHERE session_id = ? AND lane = ?")
.run(leafId, sessionId, lane);
if (result.changes !== 1) throw new SessionError("invalid_lane", `Lane not found: ${lane}`);
}
export function readLaneMoveRows(db: SqliteDatabase, sessionId: string, options: { afterSeq?: number } = {}) {
const predicates = ["session_id = ?"];
const params: unknown[] = [sessionId];
if (options.afterSeq !== undefined) {
predicates.push("seq > ?");
params.push(options.afterSeq);
}
return db
.prepare(
`SELECT session_id, seq, lane, leaf_id
FROM lane_moves
WHERE ${predicates.join(" AND ")}
ORDER BY seq`,
)
.all<LaneMoveRow>(...params);
}
export function deleteLaneRows(db: SqliteDatabase, sessionId: string) {
db.prepare("DELETE FROM lane_moves WHERE session_id = ?").run(sessionId);
db.prepare("DELETE FROM lanes WHERE session_id = ?").run(sessionId);
}
function appendLaneMove(db: SqliteDatabase, sessionId: string, seq: number, lane: string, leafId: string | null) {
db.prepare("INSERT INTO lane_moves (session_id, seq, lane, leaf_id) VALUES (?, ?, ?, ?)").run(
sessionId,
seq,
lane,
leafId,
);
}

65
packages/storage/sqlite-node/src/sqlite/storage/leases.ts Normal file
View file

@ -0,0 +1,65 @@
import type { SqliteDatabase } from "../types.ts";
export interface SessionLease {
ownerId: string;
fence: number;
expiresAtMs: number;
}
interface SessionLeaseRow {
owner_id: string;
fence: number;
expires_at_ms: number;
}
export function acquireSessionLease(
db: SqliteDatabase,
sessionId: string,
ownerId: string,
now: number,
expiresAtMs: number,
) {
const row = db
.prepare(
`INSERT INTO leases (session_id, owner_id, fence, expires_at_ms)
VALUES (?, ?, 1, ?)
ON CONFLICT(session_id) DO UPDATE SET
owner_id = excluded.owner_id,
fence = leases.fence + 1,
expires_at_ms = excluded.expires_at_ms
WHERE leases.expires_at_ms <= ?
RETURNING owner_id, fence, expires_at_ms`,
)
.get<SessionLeaseRow>(sessionId, ownerId, expiresAtMs, now);
return row === undefined ? undefined : { ownerId: row.owner_id, fence: row.fence, expiresAtMs: row.expires_at_ms };
}
export function renewSessionLease(
db: SqliteDatabase,
sessionId: string,
lease: SessionLease,
now: number,
expiresAtMs: number,
) {
const result = db
.prepare(
`UPDATE leases
SET expires_at_ms = ?
WHERE session_id = ? AND owner_id = ? AND fence = ? AND expires_at_ms > ?`,
)
.run(expiresAtMs, sessionId, lease.ownerId, lease.fence, now);
if (result.changes === 1) lease.expiresAtMs = expiresAtMs;
return result.changes === 1;
}
export function releaseSessionLease(db: SqliteDatabase, sessionId: string, lease: SessionLease) {
db.prepare("DELETE FROM leases WHERE session_id = ? AND owner_id = ? AND fence = ?").run(
sessionId,
lease.ownerId,
lease.fence,
);
}
export function deleteSessionLease(db: SqliteDatabase, sessionId: string) {
db.prepare("DELETE FROM leases WHERE session_id = ?").run(sessionId);
}

95
packages/storage/sqlite-node/src/sqlite/storage/records.ts Normal file
View file

@ -0,0 +1,95 @@
import type { SqliteDatabase } from "../types.ts";
export interface RecordRow {
session_id: string;
seq: number;
id: string;
lane: string;
run_id: string | null;
type: string;
op_kind: string | null;
timestamp: string;
payload: string;
}
export interface NewRecordRow {
seq: number;
id: string;
lane: string;
runId?: string;
type: string;
opKind?: string;
timestamp: string;
payload: string;
}
export function appendRecordRow(db: SqliteDatabase, sessionId: string, record: NewRecordRow) {
db.prepare(
`INSERT INTO records
(session_id, seq, id, lane, run_id, type, op_kind, timestamp, payload)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
).run(
sessionId,
record.seq,
record.id,
record.lane,
record.runId ?? null,
record.type,
record.opKind ?? null,
record.timestamp,
record.payload,
);
}
export function idExistsInRecords(db: SqliteDatabase, sessionId: string, id: string) {
return !!db
.prepare("SELECT 1 AS found FROM records WHERE session_id = ? AND id = ? LIMIT 1")
.get<{ found: number }>(sessionId, id);
}
export function deleteRecordRows(db: SqliteDatabase, sessionId: string) {
db.prepare("DELETE FROM records WHERE session_id = ?").run(sessionId);
}
export function readRecordRows(
db: SqliteDatabase,
sessionId: string,
query: {
lane?: string;
type?: string;
runId?: string;
afterSeq?: number;
order?: "newestFirst" | "oldestFirst";
limit?: number;
} = {},
) {
const predicates = ["session_id = ?"];
const params: unknown[] = [sessionId];
if (query.lane !== undefined) {
predicates.push("lane = ?");
params.push(query.lane);
}
if (query.type !== undefined) {
predicates.push("type = ?");
params.push(query.type);
}
if (query.runId !== undefined) {
predicates.push("run_id = ?");
params.push(query.runId);
}
if (query.afterSeq !== undefined) {
predicates.push("seq > ?");
params.push(query.afterSeq);
}
const limit = query.limit === undefined ? "" : " LIMIT ?";
if (query.limit !== undefined) params.push(query.limit);
const direction = query.order === "oldestFirst" ? "ASC" : "DESC";
return db
.prepare(
`SELECT session_id, seq, id, lane, run_id, type, op_kind, timestamp, payload
FROM records
WHERE ${predicates.join(" AND ")}
ORDER BY seq ${direction}${limit}`,
)
.all<RecordRow>(...params);
}

217
packages/storage/sqlite-node/src/sqlite/storage/session-entries.ts
View file

@ -1,217 +0,0 @@
import type { SessionTreeEntry, SessionTreeEntryBase } from "@earendil-works/pi-agent-core";
import { invalidEntry, isRecord } from "./shared.ts";
export interface SessionEntryRow {
session_id: string;
id: string;
entry_seq: number;
parent_id: string | null;
type: SessionTreeEntry["type"];
timestamp: string;
payload: string;
}
export type EncodedEntry = {
payload: string;
};
type EntryPayload<TEntry extends SessionTreeEntry> = Omit<TEntry, keyof SessionTreeEntryBase | "type">;
type MessagePayload = EntryPayload<Extract<SessionTreeEntry, { type: "message" }>>;
type ThinkingLevelChangePayload = EntryPayload<Extract<SessionTreeEntry, { type: "thinking_level_change" }>>;
type ModelChangePayload = EntryPayload<Extract<SessionTreeEntry, { type: "model_change" }>>;
type ActiveToolsChangePayload = EntryPayload<Extract<SessionTreeEntry, { type: "active_tools_change" }>>;
type CompactionPayload = EntryPayload<Extract<SessionTreeEntry, { type: "compaction" }>>;
type BranchSummaryPayload = EntryPayload<Extract<SessionTreeEntry, { type: "branch_summary" }>>;
type CustomPayload = EntryPayload<Extract<SessionTreeEntry, { type: "custom" }>>;
type CustomMessagePayload = EntryPayload<Extract<SessionTreeEntry, { type: "custom_message" }>>;
type LabelPayload = EntryPayload<Extract<SessionTreeEntry, { type: "label" }>>;
type SessionInfoPayload = EntryPayload<Extract<SessionTreeEntry, { type: "session_info" }>>;
type LeafPayload = EntryPayload<Extract<SessionTreeEntry, { type: "leaf" }>>;
function parsePayload(row: SessionEntryRow): unknown {
try {
return JSON.parse(row.payload);
} catch (error) {
throw invalidEntry(`entry ${row.id} payload is not valid JSON`, error instanceof Error ? error : undefined);
}
}
function isTextImageContentArray(value: unknown): boolean {
return (
Array.isArray(value) &&
value.every(
(item) =>
isRecord(item) && typeof item.type === "string" && (item.type !== "text" || typeof item.text === "string"),
)
);
}
export function validateSessionTreeEntry(entry: SessionTreeEntry): void {
if (typeof entry.id !== "string" || !entry.id) throw invalidEntry("entry is missing id");
if (entry.parentId !== null && typeof entry.parentId !== "string") {
throw invalidEntry(`entry ${entry.id} has invalid parentId`);
}
if (typeof entry.timestamp !== "string" || !entry.timestamp) {
throw invalidEntry(`entry ${entry.id} is missing timestamp`);
}
switch (entry.type) {
case "message":
if (!isRecord(entry.message) || typeof entry.message.role !== "string") {
throw invalidEntry(`entry ${entry.id} is missing message payload`);
}
break;
case "thinking_level_change":
if (typeof entry.thinkingLevel !== "string") throw invalidEntry(`entry ${entry.id} is missing thinkingLevel`);
break;
case "model_change":
if (typeof entry.provider !== "string" || typeof entry.modelId !== "string") {
throw invalidEntry(`entry ${entry.id} has invalid model_change payload`);
}
break;
case "active_tools_change":
if (
!Array.isArray(entry.activeToolNames) ||
entry.activeToolNames.some((value) => typeof value !== "string")
) {
throw invalidEntry(`entry ${entry.id} has invalid active_tools_change payload`);
}
break;
case "compaction":
if (
typeof entry.summary !== "string" ||
typeof entry.firstKeptEntryId !== "string" ||
typeof entry.tokensBefore !== "number" ||
(entry.retainedTail !== undefined && !Array.isArray(entry.retainedTail))
) {
throw invalidEntry(`entry ${entry.id} has invalid compaction payload`);
}
break;
case "branch_summary":
if (typeof entry.fromId !== "string" || typeof entry.summary !== "string") {
throw invalidEntry(`entry ${entry.id} has invalid branch_summary payload`);
}
break;
case "custom":
if (typeof entry.customType !== "string") throw invalidEntry(`entry ${entry.id} has invalid custom payload`);
break;
case "custom_message":
if (
typeof entry.customType !== "string" ||
typeof entry.display !== "boolean" ||
!(typeof entry.content === "string" || isTextImageContentArray(entry.content))
) {
throw invalidEntry(`entry ${entry.id} has invalid custom_message payload`);
}
break;
case "label":
if (typeof entry.targetId !== "string" || (entry.label !== undefined && typeof entry.label !== "string")) {
throw invalidEntry(`entry ${entry.id} has invalid label payload`);
}
break;
case "session_info":
if (entry.name !== undefined && typeof entry.name !== "string") {
throw invalidEntry(`entry ${entry.id} has invalid session_info payload`);
}
break;
case "leaf":
if (entry.targetId !== null && typeof entry.targetId !== "string") {
throw invalidEntry(`entry ${entry.id} has invalid leaf payload`);
}
break;
default: {
const exhaustive: never = entry;
throw invalidEntry(`unknown entry type ${(exhaustive as { type?: string }).type ?? "unknown"}`);
}
}
}
function entryToPayload<TEntry extends SessionTreeEntry>(entry: TEntry): EntryPayload<TEntry> {
const { type: _type, id: _id, parentId: _parentId, timestamp: _timestamp, ...payload } = entry;
return payload as EntryPayload<TEntry>;
}
export function encodeEntry(entry: SessionTreeEntry): EncodedEntry {
validateSessionTreeEntry(entry);
return { payload: JSON.stringify(entryToPayload(entry)) };
}
export function decodeEntry(row: SessionEntryRow): SessionTreeEntry {
const payload = parsePayload(row);
if (!isRecord(payload)) throw invalidEntry(`entry ${row.id} payload is not an object`);
const base = {
id: row.id,
parentId: row.parent_id,
timestamp: row.timestamp,
};
switch (row.type) {
case "message": {
if (!("message" in payload)) throw invalidEntry(`entry ${row.id} is missing message payload`);
const messagePayload = payload as MessagePayload;
return { ...base, type: "message", ...messagePayload };
}
case "thinking_level_change":
if (typeof payload.thinkingLevel !== "string") throw invalidEntry(`entry ${row.id} is missing thinkingLevel`);
return { ...base, type: "thinking_level_change", ...(payload as ThinkingLevelChangePayload) };
case "model_change":
if (typeof payload.provider !== "string" || typeof payload.modelId !== "string") {
throw invalidEntry(`entry ${row.id} has invalid model_change payload`);
}
return { ...base, type: "model_change", ...(payload as ModelChangePayload) };
case "active_tools_change":
if (
!Array.isArray(payload.activeToolNames) ||
payload.activeToolNames.some((value) => typeof value !== "string")
) {
throw invalidEntry(`entry ${row.id} has invalid active_tools_change payload`);
}
return { ...base, type: "active_tools_change", ...(payload as ActiveToolsChangePayload) };
case "compaction":
if (
typeof payload.summary !== "string" ||
typeof payload.firstKeptEntryId !== "string" ||
typeof payload.tokensBefore !== "number" ||
(payload.retainedTail !== undefined && !Array.isArray(payload.retainedTail))
) {
throw invalidEntry(`entry ${row.id} has invalid compaction payload`);
}
return { ...base, type: "compaction", ...(payload as CompactionPayload) };
case "branch_summary":
if (typeof payload.fromId !== "string" || typeof payload.summary !== "string") {
throw invalidEntry(`entry ${row.id} has invalid branch_summary payload`);
}
return { ...base, type: "branch_summary", ...(payload as BranchSummaryPayload) };
case "custom":
if (typeof payload.customType !== "string") throw invalidEntry(`entry ${row.id} has invalid custom payload`);
return { ...base, type: "custom", ...(payload as CustomPayload) };
case "custom_message":
if (
typeof payload.customType !== "string" ||
typeof payload.display !== "boolean" ||
!("content" in payload)
) {
throw invalidEntry(`entry ${row.id} has invalid custom_message payload`);
}
return { ...base, type: "custom_message", ...(payload as CustomMessagePayload) };
case "label":
if (typeof payload.targetId !== "string") throw invalidEntry(`entry ${row.id} has invalid label payload`);
if (payload.label !== undefined && typeof payload.label !== "string") {
throw invalidEntry(`entry ${row.id} has invalid label payload`);
}
return { ...base, type: "label", ...(payload as LabelPayload) };
case "session_info":
if (payload.name !== undefined && typeof payload.name !== "string") {
throw invalidEntry(`entry ${row.id} has invalid session_info payload`);
}
return { ...base, type: "session_info", ...(payload as SessionInfoPayload) };
case "leaf":
if (payload.targetId !== null && typeof payload.targetId !== "string") {
throw invalidEntry(`entry ${row.id} has invalid leaf payload`);
}
return { ...base, type: "leaf", ...(payload as LeafPayload) };
default:
throw invalidEntry(`unknown entry type ${row.type}`);
}
}

355
packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts
View file

@ -1,355 +0,0 @@
import type { SessionTreeEntry, ThinkingLevel } from "@earendil-works/pi-agent-core";
import { invalidSession, isRecord } from "./shared.ts";
export interface SessionMaterializedRow {
session_id: string;
payload: string;
}
export interface EntryMaterializedRow {
session_id: string;
entry_seq: number;
type: string;
payload: string;
}
export interface ModelThinkingConfig {
provider: string;
modelId: string;
thinkingLevel: ThinkingLevel;
}
export interface SessionMaterializedState {
name: string | undefined;
messageCount: number;
cachedTokens: number;
uncachedTokens: number;
totalTokens: number;
costTotal: number;
labelsById: Map<string, string>;
modelThinkingConfigs: ModelThinkingConfig[];
currentModel: { provider: string; modelId: string } | null;
currentThinkingLevel: ThinkingLevel | null;
}
interface SessionMaterializedSummary {
name?: string;
messageCount: number;
cachedTokens: number;
uncachedTokens: number;
totalTokens: number;
costTotal: number;
currentModel?: { provider: string; modelId: string } | null;
currentThinkingLevel?: ThinkingLevel | null;
}
function compareModelThinkingConfig(left: ModelThinkingConfig, right: ModelThinkingConfig): number {
return (
left.provider.localeCompare(right.provider) ||
left.modelId.localeCompare(right.modelId) ||
left.thinkingLevel.localeCompare(right.thinkingLevel)
);
}
function normalizeModelThinkingConfigs(configs: readonly ModelThinkingConfig[]): ModelThinkingConfig[] {
const unique = new Map<string, ModelThinkingConfig>();
for (const config of configs) {
unique.set(`${config.provider}\u0000${config.modelId}\u0000${config.thinkingLevel}`, config);
}
return [...unique.values()].sort(compareModelThinkingConfig);
}
function addModelThinkingConfig(
state: SessionMaterializedState,
provider: string,
modelId: string,
thinkingLevel: ThinkingLevel,
): void {
state.modelThinkingConfigs = normalizeModelThinkingConfigs([
...state.modelThinkingConfigs,
{ provider, modelId, thinkingLevel },
]);
}
export function isThinkingLevel(value: unknown): value is ThinkingLevel {
return (
value === "off" ||
value === "minimal" ||
value === "low" ||
value === "medium" ||
value === "high" ||
value === "xhigh"
);
}
function getAssistantUsage(message: unknown):
| {
provider: string;
modelId: string;
input: number;
output: number;
cacheRead: number;
cacheWrite: number;
costTotal: number;
}
| undefined {
if (!isRecord(message) || message.role !== "assistant") return undefined;
if (typeof message.provider !== "string" || typeof message.model !== "string") return undefined;
if (!isRecord(message.usage) || !isRecord(message.usage.cost)) return undefined;
const { input, output, cacheRead, cacheWrite } = message.usage;
const costTotal = message.usage.cost.total;
if (
typeof input !== "number" ||
typeof output !== "number" ||
typeof cacheRead !== "number" ||
typeof cacheWrite !== "number" ||
typeof costTotal !== "number"
) {
return undefined;
}
return {
provider: message.provider,
modelId: message.model,
input,
output,
cacheRead,
cacheWrite,
costTotal,
};
}
export function createEmptyMaterializedState(): SessionMaterializedState {
return {
name: undefined,
messageCount: 0,
cachedTokens: 0,
uncachedTokens: 0,
totalTokens: 0,
costTotal: 0,
labelsById: new Map<string, string>(),
modelThinkingConfigs: [],
currentModel: null,
currentThinkingLevel: null,
};
}
export function applyEntryToMaterializedState(state: SessionMaterializedState, entry: SessionTreeEntry): void {
switch (entry.type) {
case "session_info":
state.name = entry.name?.trim() || undefined;
break;
case "label": {
const label = entry.label?.trim();
if (label) {
state.labelsById.set(entry.targetId, label);
} else {
state.labelsById.delete(entry.targetId);
}
break;
}
case "model_change":
state.currentModel = { provider: entry.provider, modelId: entry.modelId };
if (state.currentThinkingLevel) {
addModelThinkingConfig(state, entry.provider, entry.modelId, state.currentThinkingLevel);
}
break;
case "thinking_level_change":
if (!isThinkingLevel(entry.thinkingLevel)) break;
state.currentThinkingLevel = entry.thinkingLevel;
if (state.currentModel) {
addModelThinkingConfig(state, state.currentModel.provider, state.currentModel.modelId, entry.thinkingLevel);
}
break;
case "message": {
state.messageCount += 1;
const usage = getAssistantUsage(entry.message);
if (!usage) break;
state.cachedTokens += usage.cacheRead;
state.uncachedTokens += usage.input + usage.cacheWrite;
state.totalTokens += usage.input + usage.output + usage.cacheRead + usage.cacheWrite;
state.costTotal += usage.costTotal;
state.currentModel = { provider: usage.provider, modelId: usage.modelId };
if (state.currentThinkingLevel) {
addModelThinkingConfig(state, usage.provider, usage.modelId, state.currentThinkingLevel);
}
break;
}
case "compaction":
case "branch_summary": {
const usage = entry.usage;
if (
!isRecord(usage) ||
!isRecord(usage.cost) ||
typeof usage.input !== "number" ||
typeof usage.output !== "number" ||
typeof usage.cacheRead !== "number" ||
typeof usage.cacheWrite !== "number" ||
typeof usage.cost.total !== "number"
) {
break;
}
state.cachedTokens += usage.cacheRead;
state.uncachedTokens += usage.input + usage.cacheWrite;
state.totalTokens += usage.input + usage.output + usage.cacheRead + usage.cacheWrite;
state.costTotal += usage.cost.total;
break;
}
case "active_tools_change":
case "custom":
case "custom_message":
case "leaf":
break;
default: {
const exhaustive: never = entry;
void exhaustive;
break;
}
}
}
export function serializeSummary(state: SessionMaterializedState): string {
const summary: SessionMaterializedSummary = {
name: state.name,
messageCount: state.messageCount,
cachedTokens: state.cachedTokens,
uncachedTokens: state.uncachedTokens,
totalTokens: state.totalTokens,
costTotal: state.costTotal,
currentModel: state.currentModel,
currentThinkingLevel: state.currentThinkingLevel,
};
return JSON.stringify(summary);
}
function parseSummary(json: string): SessionMaterializedSummary {
let parsed: unknown;
try {
parsed = JSON.parse(json);
} catch (error) {
throw invalidSession(
`materialized session summary is not valid JSON`,
error instanceof Error ? error : undefined,
);
}
if (!isRecord(parsed) || Array.isArray(parsed)) {
throw invalidSession("materialized session summary is not an object");
}
const currentModel = parsed.currentModel;
const currentThinkingLevel = parsed.currentThinkingLevel;
if (
(parsed.name !== undefined && typeof parsed.name !== "string") ||
typeof parsed.messageCount !== "number" ||
typeof parsed.cachedTokens !== "number" ||
typeof parsed.uncachedTokens !== "number" ||
typeof parsed.totalTokens !== "number" ||
typeof parsed.costTotal !== "number" ||
(currentModel !== undefined &&
currentModel !== null &&
(!isRecord(currentModel) ||
typeof currentModel.provider !== "string" ||
typeof currentModel.modelId !== "string")) ||
(currentThinkingLevel !== undefined && currentThinkingLevel !== null && !isThinkingLevel(currentThinkingLevel))
) {
throw invalidSession("materialized session summary has invalid fields");
}
return {
name: parsed.name?.trim() || undefined,
messageCount: parsed.messageCount,
cachedTokens: parsed.cachedTokens,
uncachedTokens: parsed.uncachedTokens,
totalTokens: parsed.totalTokens,
costTotal: parsed.costTotal,
currentModel:
currentModel && isRecord(currentModel)
? { provider: currentModel.provider as string, modelId: currentModel.modelId as string }
: (currentModel ?? undefined),
currentThinkingLevel: (currentThinkingLevel as ThinkingLevel | null | undefined) ?? undefined,
};
}
function parseEntryMaterializedPayload(row: EntryMaterializedRow): unknown {
try {
return JSON.parse(row.payload);
} catch (error) {
throw invalidSession(
`materialized entry row ${row.entry_seq} is not valid JSON`,
error instanceof Error ? error : undefined,
);
}
}
export function materializedStateFromRows(
summaryRow: SessionMaterializedRow,
entryRows: EntryMaterializedRow[],
): SessionMaterializedState {
const summary = parseSummary(summaryRow.payload);
const state: SessionMaterializedState = {
name: summary.name,
messageCount: summary.messageCount,
cachedTokens: summary.cachedTokens,
uncachedTokens: summary.uncachedTokens,
totalTokens: summary.totalTokens,
costTotal: summary.costTotal,
labelsById: new Map<string, string>(),
modelThinkingConfigs: [],
currentModel: summary.currentModel ?? null,
currentThinkingLevel: summary.currentThinkingLevel ?? null,
};
for (const row of entryRows) {
const payload = parseEntryMaterializedPayload(row);
if (!isRecord(payload)) throw invalidSession(`materialized entry row ${row.entry_seq} is not an object`);
if (row.type === "label") {
if (typeof payload.targetId !== "string") {
throw invalidSession(`materialized label row ${row.entry_seq} is missing targetId`);
}
if (payload.label !== null && payload.label !== undefined && typeof payload.label !== "string") {
throw invalidSession(`materialized label row ${row.entry_seq} has invalid label`);
}
const label = typeof payload.label === "string" ? payload.label.trim() : "";
if (label) {
state.labelsById.set(payload.targetId, label);
} else {
state.labelsById.delete(payload.targetId);
}
}
}
return state;
}
export function materializedStateValues(
sessionId: string,
state: SessionMaterializedState,
): [sessionId: string, payload: string] {
return [sessionId, serializeSummary(state)];
}
export function entryMaterializedValues(
entry: SessionTreeEntry,
): Array<{ type: EntryMaterializedRow["type"]; payload: string }> {
switch (entry.type) {
case "label":
return [
{
type: "label",
payload: JSON.stringify({ targetId: entry.targetId, label: entry.label ?? null }),
},
];
case "model_change":
case "thinking_level_change":
case "message":
return [];
case "active_tools_change":
case "branch_summary":
case "compaction":
case "custom":
case "custom_message":
case "leaf":
case "session_info":
return [];
default: {
const exhaustive: never = entry;
void exhaustive;
return [];
}
}
}

24
packages/storage/sqlite-node/src/sqlite/storage/session-sequences.ts
View file

@ -1,16 +1,28 @@
import { SessionError } from "@earendil-works/pi-agent-core/experimental";
import type { SqliteDatabase } from "../types.ts";
import { invalidSession } from "./shared.ts";
export async function getNextSequence(db: SqliteDatabase, sessionId: string): Promise<number> {
const sequenceRow = await db
export function createSequence(db: SqliteDatabase, sessionId: string, nextSeq = 1) {
db.prepare("INSERT INTO session_sequences (session_id, next_seq) VALUES (?, ?)").run(sessionId, nextSeq);
}
export function getNextSequence(db: SqliteDatabase, sessionId: string) {
const sequenceRow = db
.prepare("SELECT next_seq FROM session_sequences WHERE session_id = ?")
.get<{ next_seq: number }>(sessionId);
if (!sequenceRow) {
throw invalidSession(`missing sequence row for session ${sessionId}`);
throw new SessionError("storage", `Missing sequence row for session ${sessionId}`);
}
return sequenceRow.next_seq;
}
export async function advanceSequence(db: SqliteDatabase, sessionId: string, nextSeq: number): Promise<void> {
await db.prepare("UPDATE session_sequences SET next_seq = ? WHERE session_id = ?").run(nextSeq + 1, sessionId);
export function setNextSequence(db: SqliteDatabase, sessionId: string, nextSeq: number) {
db.prepare("UPDATE session_sequences SET next_seq = ? WHERE session_id = ?").run(nextSeq, sessionId);
}
export function advanceSequence(db: SqliteDatabase, sessionId: string, seq: number) {
setNextSequence(db, sessionId, seq + 1);
}
export function deleteSequence(db: SqliteDatabase, sessionId: string) {
db.prepare("DELETE FROM session_sequences WHERE session_id = ?").run(sessionId);
}

63
packages/storage/sqlite-node/src/sqlite/storage/session-stats.ts Normal file
View file

@ -0,0 +1,63 @@
import { SessionError, type SessionStats } from "@earendil-works/pi-agent-core/experimental";
import type { Usage } from "@earendil-works/pi-ai";
import type { SqliteDatabase } from "../types.ts";
export interface SessionStatsRow {
session_id: string;
message_count: number;
cached_tokens: number;
uncached_tokens: number;
total_tokens: number;
cost_total: number;
}
export function createStats(db: SqliteDatabase, sessionId: string): void {
db.prepare(
`INSERT INTO session_stats
(session_id, message_count, cached_tokens, uncached_tokens, total_tokens, cost_total)
VALUES (?, 0, 0, 0, 0, 0)`,
).run(sessionId);
}
export function readStats(db: SqliteDatabase, sessionId: string): SessionStats {
const row = db
.prepare(
`SELECT session_id, message_count, cached_tokens, uncached_tokens, total_tokens, cost_total
FROM session_stats
WHERE session_id = ?`,
)
.get<SessionStatsRow>(sessionId);
if (!row) throw new SessionError("storage", `Missing stats row for session ${sessionId}`);
return {
messageCount: row.message_count,
cachedTokens: row.cached_tokens,
uncachedTokens: row.uncached_tokens,
totalTokens: row.total_tokens,
costTotal: row.cost_total,
};
}
export function incrementMessageCount(db: SqliteDatabase, sessionId: string): void {
const result = db
.prepare("UPDATE session_stats SET message_count = message_count + 1 WHERE session_id = ?")
.run(sessionId);
if (result.changes !== 1) throw new SessionError("storage", `Missing stats row for session ${sessionId}`);
}
export function addUsageToStats(db: SqliteDatabase, sessionId: string, usage: Usage): void {
const result = db
.prepare(
`UPDATE session_stats
SET cached_tokens = cached_tokens + ?,
uncached_tokens = uncached_tokens + ?,
total_tokens = total_tokens + ?,
cost_total = cost_total + ?
WHERE session_id = ?`,
)
.run(usage.cacheRead, usage.input + usage.cacheWrite, usage.totalTokens, usage.cost.total, sessionId);
if (result.changes !== 1) throw new SessionError("storage", `Missing stats row for session ${sessionId}`);
}
export function deleteStats(db: SqliteDatabase, sessionId: string): void {
db.prepare("DELETE FROM session_stats WHERE session_id = ?").run(sessionId);
}

62
packages/storage/sqlite-node/src/sqlite/storage/sessions.ts
View file

@ -1,5 +1,5 @@
import { SessionError } from "@earendil-works/pi-agent-core";
import type { SqliteSessionMetadata } from "../types.ts";
import { assertJsonSerializable, SessionError } from "@earendil-works/pi-agent-core/experimental";
import type { SqliteDatabase, SqliteSessionMetadata } from "../types.ts";
export interface SessionRow {
id: string;
@ -7,7 +7,14 @@ export interface SessionRow {
metadata: string | null;
cwd: string;
parent_session_id: string | null;
active_leaf_id: string | null;
}
export interface NewSessionRow {
id: string;
createdAt: string;
cwd: string;
parentSessionId?: string;
metadata?: Record<string, unknown>;
}
function parseMetadata(metadata: string | null, sessionId: string): Record<string, unknown> | undefined {
@ -17,17 +24,62 @@ function parseMetadata(metadata: string | null, sessionId: string): Record<strin
parsed = JSON.parse(metadata);
} catch (error) {
throw new SessionError(
"invalid_session",
"storage",
`Invalid SQLite session ${sessionId}: metadata is not valid JSON`,
error instanceof Error ? error : undefined,
);
}
if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) {
throw new SessionError("invalid_session", `Invalid SQLite session ${sessionId}: metadata must be an object`);
throw new SessionError("storage", `Invalid SQLite session ${sessionId}: metadata must be an object`);
}
return parsed as Record<string, unknown>;
}
export function sessionExists(db: SqliteDatabase, sessionId: string) {
return !!db.prepare("SELECT 1 AS found FROM sessions WHERE id = ?").get<{ found: number }>(sessionId);
}
function serializeMetadata(metadata: Record<string, unknown> | undefined): string | null {
if (metadata === undefined) return null;
if (typeof metadata !== "object" || metadata === null || Array.isArray(metadata)) {
throw new SessionError("invalid_payload", "SQLite session metadata must be an object");
}
assertJsonSerializable(metadata);
return JSON.stringify(metadata);
}
export function insertSessionRow(db: SqliteDatabase, session: NewSessionRow) {
db.prepare("INSERT INTO sessions (id, created_at, metadata, cwd, parent_session_id) VALUES (?, ?, ?, ?, ?)").run(
session.id,
session.createdAt,
serializeMetadata(session.metadata),
session.cwd,
session.parentSessionId ?? null,
);
}
export function readSessionRow(db: SqliteDatabase, sessionId: string) {
return db
.prepare("SELECT id, created_at, metadata, cwd, parent_session_id FROM sessions WHERE id = ?")
.get<SessionRow>(sessionId);
}
export function readSessionRows(db: SqliteDatabase, options: { cwd?: string } = {}) {
return options.cwd
? db
.prepare(
"SELECT id, created_at, metadata, cwd, parent_session_id FROM sessions WHERE cwd = ? ORDER BY created_at DESC",
)
.all<SessionRow>(options.cwd)
: db
.prepare("SELECT id, created_at, metadata, cwd, parent_session_id FROM sessions ORDER BY created_at DESC")
.all<SessionRow>();
}
export function deleteSessionRow(db: SqliteDatabase, sessionId: string) {
db.prepare("DELETE FROM sessions WHERE id = ?").run(sessionId);
}
export function rowToMetadata(row: SessionRow, path: string): SqliteSessionMetadata {
return {
id: row.id,

17
packages/storage/sqlite-node/src/sqlite/storage/shared.ts
View file

@ -1,17 +0,0 @@
import type { SessionTreeEntry } from "@earendil-works/pi-agent-core";
import { SessionError } from "@earendil-works/pi-agent-core";
export function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === "object" && value !== null;
}
export function invalidSession(message: string, cause?: Error): SessionError {
return new SessionError("invalid_session", `Invalid SQLite session: ${message}`, cause);
}
export function invalidEntry(message: string, cause?: Error): SessionError {
return new SessionError("invalid_entry", `Invalid SQLite session entry: ${message}`, cause);
}
export function leafIdAfterEntry(entry: SessionTreeEntry): string | null {
return entry.type === "leaf" ? entry.targetId : entry.id;
}

13
packages/storage/sqlite-node/src/sqlite/types.ts
View file

@ -10,17 +10,18 @@ export interface SqliteRunResult {
/** Prepared SQLite statement capability used by the SQLite session backend. */
export interface SqliteStatement {
run(...params: unknown[]): Promise<SqliteRunResult>;
get<TRow extends object>(...params: unknown[]): Promise<TRow | undefined>;
all<TRow extends object>(...params: unknown[]): Promise<TRow[]>;
run(...params: unknown[]): SqliteRunResult;
get<TRow extends object>(...params: unknown[]): TRow | undefined;
all<TRow extends object>(...params: unknown[]): TRow[];
}
/** SQLite database capability used by the SQLite session backend. */
export interface SqliteDatabase {
exec(sql: string): Promise<void>;
exec(sql: string): void;
prepare(sql: string): SqliteStatement;
transaction<T>(fn: () => Promise<T>): Promise<T>;
close(): Promise<void>;
/** Runs a synchronous write transaction. The callback must not return a promise. */
transaction<T>(fn: () => T): T;
close(): void;
}
export interface SqliteDatabaseFactory {

1
tsconfig.json
View file

@ -11,6 +11,7 @@
"@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/experimental": ["./packages/agent/src/experimental.ts"],
"@earendil-works/pi-agent-core/*": ["./packages/agent/src/*"],
"@earendil-works/pi-agent-sqlite-node": ["./packages/storage/sqlite-node/src/index.ts"],
"@earendil-works/pi-coding-agent": ["./packages/coding-agent/src/index.ts"],