From 9e7582aa03e54f410fa9688197a3b64514e93400 Mon Sep 17 00:00:00 2001 From: Cristina Poncela Cubeiro <140309543+cristinaponcela@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:36:31 +0200 Subject: [PATCH] feat: sqlite session storage (#6594) This PR: - Adds retainedTail to compaction entries in the new agent harness so we don't have to walk up the tree for the 2000 tokens before compaction, - Changes getPathToRoot to getPathToRootOrCompaction to only load until last compaction, as unnecessary to access all nodes where it is called, - Adds a SQLite storage backend, in a separate packages/session-backend-sqlite, with a migration system and schemas as per on-site discussions: sessions to match session header messages (except for metadata, which I couldn't understand what it's used for or where it gets written, so I omitted it), session_entries for shared entry types as columns plus payload as a json for what remains, session_sequences to represent the append-only, serialized nature of the jsonl files, branch_entries to attribute nodes to branches (relationship one-to-many), and session_materialized with the session info (see /session in TUI) to act as a "cache" or quick-access for costs, message count, token info, labels, session name, and model-thinking-level config (e.g. for fast resume). - This is compatible with the new agent harness Session abstraction. --- biome.json | 2 + package-lock.json | 96 ++-- package.json | 5 +- packages/agent/README.md | 4 + packages/agent/package.json | 3 +- packages/agent/src/harness/agent-harness.ts | 1 + .../src/harness/compaction/compaction.ts | 20 +- .../src/harness/session/jsonl-storage.ts | 70 ++- .../src/harness/session/memory-storage.ts | 62 ++- .../agent/src/harness/session/repo-utils.ts | 2 +- packages/agent/src/harness/session/session.ts | 42 +- packages/agent/src/harness/types.ts | 26 +- packages/agent/src/index.ts | 1 + .../agent/test/harness/compaction.test.ts | 37 +- packages/agent/test/harness/session.test.ts | 11 +- .../test/harness/sqlite-migrations.test.ts | 441 +++++++++++++++++ .../agent/test/harness/sqlite-node.test.ts | 21 + packages/agent/test/harness/storage.test.ts | 175 ++++++- packages/agent/vitest.harness.config.ts | 2 + packages/coding-agent/docs/session-format.md | 17 +- .../install-lock/package-lock.json | 21 +- packages/coding-agent/npm-shrinkwrap.json | 21 +- packages/storage/sqlite-node/README.md | 5 + packages/storage/sqlite-node/package.json | 37 ++ .../sqlite-node/scripts/prepare-dist.mjs | 35 ++ packages/storage/sqlite-node/src/index.ts | 97 ++++ .../storage/sqlite-node/src/sqlite/index.ts | 4 + .../sqlite-node/src/sqlite/migrations.ts | 50 ++ .../src/sqlite/migrations/001_initial.sql | 59 +++ .../storage/sqlite-node/src/sqlite/repo.ts | 192 ++++++++ .../src/sqlite/storage/branch-entries.ts | 57 +++ .../sqlite-node/src/sqlite/storage/index.ts | 449 ++++++++++++++++++ .../src/sqlite/storage/session-entries.ts | 217 +++++++++ .../sqlite/storage/session-materialized.ts | 368 ++++++++++++++ .../src/sqlite/storage/session-sequences.ts | 16 + .../src/sqlite/storage/sessions.ts | 40 ++ .../sqlite-node/src/sqlite/storage/shared.ts | 29 ++ .../storage/sqlite-node/src/sqlite/types.ts | 55 +++ .../storage/sqlite-node/tsconfig.build.json | 11 + tsconfig.json | 3 +- 40 files changed, 2659 insertions(+), 145 deletions(-) create mode 100644 packages/agent/test/harness/sqlite-migrations.test.ts create mode 100644 packages/agent/test/harness/sqlite-node.test.ts create mode 100644 packages/storage/sqlite-node/README.md create mode 100644 packages/storage/sqlite-node/package.json create mode 100644 packages/storage/sqlite-node/scripts/prepare-dist.mjs create mode 100644 packages/storage/sqlite-node/src/index.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/index.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/migrations.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql create mode 100644 packages/storage/sqlite-node/src/sqlite/repo.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/index.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/session-entries.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/session-sequences.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/sessions.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/shared.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/types.ts create mode 100644 packages/storage/sqlite-node/tsconfig.build.json diff --git a/biome.json b/biome.json index b5451402..ef3cec42 100644 --- a/biome.json +++ b/biome.json @@ -27,6 +27,8 @@ "includes": [ "packages/*/src/**/*.ts", "packages/*/test/**/*.ts", + "packages/storage/*/src/**/*.ts", + "packages/storage/*/test/**/*.ts", "packages/coding-agent/examples/**/*.ts", "!**/node_modules/**/*", "!**/test-sessions.ts", diff --git a/package-lock.json b/package-lock.json index 1190eb57..2799f6d7 100644 --- a/package-lock.json +++ b/package-lock.json @@ -9,6 +9,7 @@ "version": "0.0.3", "workspaces": [ "packages/*", + "packages/storage/*", "packages/coding-agent/examples/extensions/with-deps", "packages/coding-agent/examples/extensions/custom-provider-anthropic", "packages/coding-agent/examples/extensions/custom-provider-gitlab-duo", @@ -785,6 +786,10 @@ "resolved": "packages/agent", "link": true }, + "node_modules/@earendil-works/pi-agent-sqlite-node": { + "resolved": "packages/storage/sqlite-node", + "link": true + }, "node_modules/@earendil-works/pi-ai": { "resolved": "packages/ai", "link": true @@ -801,29 +806,6 @@ "resolved": "packages/tui", "link": true }, - "node_modules/@emnapi/core": { - "version": "1.10.0", - "resolved": "https://registry.npmjs.org/@emnapi/core/-/core-1.10.0.tgz", - "integrity": "sha512-yq6OkJ4p82CAfPl0u9mQebQHKPJkY7WrIuk205cTYnYe+k2Z8YBh11FrbRG/H6ihirqcacOgl2BIO8oyMQLeXw==", - "dev": true, - "license": "MIT", - "optional": true, - "dependencies": { - "@emnapi/wasi-threads": "1.2.1", - "tslib": "^2.4.0" - } - }, - "node_modules/@emnapi/runtime": { - "version": "1.10.0", - "resolved": "https://registry.npmjs.org/@emnapi/runtime/-/runtime-1.10.0.tgz", - "integrity": "sha512-ewvYlk86xUoGI0zQRNq/mC+16R1QeDlKQy21Ki3oSYXNgLb45GV1P6A0M+/s6nyCuNDqe5VpaY84BzXGwVbwFA==", - "dev": true, - "license": "MIT", - "optional": true, - "dependencies": { - "tslib": "^2.4.0" - } - }, "node_modules/@emnapi/wasi-threads": { "version": "1.2.1", "resolved": "https://registry.npmjs.org/@emnapi/wasi-threads/-/wasi-threads-1.2.1.tgz", @@ -1403,9 +1385,6 @@ "cpu": [ "arm64" ], - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1422,9 +1401,6 @@ "cpu": [ "arm64" ], - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -1441,9 +1417,6 @@ "cpu": [ "riscv64" ], - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1460,9 +1433,6 @@ "cpu": [ "x64" ], - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1479,9 +1449,6 @@ "cpu": [ "x64" ], - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -1617,6 +1584,7 @@ "resolved": "https://registry.npmjs.org/@opentelemetry/api/-/api-1.9.0.tgz", "integrity": "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg==", "license": "Apache-2.0", + "peer": true, "engines": { "node": ">=8.0.0" } @@ -1796,9 +1764,6 @@ "arm64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1816,9 +1781,6 @@ "arm64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -1836,9 +1798,6 @@ "ppc64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1856,9 +1815,6 @@ "s390x" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1876,9 +1832,6 @@ "x64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1896,9 +1849,6 @@ "x64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -2712,6 +2662,7 @@ "dev": true, "hasInstallScript": true, "license": "MIT", + "peer": true, "bin": { "esbuild": "bin/esbuild" }, @@ -3607,9 +3558,6 @@ "arm64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -3631,9 +3579,6 @@ "arm64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -3655,9 +3600,6 @@ "x64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -3679,9 +3621,6 @@ "x64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -4793,6 +4732,7 @@ "integrity": "sha512-QP88BAKvMam/3NxH6vj2o21R6MjxZUAd6nlwAS/pnGvN9IVLocLHxGYIzFhg6fUQ+5th6P4dv4eW9jX3DSIj7A==", "dev": true, "license": "MIT", + "peer": true, "engines": { "node": ">=12" }, @@ -4831,6 +4771,7 @@ "integrity": "sha512-TvncJykhxAzFCk0VQZKBTClall4Pm7qXDSodb6uxi8QFa8X8mT6ABjxxsQ2opDRYxG7AzcRWXaFtruz5HJKuWg==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "esbuild": "~0.28.0" }, @@ -4911,6 +4852,7 @@ "integrity": "sha512-h9bXPmJichP5fLmVQo3PyaGSDE2n3aPuomeAlVRm0JLmt4rY6zmPKd59HYI4LNW8oTK7tlTsuC7l/m7awx9Jcw==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "lightningcss": "^1.32.0", "picomatch": "^4.0.4", @@ -5100,6 +5042,7 @@ "resolved": "https://registry.npmjs.org/zod/-/zod-3.25.76.tgz", "integrity": "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==", "license": "MIT", + "peer": true, "funding": { "url": "https://github.com/sponsors/colinhacks" } @@ -5139,6 +5082,7 @@ "integrity": "sha512-GUUEShf+PBCGW2KaXwcIt3Yk+e3pkKwWKb9GSyM9WQVE+ep2jzmHdGsHzu4wgcZy5fN9FBdVzjpBQsYlpfpgLA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "undici-types": "~7.16.0" } @@ -5149,6 +5093,7 @@ "integrity": "sha512-G9/lgqibheLVBDRuya45EbsEXTYcWoSG+TLg7i2axuzx0Eq62eXn+aWXyaVdV5vKvFSWd6ywcX8hA7la9Pvu8g==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@bcoe/v8-coverage": "^1.0.2", "@vitest/utils": "4.1.9", @@ -5354,6 +5299,7 @@ "integrity": "sha512-nE3/LEyc0z87uHYLZebqCUOaJr2hdtuPp7BQ4BosVFnfltxgAvMG08NyrSGlPpOUWvR27c5flSmYFTNr78L9GQ==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@vitest/expect": "4.1.9", "@vitest/mocker": "4.1.9", @@ -5500,6 +5446,7 @@ "integrity": "sha512-GUUEShf+PBCGW2KaXwcIt3Yk+e3pkKwWKb9GSyM9WQVE+ep2jzmHdGsHzu4wgcZy5fN9FBdVzjpBQsYlpfpgLA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "undici-types": "~7.16.0" } @@ -5867,6 +5814,7 @@ "integrity": "sha512-GUUEShf+PBCGW2KaXwcIt3Yk+e3pkKwWKb9GSyM9WQVE+ep2jzmHdGsHzu4wgcZy5fN9FBdVzjpBQsYlpfpgLA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "undici-types": "~7.16.0" } @@ -6152,6 +6100,18 @@ "node": ">=22.19.0" } }, + "packages/storage/sqlite-node": { + "name": "@earendil-works/pi-agent-sqlite-node", + "version": "0.80.10", + "license": "MIT", + "dependencies": { + "@earendil-works/pi-agent-core": "^0.80.10", + "@earendil-works/pi-ai": "^0.80.10" + }, + "engines": { + "node": ">=22.19.0" + } + }, "packages/tui": { "name": "@earendil-works/pi-tui", "version": "0.80.10", diff --git a/package.json b/package.json index fc3dd5db..8cdca855 100644 --- a/package.json +++ b/package.json @@ -4,6 +4,7 @@ "type": "module", "workspaces": [ "packages/*", + "packages/storage/*", "packages/coding-agent/examples/extensions/with-deps", "packages/coding-agent/examples/extensions/custom-provider-anthropic", "packages/coding-agent/examples/extensions/custom-provider-gitlab-duo", @@ -12,8 +13,8 @@ ], "scripts": { "clean": "npm run clean --workspaces", - "build": "cd packages/tui && npm run build && cd ../ai && npm run build && cd ../agent && npm run build && cd ../coding-agent && npm run build && cd ../orchestrator && npm run build", - "build:offline": "cd packages/tui && npm run build && cd ../ai && npm run build:offline && cd ../agent && npm run build && cd ../coding-agent && npm run build && cd ../orchestrator && npm run build", + "build": "cd packages/tui && npm run build && cd ../ai && npm run build && cd ../agent && npm run build && cd ../storage/sqlite-node && npm run build && cd ../../coding-agent && npm run build && cd ../orchestrator && npm run build", + "build:offline": "cd packages/tui && npm run build && cd ../ai && npm run build:offline && cd ../agent && npm run build && cd ../storage/sqlite-node && npm run build && cd ../../coding-agent && npm run build && cd ../orchestrator && npm run build", "check": "biome check --write --error-on-warnings . && npm run check:pinned-deps && npm run check:ts-imports && npm run check:shrinkwrap && npm run check:install-lock:coding-agent && tsgo --noEmit && npm run check:browser-smoke", "check:browser-smoke": "node scripts/check-browser-smoke.mjs", "check:pinned-deps": "node scripts/check-pinned-deps.mjs", diff --git a/packages/agent/README.md b/packages/agent/README.md index 3c935f44..00dac621 100644 --- a/packages/agent/README.md +++ b/packages/agent/README.md @@ -8,6 +8,10 @@ Stateful agent with tool execution and event streaming. Built on `@earendil-work npm install @earendil-works/pi-agent-core ``` +### SQLite session backends + +The SQLite session backend and the `node:sqlite` adapter live in a separate package, `@earendil-works/pi-agent-sqlite-node`, so the core package does not pull in runtime builtins or native SQLite dependencies by default. The backend accepts a runtime-specific SQLite factory, allowing other storage backends to ship as their own packages in the future. + ## Quick Start ```typescript diff --git a/packages/agent/package.json b/packages/agent/package.json index 7334e3a6..d130df50 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -21,12 +21,11 @@ "README.md" ], "scripts": { - "clean": "shx rm -rf dist", "build": "tsgo -p tsconfig.build.json", "test": "vitest --run", "test:harness": "vitest --run --config vitest.harness.config.ts", "coverage:harness": "vitest --run --config vitest.harness.config.ts --coverage", - "prepublishOnly": "npm run clean && npm run build" + "prepublishOnly": "npm run build" }, "dependencies": { "@earendil-works/pi-ai": "^0.80.10", diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index 85af6b28..5ec7d574 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -730,6 +730,7 @@ export class AgentHarness< result.details, provided !== undefined, result.usage, + result.retainedTail, ); const entry = await this.session.getEntry(entryId); if (entry?.type === "compaction") { diff --git a/packages/agent/src/harness/compaction/compaction.ts b/packages/agent/src/harness/compaction/compaction.ts index 6c21a492..0dc82c7c 100644 --- a/packages/agent/src/harness/compaction/compaction.ts +++ b/packages/agent/src/harness/compaction/compaction.ts @@ -97,12 +97,14 @@ function getMessageFromEntryForCompaction(entry: SessionTreeEntry): AgentMessage export interface CompactionResult { /** Summary text that replaces compacted history in future context. */ summary: string; - /** Entry id where retained history starts. */ - firstKeptEntryId: string; + /** Entry id where retained history starts. Optional during Pi 2.0 transition. */ + firstKeptEntryId?: string; /** Estimated context tokens before compaction. */ tokensBefore: number; /** Usage from the LLM call(s) that generated this summary, if available. */ usage?: Usage; + /** Retained recent messages stored directly on the compaction entry. Optional during Pi 2.0 transition. */ + retainedTail?: AgentMessage[]; /** Optional implementation-specific details stored with the compaction entry. */ details?: T; } @@ -583,6 +585,8 @@ export interface CompactionPreparation { messagesToSummarize: AgentMessage[]; /** Prefix messages summarized separately when compaction splits a turn. */ turnPrefixMessages: AgentMessage[]; + /** Recent messages retained after compaction and stored on the compaction entry. */ + retainedTail: AgentMessage[]; /** Whether compaction splits a turn. */ isSplitTurn: boolean; /** Estimated context tokens before compaction. */ @@ -617,7 +621,9 @@ export function prepareCompaction( if (prevCompactionIndex >= 0) { const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry; previousSummary = prevCompaction.summary; - const firstKeptEntryIndex = pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId); + const firstKeptEntryIndex = prevCompaction.firstKeptEntryId + ? pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId) + : -1; boundaryStart = firstKeptEntryIndex >= 0 ? firstKeptEntryIndex : prevCompactionIndex + 1; } const boundaryEnd = pathEntries.length; @@ -644,6 +650,11 @@ export function prepareCompaction( if (msg) turnPrefixMessages.push(msg); } } + const retainedTail: AgentMessage[] = []; + for (let i = cutPoint.firstKeptEntryIndex; i < boundaryEnd; i++) { + const msg = getMessageFromEntryForCompaction(pathEntries[i]); + if (msg) retainedTail.push(msg); + } const fileOps = extractFileOperations(messagesToSummarize, pathEntries, prevCompactionIndex); if (cutPoint.isSplitTurn) { for (const msg of turnPrefixMessages) { @@ -655,6 +666,7 @@ export function prepareCompaction( firstKeptEntryId, messagesToSummarize, turnPrefixMessages, + retainedTail, isSplitTurn: cutPoint.isSplitTurn, tokensBefore, previousSummary, @@ -693,6 +705,7 @@ export async function compact( firstKeptEntryId, messagesToSummarize, turnPrefixMessages, + retainedTail, isSplitTurn, tokensBefore, previousSummary, @@ -762,6 +775,7 @@ export async function compact( firstKeptEntryId, tokensBefore, usage: summaryUsage, + retainedTail, details: { readFiles, modifiedFiles } as CompactionDetails, }); } diff --git a/packages/agent/src/harness/session/jsonl-storage.ts b/packages/agent/src/harness/session/jsonl-storage.ts index d57398d5..fea97fbd 100644 --- a/packages/agent/src/harness/session/jsonl-storage.ts +++ b/packages/agent/src/harness/session/jsonl-storage.ts @@ -1,5 +1,12 @@ import { uuidv7 } from "@earendil-works/pi-ai"; -import type { FileSystem, JsonlSessionMetadata, LeafEntry, SessionStorage, SessionTreeEntry } from "../types.ts"; +import type { + FileSystem, + JsonlSessionMetadata, + LeafEntry, + SessionEntryCursorOptions, + SessionStorage, + SessionTreeEntry, +} from "../types.ts"; import { SessionError, toError } from "../types.ts"; import { getFileSystemResultOrThrow } from "./repo-utils.ts"; @@ -293,13 +300,66 @@ export class JsonlSessionStorage implements SessionStorage return this.labelsById.get(id); } - async getPathToRoot(leafId: string | null): Promise { + async getSessionName(): Promise { + const entries = await this.findEntries("session_info"); + return entries[entries.length - 1]?.name?.trim() || undefined; + } + + async getSessionStats() { + let messageCount = 0; + let cachedTokens = 0; + let uncachedTokens = 0; + let totalTokens = 0; + let costTotal = 0; + for (const entry of this.entries) { + if (entry.type === "message") { + messageCount += 1; + } + const usage = + entry.type === "message" + ? entry.message.role === "assistant" + ? entry.message.usage + : undefined + : entry.type === "compaction" || entry.type === "branch_summary" + ? entry.usage + : undefined; + if ( + !usage || + typeof usage.input !== "number" || + typeof usage.output !== "number" || + typeof usage.cacheRead !== "number" || + typeof usage.cacheWrite !== "number" || + typeof usage.cost?.total !== "number" + ) { + continue; + } + cachedTokens += usage.cacheRead; + uncachedTokens += usage.input + usage.cacheWrite; + totalTokens += usage.input + usage.output + usage.cacheRead + usage.cacheWrite; + costTotal += usage.cost.total; + } + return { + messageCount, + cachedTokens, + uncachedTokens, + totalTokens, + costTotal, + }; + } + + async getPathToRootOrCompaction(leafId: string | null): Promise { if (leafId === null) return []; const path: SessionTreeEntry[] = []; + let stopAtEntryId: string | null = null; let current = this.byId.get(leafId); if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`); while (current) { path.unshift(current); + if (stopAtEntryId !== null && current.id === stopAtEntryId) break; + if (current.type === "compaction") { + if (current.retainedTail) break; + stopAtEntryId = current.firstKeptEntryId ?? null; + } if (!current.parentId) break; const parent = this.byId.get(current.parentId); if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); @@ -308,7 +368,9 @@ export class JsonlSessionStorage implements SessionStorage return path; } - async getEntries(): Promise { - return [...this.entries]; + async getEntries(options?: SessionEntryCursorOptions): Promise { + const start = options?.afterEntrySeq ?? 0; + const end = options?.limit === undefined ? undefined : start + options.limit; + return this.entries.slice(start, end); } } diff --git a/packages/agent/src/harness/session/memory-storage.ts b/packages/agent/src/harness/session/memory-storage.ts index 68ec7623..1bc24f04 100644 --- a/packages/agent/src/harness/session/memory-storage.ts +++ b/packages/agent/src/harness/session/memory-storage.ts @@ -1,6 +1,7 @@ import { uuidv7 } from "@earendil-works/pi-ai"; import { type LeafEntry, + type SessionEntryCursorOptions, SessionError, type SessionMetadata, type SessionStorage, @@ -112,13 +113,66 @@ export class InMemorySessionStorage { + async getSessionName(): Promise { + const entries = await this.findEntries("session_info"); + return entries[entries.length - 1]?.name?.trim() || undefined; + } + + async getSessionStats() { + let messageCount = 0; + let cachedTokens = 0; + let uncachedTokens = 0; + let totalTokens = 0; + let costTotal = 0; + for (const entry of this.entries) { + if (entry.type === "message") { + messageCount += 1; + } + const usage = + entry.type === "message" + ? entry.message.role === "assistant" + ? entry.message.usage + : undefined + : entry.type === "compaction" || entry.type === "branch_summary" + ? entry.usage + : undefined; + if ( + !usage || + typeof usage.input !== "number" || + typeof usage.output !== "number" || + typeof usage.cacheRead !== "number" || + typeof usage.cacheWrite !== "number" || + typeof usage.cost?.total !== "number" + ) { + continue; + } + cachedTokens += usage.cacheRead; + uncachedTokens += usage.input + usage.cacheWrite; + totalTokens += usage.input + usage.output + usage.cacheRead + usage.cacheWrite; + costTotal += usage.cost.total; + } + return { + messageCount, + cachedTokens, + uncachedTokens, + totalTokens, + costTotal, + }; + } + + async getPathToRootOrCompaction(leafId: string | null): Promise { if (leafId === null) return []; const path: SessionTreeEntry[] = []; + let stopAtEntryId: string | null = null; let current = this.byId.get(leafId); if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`); while (current) { path.unshift(current); + if (stopAtEntryId !== null && current.id === stopAtEntryId) break; + if (current.type === "compaction") { + if (current.retainedTail) break; + stopAtEntryId = current.firstKeptEntryId ?? null; + } if (!current.parentId) break; const parent = this.byId.get(current.parentId); if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); @@ -127,7 +181,9 @@ export class InMemorySessionStorage { - return [...this.entries]; + async getEntries(options?: SessionEntryCursorOptions): Promise { + const start = options?.afterEntrySeq ?? 0; + const end = options?.limit === undefined ? undefined : start + options.limit; + return this.entries.slice(start, end); } } diff --git a/packages/agent/src/harness/session/repo-utils.ts b/packages/agent/src/harness/session/repo-utils.ts index dd2e0f9a..fd6ee4f2 100644 --- a/packages/agent/src/harness/session/repo-utils.ts +++ b/packages/agent/src/harness/session/repo-utils.ts @@ -47,5 +47,5 @@ export async function getEntriesToFork( } effectiveLeafId = target.parentId; } - return storage.getPathToRoot(effectiveLeafId); + return storage.getPathToRootOrCompaction(effectiveLeafId); } diff --git a/packages/agent/src/harness/session/session.ts b/packages/agent/src/harness/session/session.ts index bff41fdd..13007bc6 100644 --- a/packages/agent/src/harness/session/session.ts +++ b/packages/agent/src/harness/session/session.ts @@ -11,8 +11,10 @@ import type { MessageEntry, ModelChangeEntry, SessionContext, + SessionEntryCursorOptions, SessionInfoEntry, SessionMetadata, + SessionStats, SessionStorage, SessionTreeEntry, ThinkingLevelChangeEntry, @@ -67,11 +69,19 @@ export function defaultContextEntryTransform(pathEntries: readonly SessionTreeEn const entries: SessionTreeEntry[] = [compaction]; const compactionIdx = pathEntries.findIndex((entry) => entry.type === "compaction" && entry.id === compaction.id); - let foundFirstKept = false; - for (let i = 0; i < compactionIdx; i++) { - const entry = pathEntries[i]!; - if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true; - if (foundFirstKept) entries.push(entry); + if (compaction.retainedTail) { + for (let i = compactionIdx + 1; i < pathEntries.length; i++) { + entries.push(pathEntries[i]!); + } + return entries; + } + if (compaction.firstKeptEntryId) { + let foundFirstKept = false; + for (let i = 0; i < compactionIdx; i++) { + const entry = pathEntries[i]!; + if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true; + if (foundFirstKept) entries.push(entry); + } } for (let i = compactionIdx + 1; i < pathEntries.length; i++) { entries.push(pathEntries[i]!); @@ -111,7 +121,10 @@ export function sessionEntryToContextMessages( ]; } if (entry.type === "compaction") { - return [createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp)]; + return [ + createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp), + ...(entry.retainedTail ?? []), + ]; } if (entry.type === "branch_summary" && entry.summary) { return [createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp)]; @@ -159,13 +172,13 @@ export class Session { return this.storage.getEntry(id); } - getEntries(): Promise { - return this.storage.getEntries(); + getEntries(options?: SessionEntryCursorOptions): Promise { + return this.storage.getEntries(options); } async getBranch(fromId?: string): Promise { const leafId = fromId ?? (await this.storage.getLeafId()); - return this.storage.getPathToRoot(leafId); + return this.storage.getPathToRootOrCompaction(leafId); } async buildContextEntries(options: SessionContextBuildOptions = {}): Promise { @@ -190,9 +203,12 @@ export class Session { return this.storage.getLabel(id); } + getSessionStats(): Promise { + return this.storage.getSessionStats(); + } + async getSessionName(): Promise { - const entries = await this.storage.findEntries("session_info"); - return entries[entries.length - 1]?.name?.trim() || undefined; + return this.storage.getSessionName(); } private async appendTypedEntry(entry: TEntry): Promise { @@ -243,11 +259,12 @@ export class Session { async appendCompaction( summary: string, - firstKeptEntryId: string, + firstKeptEntryId: string | undefined, tokensBefore: number, details?: T, fromHook?: boolean, usage?: Usage, + retainedTail?: AgentMessage[], ): Promise { return this.appendTypedEntry({ type: "compaction", @@ -257,6 +274,7 @@ export class Session { summary, firstKeptEntryId, tokensBefore, + retainedTail, details, usage, fromHook, diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index cdc7a542..4a2f31d9 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -370,8 +370,9 @@ export interface ActiveToolsChangeEntry extends SessionTreeEntryBase { export interface CompactionEntry extends SessionTreeEntryBase { type: "compaction"; summary: string; - firstKeptEntryId: string; + firstKeptEntryId?: string; tokensBefore: number; + retainedTail?: AgentMessage[]; details?: T; usage?: Usage; fromHook?: boolean; @@ -436,6 +437,14 @@ export interface SessionContext { activeToolNames: string[] | null; } +export interface SessionStats { + messageCount: number; + cachedTokens: number; + uncachedTokens: number; + totalTokens: number; + costTotal: number; +} + export interface SessionMetadata { id: string; createdAt: string; @@ -448,6 +457,11 @@ export interface JsonlSessionMetadata extends SessionMetadata { metadata?: Record; } +export interface SessionEntryCursorOptions { + afterEntrySeq?: number; + limit?: number; +} + export interface SessionStorage { getMetadata(): Promise; getLeafId(): Promise; @@ -460,8 +474,10 @@ export interface SessionStorage>>; getLabel(id: string): Promise; - getPathToRoot(leafId: string | null): Promise; - getEntries(): Promise; + getSessionName(): Promise; + getSessionStats(): Promise; + getPathToRootOrCompaction(leafId: string | null): Promise; + getEntries(options?: SessionEntryCursorOptions): Promise; } export type { Session } from "./session/session.ts"; @@ -753,10 +769,11 @@ export interface AbortResult { export interface CompactResult { summary: string; - firstKeptEntryId: string; + firstKeptEntryId?: string; tokensBefore: number; /** Usage from the LLM call(s) that generated this summary, if available. */ usage?: Usage; + retainedTail?: AgentMessage[]; details?: unknown; } @@ -776,6 +793,7 @@ export interface CompactionPreparation { firstKeptEntryId: string; messagesToSummarize: AgentMessage[]; turnPrefixMessages: AgentMessage[]; + retainedTail: AgentMessage[]; isSplitTurn: boolean; tokensBefore: number; previousSummary?: string; diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index b5ec7807..b98023b0 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -1,4 +1,5 @@ // Core Agent +export { uuidv7 } from "@earendil-works/pi-ai"; export * from "./agent.ts"; // Loop functions export * from "./agent-loop.ts"; diff --git a/packages/agent/test/harness/compaction.test.ts b/packages/agent/test/harness/compaction.test.ts index c8160307..e55be45a 100644 --- a/packages/agent/test/harness/compaction.test.ts +++ b/packages/agent/test/harness/compaction.test.ts @@ -91,6 +91,7 @@ function createCompactionEntry( summary: string, firstKeptEntryId: string, parentId: string | null = null, + retainedTail?: AgentMessage[], ): CompactionEntry { return { type: "compaction", @@ -100,6 +101,7 @@ function createCompactionEntry( summary, firstKeptEntryId, tokensBefore: 1234, + retainedTail, }; } @@ -343,12 +345,38 @@ describe("harness compaction", () => { const a1 = createMessageEntry(createAssistantMessage("a"), u1.id); const u2 = createMessageEntry(createUserMessage("2"), a1.id); const a2 = createMessageEntry(createAssistantMessage("b"), u2.id); - const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id, a2.id); + const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id, a2.id, [ + createUserMessage("2"), + createAssistantMessage("b"), + ]); const u3 = createMessageEntry(createUserMessage("3"), compaction.id); const a3 = createMessageEntry(createAssistantMessage("c"), u3.id); const loaded = buildSessionContext([u1, a1, u2, a2, compaction, u3, a3]); expect(loaded.messages).toHaveLength(5); expect(loaded.messages[0]?.role).toBe("compactionSummary"); + expect(loaded.messages.map((message) => message.role)).toEqual([ + "compactionSummary", + "user", + "assistant", + "user", + "assistant", + ]); + }); + + it("falls back to firstKeptEntryId when a compaction has no retained tail", () => { + const u1 = createMessageEntry(createUserMessage("1")); + const a1 = createMessageEntry(createAssistantMessage("a"), u1.id); + const u2 = createMessageEntry(createUserMessage("2"), a1.id); + const a2 = createMessageEntry(createAssistantMessage("b"), u2.id); + const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id, a2.id); + const u3 = createMessageEntry(createUserMessage("3"), compaction.id); + const loaded = buildSessionContext([u1, a1, u2, a2, compaction, u3]); + expect(loaded.messages.map((message) => message.role)).toEqual([ + "compactionSummary", + "user", + "assistant", + "user", + ]); }); it("tracks model and thinking level changes in built context", () => { @@ -374,6 +402,7 @@ describe("harness compaction", () => { expect(preparation).toBeDefined(); expect(preparation?.previousSummary).toBe("First summary"); expect(preparation?.firstKeptEntryId).toBeTruthy(); + expect(preparation?.retainedTail.length).toBeGreaterThan(0); expect(preparation?.tokensBefore).toBe(estimateContextTokens(buildSessionContext(pathEntries).messages).tokens); }); @@ -566,6 +595,7 @@ describe("harness compaction", () => { firstKeptEntryId: "entry-keep", messagesToSummarize: messages, turnPrefixMessages: messages, + retainedTail: messages, isSplitTurn: true, tokensBefore: 600000, fileOps: { read: new Set(), written: new Set(), edited: new Set() }, @@ -583,6 +613,7 @@ describe("harness compaction", () => { firstKeptEntryId: "entry-keep", messagesToSummarize: messages, turnPrefixMessages: [], + retainedTail: messages, isSplitTurn: false, tokensBefore: 100, fileOps: { read: new Set(), written: new Set(), edited: new Set() }, @@ -619,6 +650,7 @@ describe("harness compaction", () => { turnPrefixMessages: messages, isSplitTurn: true, tokensBefore: 100, + retainedTail: messages, fileOps: { read: new Set(), written: new Set(), edited: new Set() }, settings: { enabled: true, reserveTokens: 2000, keepRecentTokens: 20 }, }; @@ -642,6 +674,7 @@ describe("harness compaction", () => { firstKeptEntryId: "entry-keep", messagesToSummarize: [], turnPrefixMessages: messages, + retainedTail: messages, isSplitTurn: true, tokensBefore: 100, fileOps: { read: new Set(), written: new Set(), edited: new Set() }, @@ -659,6 +692,7 @@ describe("harness compaction", () => { firstKeptEntryId: "entry-keep", messagesToSummarize: [], turnPrefixMessages: messages, + retainedTail: messages, isSplitTurn: true, tokensBefore: 100, fileOps: { read: new Set(), written: new Set(), edited: new Set() }, @@ -697,6 +731,7 @@ describe("harness compaction", () => { expect(result.summary.length).toBeGreaterThan(0); expect(result.firstKeptEntryId).toBeTruthy(); expect(result.usage?.totalTokens).toBeGreaterThan(0); + expect(result.retainedTail?.length).toBeGreaterThan(0); expect(result.details).toBeDefined(); }); }); diff --git a/packages/agent/test/harness/session.test.ts b/packages/agent/test/harness/session.test.ts index 26af52e2..70adbb38 100644 --- a/packages/agent/test/harness/session.test.ts +++ b/packages/agent/test/harness/session.test.ts @@ -68,11 +68,20 @@ async function runSessionSuite( await session.appendMessage(createAssistantMessage("two")); const user2 = await session.appendMessage(createUserMessage("three")); await session.appendMessage(createAssistantMessage("four")); - await session.appendCompaction("summary", user2, 1234); + await session.appendCompaction("summary", user2, 1234, undefined, undefined, undefined, [ + createUserMessage("three"), + createAssistantMessage("four"), + ]); await session.appendMessage(createUserMessage("five")); const context = await session.buildContext(); expect(context.messages[0]?.role).toBe("compactionSummary"); expect(context.messages).toHaveLength(4); + expect(context.messages.map((message) => message.role)).toEqual([ + "compactionSummary", + "user", + "assistant", + "user", + ]); }); it("supports moving with branch summary entries in context", async () => { diff --git a/packages/agent/test/harness/sqlite-migrations.test.ts b/packages/agent/test/harness/sqlite-migrations.test.ts new file mode 100644 index 00000000..0cf74938 --- /dev/null +++ b/packages/agent/test/harness/sqlite-migrations.test.ts @@ -0,0 +1,441 @@ +import { mkdtempSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { describe, expect, it } from "vitest"; +import { + applyMigrations, + createNodeSqliteFactory, + type SqliteDatabase, + type SqliteDatabaseFactory, + type SqliteRunResult, + type SqliteSessionMetadata, + SqliteSessionRepo, + SqliteSessionStorage, + type SqliteStatement, +} from "../../../storage/sqlite-node/src/index.ts"; +import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; +import { createAssistantMessage, createUserMessage } from "./session-test-utils.ts"; + +function createTempDir(): string { + return mkdtempSync(join(tmpdir(), "pi-agent-sqlite-")); +} + +class ThrowingStatement implements SqliteStatement { + private readonly onRun: () => Promise; + + constructor(onRun: () => Promise) { + this.onRun = onRun; + } + + async run(..._params: unknown[]): Promise { + return this.onRun(); + } + + async get(..._params: unknown[]): Promise { + return undefined; + } + + async all(..._params: unknown[]): Promise { + return []; + } +} + +class CountingDatabase implements SqliteDatabase { + closeCount = 0; + private readonly statementFactory: (sql: string) => SqliteStatement; + + constructor(statementFactory: (sql: string) => SqliteStatement) { + this.statementFactory = statementFactory; + } + + async exec(_sql: string): Promise {} + + prepare(sql: string): SqliteStatement { + return this.statementFactory(sql); + } + + async transaction(fn: () => Promise): Promise { + return fn(); + } + + async close(): Promise { + this.closeCount += 1; + } +} + +describe("SQLite migrations", () => { + it("applies file-based migrations and records them", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const repo = new SqliteSessionRepo({ env, sqlite, databasePath }); + await repo.create({ cwd: root, id: "session-1" }); + + const db = await sqlite.open(databasePath); + try { + const rows = await db.prepare("SELECT id FROM migrations ORDER BY id").all<{ id: string }>(); + expect(rows.map((row) => row.id)).toEqual(["001_initial.sql"]); + const tables = await db + .prepare("SELECT name, sql FROM sqlite_master WHERE type = 'table' ORDER BY name") + .all<{ name: string; sql: string | null }>(); + expect(tables.map((row) => row.name)).toEqual( + expect.arrayContaining([ + "migrations", + "sessions", + "session_entries", + "session_sequences", + "branch_entries", + "session_materialized", + "entry_materialized", + ]), + ); + const sessionColumns = await db.prepare("PRAGMA table_info(sessions)").all<{ name: string }>(); + expect(sessionColumns.map((column) => column.name)).toContain("active_leaf_id"); + for (const tableName of [ + "sessions", + "session_sequences", + "branch_entries", + "session_materialized", + "entry_materialized", + ]) { + const table = tables.find((row) => row.name === tableName); + expect(table?.sql).toContain("WITHOUT ROWID"); + } + } finally { + await db.close(); + } + }); + + it("persists session metadata through create, list, open, and fork", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const repo = new SqliteSessionRepo({ env, sqlite: createNodeSqliteFactory(), databasePath }); + const source = await repo.create({ + cwd: root, + id: "session-1", + metadata: { profile: "reviewer" }, + }); + const sourceMetadata = await source.getMetadata(); + expect(sourceMetadata.metadata).toEqual({ profile: "reviewer" }); + expect((await repo.list({ cwd: root })).map((listed) => listed.metadata)).toEqual([{ profile: "reviewer" }]); + expect((await (await repo.open(sourceMetadata)).getMetadata()).metadata).toEqual({ profile: "reviewer" }); + const fork = await repo.fork(sourceMetadata, { cwd: root, id: "session-2" }); + expect((await fork.getMetadata()).metadata).toEqual({ profile: "reviewer" }); + const overridden = await repo.fork(sourceMetadata, { + cwd: root, + id: "session-3", + metadata: { profile: "writer" }, + }); + expect((await overridden.getMetadata()).metadata).toEqual({ profile: "writer" }); + }); + + it("materializes active leaf id in sessions transactionally", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const repo = new SqliteSessionRepo({ env, sqlite, databasePath }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const rootId = await session.appendMessage(createUserMessage("root")); + const childId = await session.appendMessage(createAssistantMessage("child")); + await session.getStorage().setLeafId(rootId); + + const db = await sqlite.open(databasePath); + try { + const row = await db + .prepare("SELECT active_leaf_id FROM sessions WHERE id = ?") + .get<{ active_leaf_id: string | null }>("session-1"); + expect(row?.active_leaf_id).toBe(rootId); + const latestBranchRow = await db + .prepare( + "SELECT branch_id, entry_id, entry_seq FROM branch_entries WHERE session_id = ? ORDER BY entry_seq DESC LIMIT 1", + ) + .get<{ branch_id: string; entry_id: string; entry_seq: number }>("session-1"); + const latestSessionEntry = await db + .prepare("SELECT id, type FROM session_entries WHERE session_id = ? ORDER BY entry_seq DESC LIMIT 1") + .get<{ id: string; type: string }>("session-1"); + expect(latestSessionEntry?.type).toBe("leaf"); + expect(latestBranchRow?.entry_id).toBe(latestSessionEntry?.id); + } finally { + await db.close(); + } + + const reopened = await repo.open(await session.getMetadata()); + expect(await reopened.getLeafId()).toBe(rootId); + expect(childId).not.toBe(rootId); + }); + + it("materializes a new branch when appending from a parent with an existing child", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const repo = new SqliteSessionRepo({ env, sqlite, databasePath }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const rootId = await session.appendMessage(createUserMessage("root")); + const firstChildId = await session.appendMessage(createAssistantMessage("first child")); + await session.getStorage().setLeafId(rootId); + const secondChildId = await session.appendMessage(createAssistantMessage("second child")); + + const db = await sqlite.open(databasePath); + try { + const branchRows = await db + .prepare( + "SELECT branch_id, entry_id, entry_seq FROM branch_entries WHERE session_id = ? ORDER BY branch_id, entry_seq", + ) + .all<{ branch_id: string; entry_id: string; entry_seq: number }>("session-1"); + const branchIds = [...new Set(branchRows.map((row) => row.branch_id))]; + expect(branchIds).toHaveLength(3); + expect(branchRows.filter((row) => row.entry_id === rootId)).toHaveLength(3); + expect(branchRows.filter((row) => row.entry_id === firstChildId)).toHaveLength(1); + expect(branchRows.filter((row) => row.entry_id === secondChildId)).toHaveLength(1); + } finally { + await db.close(); + } + }); + + it("reopens using branch materialization and session summary state", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const repo = new SqliteSessionRepo({ env, sqlite: createNodeSqliteFactory(), databasePath }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const rootId = await session.appendMessage(createUserMessage("root")); + await session.appendMessage(createAssistantMessage("first child")); + await session.appendSessionName(" Reopened Session "); + await session.getStorage().setLeafId(rootId); + await session.appendMessage(createAssistantMessage("branched child")); + + const reopened = await repo.open(await session.getMetadata()); + expect(await reopened.getSessionName()).toBe("Reopened Session"); + expect((await reopened.buildContext()).messages.map((message) => message.role)).toEqual(["user", "assistant"]); + expect((await reopened.buildContext()).messages.at(-1)).toMatchObject({ + content: [{ type: "text", text: "branched child" }], + }); + }); + + it("pages entries by entry_seq cursor", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const repo = new SqliteSessionRepo({ env, sqlite: createNodeSqliteFactory(), databasePath }); + const session = await repo.create({ cwd: root, id: "session-1" }); + await session.appendMessage(createUserMessage("one")); + await session.appendMessage(createAssistantMessage("two")); + await session.appendMessage(createUserMessage("three")); + + expect((await session.getEntries({ limit: 2 })).map((entry) => entry.type)).toEqual(["message", "message"]); + expect((await session.getEntries({ afterEntrySeq: 2, limit: 2 })).map((entry) => entry.type)).toEqual([ + "message", + "message", + ]); + }); + + it("closes the database when create fails after openDatabase succeeds", async () => { + const root = createTempDir(); + const db = new CountingDatabase((sql) => { + if (sql.startsWith("INSERT INTO sessions")) { + return new ThrowingStatement(async () => { + throw new Error("insert failed"); + }); + } + return new ThrowingStatement(async () => ({ changes: 1 })); + }); + const sqlite: SqliteDatabaseFactory = { + open: async () => db, + }; + const env = new NodeExecutionEnv({ cwd: root }); + const repo = new SqliteSessionRepo({ env, sqlite, databasePath: join(root, "sessions.sqlite") }); + + await expect(repo.create({ cwd: root, id: "session-1" })).rejects.toThrow("insert failed"); + expect(db.closeCount).toBe(1); + }); + + it("closes the database when open fails after openDatabase succeeds", async () => { + const root = createTempDir(); + const db = new CountingDatabase((sql) => { + if (sql.includes("FROM sessions WHERE id = ?")) { + return new ThrowingStatement(async () => ({ changes: 0 })); + } + return new ThrowingStatement(async () => ({ changes: 1 })); + }); + const sqlite: SqliteDatabaseFactory = { + open: async () => db, + }; + const env = new NodeExecutionEnv({ cwd: root }); + const repo = new SqliteSessionRepo({ env, sqlite, databasePath: join(root, "sessions.sqlite") }); + const metadata: SqliteSessionMetadata = { + id: "missing", + createdAt: new Date().toISOString(), + cwd: root, + path: join(root, "sessions.sqlite"), + }; + writeFileSync(metadata.path, ""); + + await expect(repo.open(metadata)).rejects.toThrow("Session not found: missing"); + expect(db.closeCount).toBe(1); + }); + + it("closes the source storage after fork reads its entries", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const repo = new SqliteSessionRepo({ env, sqlite: createNodeSqliteFactory(), databasePath }); + let cleanupCount = 0; + const sourceStorage = { + async getEntries() { + return []; + }, + async getPathToRootOrCompaction() { + return []; + }, + async cleanup() { + cleanupCount += 1; + }, + } as const; + const originalOpen = repo.open.bind(repo); + repo.open = async () => + ({ + getStorage() { + return sourceStorage; + }, + }) as never; + + try { + await repo.fork( + { + id: "session-1", + createdAt: new Date().toISOString(), + cwd: root, + path: databasePath, + }, + { cwd: root, id: "session-2" }, + ); + } finally { + repo.open = originalOpen; + } + + expect(cleanupCount).toBe(1); + }); + + it("restores in-memory state when appendEntry fails after mutating caches", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const sqlite = createNodeSqliteFactory(); + const db = await sqlite.open(databasePath); + await applyMigrations(db); + const storage = await SqliteSessionStorage.create(db, databasePath, { + cwd: root, + sessionId: "session-1", + }); + const originalPrepare = db.prepare.bind(db); + db.prepare = (sql: string) => { + if (sql.startsWith("UPDATE sessions SET active_leaf_id = ?")) { + return new ThrowingStatement(async () => { + throw new Error("active leaf update failed"); + }); + } + return originalPrepare(sql); + }; + + await expect( + storage.appendEntry({ + type: "message", + id: "root", + parentId: null, + timestamp: new Date().toISOString(), + message: createUserMessage("root"), + }), + ).rejects.toMatchObject({ code: "storage" }); + expect(await storage.getLeafId()).toBeNull(); + expect(await storage.getEntry("root")).toBeUndefined(); + expect(await storage.getEntries()).toEqual([]); + await db.close(); + }); + + it("materializes session summary fields transactionally", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const repo = new SqliteSessionRepo({ env, sqlite, databasePath }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const userId = await session.appendMessage(createUserMessage("one")); + await session.appendThinkingLevelChange("high"); + await session.appendModelChange("anthropic", "claude-sonnet-4-5"); + const assistant = { + ...createAssistantMessage("two"), + provider: "anthropic", + model: "claude-sonnet-4-5", + usage: { + input: 100, + output: 25, + cacheRead: 40, + cacheWrite: 10, + totalTokens: 175, + cost: { input: 0.1, output: 0.2, cacheRead: 0.03, cacheWrite: 0.04, total: 0.37 }, + }, + }; + await session.appendMessage(assistant); + await session.appendCompaction("summary", userId, 200, undefined, false, { + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + cost: { input: 0.01, output: 0.02, cacheRead: 0.03, cacheWrite: 0.04, total: 0.1 }, + }); + await session.moveTo(userId, { + summary: "branch summary", + usage: { + input: 5, + output: 6, + cacheRead: 7, + cacheWrite: 8, + totalTokens: 26, + cost: { input: 0.05, output: 0.06, cacheRead: 0.07, cacheWrite: 0.08, total: 0.26 }, + }, + }); + await session.appendSessionName(" My Session "); + await session.appendLabel(userId, "checkpoint"); + + const db = await sqlite.open(databasePath); + try { + const row = await db.prepare("SELECT session_id, payload FROM session_materialized WHERE session_id = ?").get<{ + session_id: string; + payload: string; + }>("session-1"); + expect(row).toBeDefined(); + expect(row?.session_id).toBe("session-1"); + expect(JSON.parse(row?.payload ?? "null")).toMatchObject({ + name: "My Session", + messageCount: 2, + cachedTokens: 50, + uncachedTokens: 128, + totalTokens: 211, + costTotal: 0.73, + currentModel: { provider: "anthropic", modelId: "claude-sonnet-4-5" }, + currentThinkingLevel: "high", + }); + const entryRows = await db + .prepare( + "SELECT session_id, entry_seq, type, payload FROM entry_materialized WHERE session_id = ? ORDER BY entry_seq, type", + ) + .all<{ + session_id: string; + entry_seq: number; + type: string; + payload: string; + }>("session-1"); + expect( + entryRows.some((entryRow) => entryRow.type === "label" && JSON.parse(entryRow.payload).targetId === userId), + ).toBe(true); + expect(entryRows.some((entryRow) => entryRow.type === "thinking")).toBe(false); + expect(entryRows.some((entryRow) => entryRow.type === "model")).toBe(false); + } finally { + await db.close(); + } + }); +}); diff --git a/packages/agent/test/harness/sqlite-node.test.ts b/packages/agent/test/harness/sqlite-node.test.ts new file mode 100644 index 00000000..a94d3785 --- /dev/null +++ b/packages/agent/test/harness/sqlite-node.test.ts @@ -0,0 +1,21 @@ +import { join } from "node:path"; +import { describe, expect, it } from "vitest"; +import { createNodeSqliteFactory } from "../../../storage/sqlite-node/src/index.ts"; +import { createTempDir } from "./session-test-utils.ts"; + +describe("sqlite-node adapter", () => { + it("supports node:sqlite-style named parameters", async () => { + const root = createTempDir(); + const databasePath = join(root, "adapter.sqlite"); + const sqlite = createNodeSqliteFactory(); + const db = await sqlite.open(databasePath); + try { + await db.exec("CREATE TABLE items (id INTEGER PRIMARY KEY, text TEXT NOT NULL)"); + await db.prepare("INSERT INTO items (id, text) VALUES ($id, $text)").run({ $id: 1, $text: "hello" }); + const row = await db.prepare("SELECT text FROM items WHERE id = $id").get<{ text: string }>({ $id: 1 }); + expect(row).toEqual({ text: "hello" }); + } finally { + await db.close(); + } + }); +}); diff --git a/packages/agent/test/harness/storage.test.ts b/packages/agent/test/harness/storage.test.ts index 07f07df2..50f4fc59 100644 --- a/packages/agent/test/harness/storage.test.ts +++ b/packages/agent/test/harness/storage.test.ts @@ -4,7 +4,13 @@ import { describe, expect, it } from "vitest"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; import { JsonlSessionStorage, loadJsonlSessionMetadata } from "../../src/harness/session/jsonl-storage.ts"; import { InMemorySessionStorage } from "../../src/harness/session/memory-storage.ts"; -import { type MessageEntry, ok, type SessionMetadata } from "../../src/harness/types.ts"; +import { + type BranchSummaryEntry, + type CompactionEntry, + type MessageEntry, + ok, + type SessionMetadata, +} from "../../src/harness/types.ts"; import { createAssistantMessage, createTempDir, createUserMessage } from "./session-test-utils.ts"; describe("InMemorySessionStorage", () => { @@ -80,7 +86,74 @@ describe("InMemorySessionStorage", () => { expect(await storage.getLabel("entry-1")).toBeUndefined(); }); - it("walks paths to root", async () => { + it("includes summary-entry usage in session stats", async () => { + const assistant: MessageEntry = { + type: "message", + id: "assistant", + parentId: null, + timestamp: "2026-01-01T00:00:00.000Z", + message: { + role: "assistant", + content: [{ type: "text", text: "reply" }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet-4-5", + usage: { + input: 10, + output: 20, + cacheRead: 30, + cacheWrite: 40, + totalTokens: 100, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }, + stopReason: "stop", + timestamp: 0, + }, + }; + const compaction: CompactionEntry = { + type: "compaction", + id: "compaction", + parentId: "assistant", + timestamp: "2026-01-01T00:00:01.000Z", + summary: "summary", + firstKeptEntryId: "assistant", + tokensBefore: 1234, + usage: { + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + cost: { input: 0.01, output: 0.02, cacheRead: 0.03, cacheWrite: 0.04, total: 0.1 }, + }, + }; + const branchSummary: BranchSummaryEntry = { + type: "branch_summary", + id: "branch-summary", + parentId: "compaction", + timestamp: "2026-01-01T00:00:02.000Z", + fromId: "assistant", + summary: "branch", + usage: { + input: 5, + output: 6, + cacheRead: 7, + cacheWrite: 8, + totalTokens: 26, + cost: { input: 0.05, output: 0.06, cacheRead: 0.07, cacheWrite: 0.08, total: 0.26 }, + }, + }; + const storage = new InMemorySessionStorage({ entries: [assistant, compaction, branchSummary] }); + expect(await storage.getSessionStats()).toEqual({ + messageCount: 1, + cachedTokens: 40, + uncachedTokens: 68, + totalTokens: 136, + costTotal: 1.36, + }); + }); + + it("walks paths to root or retained-tail compaction", async () => { const root: MessageEntry = { type: "message", id: "root", @@ -94,9 +167,29 @@ describe("InMemorySessionStorage", () => { parentId: "root", message: createAssistantMessage("child"), }; - const storage = new InMemorySessionStorage({ entries: [root, child] }); - expect((await storage.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]); - expect(await storage.getPathToRoot(null)).toEqual([]); + const compaction: CompactionEntry = { + type: "compaction", + id: "compaction", + parentId: "child", + timestamp: "2026-01-01T00:00:01.000Z", + summary: "summary", + firstKeptEntryId: "child", + tokensBefore: 1234, + retainedTail: [createAssistantMessage("child")], + }; + const afterCompaction: MessageEntry = { + ...root, + id: "after-compaction", + parentId: "compaction", + message: createUserMessage("after"), + }; + const storage = new InMemorySessionStorage({ entries: [root, child, compaction, afterCompaction] }); + expect((await storage.getPathToRootOrCompaction("child")).map((entry) => entry.id)).toEqual(["root", "child"]); + expect((await storage.getPathToRootOrCompaction("after-compaction")).map((entry) => entry.id)).toEqual([ + "compaction", + "after-compaction", + ]); + expect(await storage.getPathToRootOrCompaction(null)).toEqual([]); }); }); @@ -255,7 +348,7 @@ describe("JsonlSessionStorage", () => { const reloaded = await JsonlSessionStorage.open(env, filePath); expect(await reloaded.getLeafId()).toBe("root"); expect((await reloaded.getEntries()).at(-1)).toMatchObject({ type: "leaf", targetId: "root" }); - expect((await loaded.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]); + expect((await loaded.getPathToRootOrCompaction("child")).map((entry) => entry.id)).toEqual(["root", "child"]); }); it("finds entries by type", async () => { @@ -309,6 +402,76 @@ describe("JsonlSessionStorage", () => { expect(await loaded.getLabel("entry-1")).toBeUndefined(); }); + it("includes summary-entry usage in session stats", async () => { + const dir = createTempDir(); + const env = new NodeExecutionEnv({ cwd: dir }); + const filePath = join(dir, "session.jsonl"); + const storage = await JsonlSessionStorage.create(env, filePath, { cwd: dir, sessionId: "session-1" }); + await storage.appendEntry({ + type: "message", + id: "assistant", + parentId: null, + timestamp: "2026-01-01T00:00:00.000Z", + message: { + role: "assistant", + content: [{ type: "text", text: "reply" }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet-4-5", + usage: { + input: 10, + output: 20, + cacheRead: 30, + cacheWrite: 40, + totalTokens: 100, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }, + stopReason: "stop", + timestamp: 0, + }, + }); + await storage.appendEntry({ + type: "compaction", + id: "compaction", + parentId: "assistant", + timestamp: "2026-01-01T00:00:01.000Z", + summary: "summary", + firstKeptEntryId: "assistant", + tokensBefore: 1234, + usage: { + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + cost: { input: 0.01, output: 0.02, cacheRead: 0.03, cacheWrite: 0.04, total: 0.1 }, + }, + }); + await storage.appendEntry({ + type: "branch_summary", + id: "branch-summary", + parentId: "compaction", + timestamp: "2026-01-01T00:00:02.000Z", + fromId: "assistant", + summary: "branch", + usage: { + input: 5, + output: 6, + cacheRead: 7, + cacheWrite: 8, + totalTokens: 26, + cost: { input: 0.05, output: 0.06, cacheRead: 0.07, cacheWrite: 0.08, total: 0.26 }, + }, + }); + expect(await storage.getSessionStats()).toEqual({ + messageCount: 1, + cachedTokens: 40, + uncachedTokens: 68, + totalTokens: 136, + costTotal: 1.36, + }); + }); + it("reads session metadata through the line-reading filesystem operation", async () => { const dir = createTempDir(); const filePath = join(dir, "session.jsonl"); diff --git a/packages/agent/vitest.harness.config.ts b/packages/agent/vitest.harness.config.ts index 91c0d471..8045bb1a 100644 --- a/packages/agent/vitest.harness.config.ts +++ b/packages/agent/vitest.harness.config.ts @@ -3,6 +3,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)); export default defineConfig({ test: { @@ -21,6 +22,7 @@ export default defineConfig({ }, resolve: { alias: [ + { find: /^@earendil-works\/pi-agent-core$/, replacement: agentSrcIndex }, { find: /^@earendil-works\/pi-ai$/, replacement: aiSrcIndex }, { find: /^@earendil-works\/pi-ai\/compat$/, replacement: aiSrcCompat }, ], diff --git a/packages/coding-agent/docs/session-format.md b/packages/coding-agent/docs/session-format.md index b27b8b70..e2c49f1d 100644 --- a/packages/coding-agent/docs/session-format.md +++ b/packages/coding-agent/docs/session-format.md @@ -232,10 +232,18 @@ Created when context is compacted. Stores a summary of earlier messages. {"type":"compaction","id":"f6g7h8i9","parentId":"e5f6g7h8","timestamp":"2024-12-03T14:10:00.000Z","summary":"User discussed X, Y, Z...","firstKeptEntryId":"c3d4e5f6","tokensBefore":50000} ``` +Newer harness-generated compactions embed the retained post-compaction context directly on the entry, instead of `firstKeptEntryId`: + +```json +{"type":"compaction","id":"f6g7h8i9","parentId":"e5f6g7h8","timestamp":"2024-12-03T14:10:00.000Z","summary":"User discussed X, Y, Z...","tokensBefore":50000,"retainedTail":[{"role":"user","content":"latest request"},{"role":"assistant","content":[{"type":"text","text":"latest reply"}],"provider":"anthropic","model":"claude-sonnet-4-5","usage":{...},"stopReason":"stop"}]} +``` + Optional fields: - `usage`: LLM usage from generating the summary; included in session token and cost totals +- `retainedTail`: Materialized `AgentMessage[]` kept after compaction. This is optional only for backward compatibility with older sessions. Newer harness-generated compactions include it so we can rebuild context from this checkpoint without walking older entries before the compaction entry. - `details`: Implementation-specific data (e.g., `{ readFiles: string[], modifiedFiles: string[] }` for default, or custom data for extensions) - `fromHook`: `true` if generated by an extension, `false`/`undefined` if pi-generated (legacy field name) +- `firstKeptEntryId`: for compatibility with old entry format. ### BranchSummaryEntry @@ -314,8 +322,9 @@ Entries form a tree: 1. Collects all entries on the path 2. If a `CompactionEntry` is on the path: - Includes the compaction entry first - - Then entries from `firstKeptEntryId` to compaction - - Then entries after compaction + - If `retainedTail` is present, it acts as a self-contained checkpoint and entries after the compaction are included + - Otherwise entries from `firstKeptEntryId` to the compaction are included + - Then entries after compaction are included 3. Preserves non-message entries in the selected range so interactive mode can render them `buildSessionContext()` builds on that entry list to produce the message list for the LLM: @@ -323,11 +332,13 @@ Entries form a tree: 1. Extracts current model and thinking level settings from the full path 2. Converts selected entries to messages: - `message` -> stored `AgentMessage` - - `compaction` -> `compactionSummary` + - `compaction` -> `compactionSummary` plus `retainedTail` when present - `branch_summary` -> `branchSummary` - `custom_message` -> `CustomMessage` - `custom` -> no context message +This makes newer compactions act like self-contained checkpoints. `retainedTail` is optional only so older sessions that only store `firstKeptEntryId` continue to load correctly. + ## Parsing Example ```typescript diff --git a/packages/coding-agent/install-lock/package-lock.json b/packages/coding-agent/install-lock/package-lock.json index f66fd9b3..5bcebab8 100644 --- a/packages/coding-agent/install-lock/package-lock.json +++ b/packages/coding-agent/install-lock/package-lock.json @@ -638,9 +638,6 @@ "cpu": [ "arm64" ], - "libc": [ - "glibc" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-arm64-musl": { @@ -657,9 +654,6 @@ "cpu": [ "arm64" ], - "libc": [ - "musl" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-riscv64-gnu": { @@ -676,9 +670,6 @@ "cpu": [ "riscv64" ], - "libc": [ - "glibc" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-x64-gnu": { @@ -695,9 +686,6 @@ "cpu": [ "x64" ], - "libc": [ - "glibc" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-x64-musl": { @@ -714,9 +702,6 @@ "cpu": [ "x64" ], - "libc": [ - "musl" - ], "optional": true }, "node_modules/@mariozechner/clipboard-win32-arm64-msvc": { @@ -790,7 +775,8 @@ "license": "Apache-2.0", "engines": { "node": ">=8.0.0" - } + }, + "peer": true }, "node_modules/@opentelemetry/semantic-conventions": { "version": "1.41.1", @@ -1834,7 +1820,8 @@ "license": "MIT", "funding": { "url": "https://github.com/sponsors/colinhacks" - } + }, + "peer": true }, "node_modules/zod-to-json-schema": { "version": "3.25.2", diff --git a/packages/coding-agent/npm-shrinkwrap.json b/packages/coding-agent/npm-shrinkwrap.json index 96c93de8..af3f2726 100644 --- a/packages/coding-agent/npm-shrinkwrap.json +++ b/packages/coding-agent/npm-shrinkwrap.json @@ -628,9 +628,6 @@ "cpu": [ "arm64" ], - "libc": [ - "glibc" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-arm64-musl": { @@ -647,9 +644,6 @@ "cpu": [ "arm64" ], - "libc": [ - "musl" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-riscv64-gnu": { @@ -666,9 +660,6 @@ "cpu": [ "riscv64" ], - "libc": [ - "glibc" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-x64-gnu": { @@ -685,9 +676,6 @@ "cpu": [ "x64" ], - "libc": [ - "glibc" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-x64-musl": { @@ -704,9 +692,6 @@ "cpu": [ "x64" ], - "libc": [ - "musl" - ], "optional": true }, "node_modules/@mariozechner/clipboard-win32-arm64-msvc": { @@ -780,7 +765,8 @@ "license": "Apache-2.0", "engines": { "node": ">=8.0.0" - } + }, + "peer": true }, "node_modules/@opentelemetry/semantic-conventions": { "version": "1.41.1", @@ -1824,7 +1810,8 @@ "license": "MIT", "funding": { "url": "https://github.com/sponsors/colinhacks" - } + }, + "peer": true }, "node_modules/zod-to-json-schema": { "version": "3.25.2", diff --git a/packages/storage/sqlite-node/README.md b/packages/storage/sqlite-node/README.md new file mode 100644 index 00000000..84915738 --- /dev/null +++ b/packages/storage/sqlite-node/README.md @@ -0,0 +1,5 @@ +# @earendil-works/pi-agent-sqlite-node + +Node sqlite storage backend for `@earendil-works/pi-agent-core` sessions. Provides the +`node:sqlite` adapter (`SqliteDatabase` implementation) and the SQLite session +repo/storage implementation (`SqliteSessionRepo`, migrations, materialized views). diff --git a/packages/storage/sqlite-node/package.json b/packages/storage/sqlite-node/package.json new file mode 100644 index 00000000..66069bd1 --- /dev/null +++ b/packages/storage/sqlite-node/package.json @@ -0,0 +1,37 @@ +{ + "name": "@earendil-works/pi-agent-sqlite-node", + "version": "0.80.10", + "description": "Node sqlite storage backend for @earendil-works/pi-agent-core sessions", + "type": "module", + "main": "./dist/index.js", + "types": "./dist/index.d.ts", + "exports": { + ".": { + "types": "./dist/index.d.ts", + "import": "./dist/index.js" + } + }, + "files": [ + "dist", + "README.md" + ], + "scripts": { + "clean": "node -e \"import('node:fs/promises').then(fs=>fs.rm('dist',{recursive:true,force:true}))\"", + "build": "tsgo -p tsconfig.build.json && node ./scripts/prepare-dist.mjs copy-sqlite-migrations", + "prepublishOnly": "npm run clean && npm run build" + }, + "author": "Earendil Works", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/earendil-works/pi.git", + "directory": "packages/storage/sqlite-node" + }, + "engines": { + "node": ">=22.19.0" + }, + "dependencies": { + "@earendil-works/pi-ai": "^0.80.10", + "@earendil-works/pi-agent-core": "^0.80.10" + } +} diff --git a/packages/storage/sqlite-node/scripts/prepare-dist.mjs b/packages/storage/sqlite-node/scripts/prepare-dist.mjs new file mode 100644 index 00000000..52328367 --- /dev/null +++ b/packages/storage/sqlite-node/scripts/prepare-dist.mjs @@ -0,0 +1,35 @@ +#!/usr/bin/env node + +import { cp, mkdir, rm } from "node:fs/promises"; +import { dirname, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +const scriptDir = dirname(fileURLToPath(import.meta.url)); +const packageDir = resolve(scriptDir, ".."); +const distDir = resolve(packageDir, "dist"); +const migrationSourceDir = resolve(packageDir, "src/sqlite/migrations"); +const migrationDestDir = resolve(distDir, "sqlite/migrations"); + +async function clean() { + await rm(distDir, { recursive: true, force: true }); +} + +async function copySqliteMigrations() { + await mkdir(migrationDestDir, { recursive: true }); + await cp(migrationSourceDir, migrationDestDir, { recursive: true }); +} + +const command = process.argv[2]; + +if (command === "clean") { + await clean(); + process.exit(0); +} + +if (command === "copy-sqlite-migrations") { + await copySqliteMigrations(); + process.exit(0); +} + +console.error("Usage: node scripts/prepare-dist.mjs "); +process.exit(1); diff --git a/packages/storage/sqlite-node/src/index.ts b/packages/storage/sqlite-node/src/index.ts new file mode 100644 index 00000000..399eadb6 --- /dev/null +++ b/packages/storage/sqlite-node/src/index.ts @@ -0,0 +1,97 @@ +import type { SQLInputValue } from "node:sqlite"; +import { DatabaseSync } from "node:sqlite"; +import type { SqliteDatabase, SqliteDatabaseFactory, SqliteRunResult, SqliteStatement } from "./sqlite/types.ts"; + +function isNamedParameters(value: unknown): value is Record { + if (value === null || typeof value !== "object") return false; + if (Array.isArray(value) || ArrayBuffer.isView(value)) return false; + return true; +} + +class NodeSqliteStatement implements SqliteStatement { + private readonly statement: ReturnType; + + constructor(statement: ReturnType) { + this.statement = statement; + } + + async run(...params: unknown[]): Promise { + const [first, ...rest] = params; + const result = isNamedParameters(first) + ? this.statement.run(first, ...(rest as SQLInputValue[])) + : this.statement.run(...(params as SQLInputValue[])); + return { + changes: Number(result.changes), + lastInsertRowid: result.lastInsertRowid === undefined ? undefined : Number(result.lastInsertRowid), + }; + } + + async get(...params: unknown[]): Promise { + const [first, ...rest] = params; + return ( + isNamedParameters(first) + ? this.statement.get(first, ...(rest as SQLInputValue[])) + : this.statement.get(...(params as SQLInputValue[])) + ) as TRow | undefined; + } + + async all(...params: unknown[]): Promise { + const [first, ...rest] = params; + return ( + isNamedParameters(first) + ? this.statement.all(first, ...(rest as SQLInputValue[])) + : this.statement.all(...(params as SQLInputValue[])) + ) as TRow[]; + } +} + +class NodeSqliteDatabase implements SqliteDatabase { + private readonly db: DatabaseSync; + + constructor(db: DatabaseSync) { + this.db = db; + } + + async exec(sql: string): Promise { + this.db.exec(sql); + } + + prepare(sql: string): SqliteStatement { + return new NodeSqliteStatement(this.db.prepare(sql)); + } + + async transaction(fn: () => Promise): Promise { + this.db.exec("BEGIN"); + try { + const result = await fn(); + this.db.exec("COMMIT"); + return result; + } catch (error) { + try { + this.db.exec("ROLLBACK"); + } catch { + // Ignore rollback errors to rethrow original error. + } + throw error; + } + } + + async close(): Promise { + this.db.close(); + } +} + +export function wrapNodeSqliteDatabase(db: DatabaseSync): SqliteDatabase { + return new NodeSqliteDatabase(db); +} + +export function createNodeSqliteFactory(): SqliteDatabaseFactory { + return { + async open(path: string): Promise { + return new NodeSqliteDatabase(new DatabaseSync(path)); + }, + }; +} + +// Re-export the SQLite session storage backend and types so this package is a complete node-sqlite backend. +export * from "./sqlite/index.ts"; diff --git a/packages/storage/sqlite-node/src/sqlite/index.ts b/packages/storage/sqlite-node/src/sqlite/index.ts new file mode 100644 index 00000000..279a6083 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/index.ts @@ -0,0 +1,4 @@ +export * from "./migrations.ts"; +export * from "./repo.ts"; +export * from "./storage/index.ts"; +export * from "./types.ts"; diff --git a/packages/storage/sqlite-node/src/sqlite/migrations.ts b/packages/storage/sqlite-node/src/sqlite/migrations.ts new file mode 100644 index 00000000..2ccd70fc --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/migrations.ts @@ -0,0 +1,50 @@ +import { readFile } from "node:fs/promises"; +import { fileURLToPath } from "node:url"; +import type { SqliteDatabase } from "./types.ts"; + +export interface SqliteMigration { + id: string; + order: number; + sql: string; +} + +async function loadMigrationSql(relativePath: string): Promise { + return readFile(fileURLToPath(new URL(relativePath, import.meta.url)), "utf8"); +} + +export async function loadMigrations(): Promise { + return [ + { + id: "001_initial.sql", + order: 1, + sql: await loadMigrationSql("./migrations/001_initial.sql"), + }, + ]; +} + +async function ensureMigrationsTable(db: SqliteDatabase): Promise { + await db.exec(` +CREATE TABLE IF NOT EXISTS migrations ( + id TEXT PRIMARY KEY, + applied_at TEXT NOT NULL +); +`); +} + +export async function applyMigrations(db: SqliteDatabase): Promise { + await ensureMigrationsTable(db); + const migrations = await loadMigrations(); + const appliedRows = await 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()); + }); + applied.add(migration.id); + } +} diff --git a/packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql b/packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql new file mode 100644 index 00000000..6e6f397d --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql @@ -0,0 +1,59 @@ +CREATE TABLE IF NOT EXISTS sessions ( + id TEXT PRIMARY KEY, + created_at TEXT NOT NULL, + cwd TEXT NOT NULL, + parent_session_id TEXT NULL, + metadata TEXT NULL, + active_leaf_id 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 ( + session_id TEXT 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) +); + +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 TABLE IF NOT EXISTS session_sequences ( + session_id TEXT PRIMARY KEY, + next_seq INTEGER NOT NULL +) WITHOUT ROWID; + +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, + 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 TABLE IF NOT EXISTS session_materialized ( + session_id TEXT PRIMARY KEY, + payload TEXT NOT NULL +) WITHOUT ROWID; + +CREATE TABLE IF NOT EXISTS entry_materialized ( + 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) +) WITHOUT ROWID; + +CREATE INDEX IF NOT EXISTS idx_entry_materialized_session_type_seq ON entry_materialized(session_id, type, entry_seq); diff --git a/packages/storage/sqlite-node/src/sqlite/repo.ts b/packages/storage/sqlite-node/src/sqlite/repo.ts new file mode 100644 index 00000000..12a4f9d0 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/repo.ts @@ -0,0 +1,192 @@ +import type { Session, SessionStorage, SessionTreeEntry } from "@earendil-works/pi-agent-core"; +import { + createSessionId, + getEntriesToFork, + getFileSystemResultOrThrow, + SessionError, + toSession, +} from "@earendil-works/pi-agent-core"; +import { applyMigrations } from "./migrations.ts"; +import { SqliteSessionStorage } from "./storage/index.ts"; +import { rowToMetadata, type SessionRow } from "./storage/sessions.ts"; +import type { + SqliteDatabase, + SqliteDatabaseFactory, + SqliteSessionCreateOptions, + SqliteSessionListOptions, + SqliteSessionMetadata, + SqliteSessionRepoApi, + SqliteSessionRepoEnv, +} from "./types.ts"; + +function getParentPath(path: string): string { + const normalized = path.replace(/[\\/]+$/, ""); + const lastSlash = Math.max(normalized.lastIndexOf("/"), normalized.lastIndexOf("\\")); + if (lastSlash < 0) return "."; + if (lastSlash === 0) return normalized.slice(0, 1); + return normalized.slice(0, lastSlash); +} + +async function configureSqliteDatabase(db: SqliteDatabase): Promise { + await db.exec("PRAGMA journal_mode=WAL"); + await db.exec("PRAGMA synchronous=FULL"); + await db.exec("PRAGMA busy_timeout=5000"); +} + +async function cleanupSessionStorage(storage: SessionStorage): Promise { + const maybeClosable = storage as SessionStorage & { cleanup?: () => Promise }; + if (typeof maybeClosable.cleanup === "function") { + await maybeClosable.cleanup(); + } +} + +export class SqliteSessionRepo implements SqliteSessionRepoApi { + private readonly env: SqliteSessionRepoEnv; + private readonly sqlite: SqliteDatabaseFactory; + private readonly databasePathInput: string; + private databasePath: string | undefined; + + constructor(options: { env: SqliteSessionRepoEnv; sqlite: SqliteDatabaseFactory; databasePath: string }) { + this.env = options.env; + this.sqlite = options.sqlite; + this.databasePathInput = options.databasePath; + } + + private async getDatabasePath(): Promise { + if (!this.databasePath) { + this.databasePath = getFileSystemResultOrThrow( + await this.env.absolutePath(this.databasePathInput), + `Failed to resolve SQLite sessions database ${this.databasePathInput}`, + ); + } + return this.databasePath; + } + + private async ensureDatabaseDir(): Promise { + const path = await this.getDatabasePath(); + const directory = getParentPath(path); + getFileSystemResultOrThrow( + await this.env.createDir(directory, { recursive: true }), + `Failed to create SQLite sessions directory ${directory}`, + ); + } + + private async openDatabase(): Promise { + await this.ensureDatabaseDir(); + const db = await this.sqlite.open(await this.getDatabasePath()); + try { + await configureSqliteDatabase(db); + await applyMigrations(db); + return db; + } catch (error) { + await db.close(); + throw error; + } + } + + async create(options: SqliteSessionCreateOptions): Promise> { + const db = await this.openDatabase(); + try { + const id = options.id ?? createSessionId(); + const storage = await SqliteSessionStorage.create(db, await this.getDatabasePath(), { + cwd: options.cwd, + sessionId: id, + parentSessionId: options.parentSessionId, + metadata: options.metadata, + }); + return toSession(storage); + } catch (error) { + await db.close(); + throw error; + } + } + + async open(metadata: SqliteSessionMetadata): Promise> { + if ( + !getFileSystemResultOrThrow(await this.env.exists(metadata.path), `Failed to check database ${metadata.path}`) + ) { + throw new SessionError("not_found", `Session not found: ${metadata.id}`); + } + const db = await this.openDatabase(); + try { + const storage = await SqliteSessionStorage.open(db, metadata); + return toSession(storage); + } catch (error) { + await db.close(); + throw error; + } + } + + async list(options: SqliteSessionListOptions = {}): Promise { + const path = await this.getDatabasePath(); + if (!getFileSystemResultOrThrow(await this.env.exists(path), `Failed to check database ${path}`)) { + return []; + } + const db = await this.openDatabase(); + try { + const rows = options.cwd + ? await db + .prepare( + "SELECT id, created_at, metadata, cwd, parent_session_id, active_leaf_id FROM sessions WHERE cwd = ? ORDER BY created_at DESC", + ) + .all(options.cwd) + : await db + .prepare( + "SELECT id, created_at, metadata, cwd, parent_session_id, active_leaf_id FROM sessions ORDER BY created_at DESC", + ) + .all(); + return rows.map((row) => rowToMetadata(row, path)); + } finally { + await db.close(); + } + } + + async delete(metadata: SqliteSessionMetadata): Promise { + const db = await this.openDatabase(); + try { + await db.transaction(async () => { + await db.prepare("DELETE FROM branch_entries WHERE session_id = ?").run(metadata.id); + await db.prepare("DELETE FROM session_entries WHERE session_id = ?").run(metadata.id); + await db.prepare("DELETE FROM entry_materialized WHERE session_id = ?").run(metadata.id); + await db.prepare("DELETE FROM session_materialized WHERE session_id = ?").run(metadata.id); + await db.prepare("DELETE FROM session_sequences WHERE session_id = ?").run(metadata.id); + const result = await db.prepare("DELETE FROM sessions WHERE id = ?").run(metadata.id); + if (result.changes === 0) { + throw new SessionError("not_found", `Session not found: ${metadata.id}`); + } + }); + } finally { + await db.close(); + } + } + + async fork( + sourceMetadata: SqliteSessionMetadata, + options: SqliteSessionCreateOptions & { entryId?: string; position?: "before" | "at"; id?: string }, + ): Promise> { + const source = await this.open(sourceMetadata); + let forkedEntries: SessionTreeEntry[]; + try { + forkedEntries = await getEntriesToFork(source.getStorage(), options); + } finally { + await cleanupSessionStorage(source.getStorage()); + } + const db = await this.openDatabase(); + try { + const id = options.id ?? createSessionId(); + const storage = await SqliteSessionStorage.create(db, await this.getDatabasePath(), { + cwd: options.cwd, + sessionId: id, + parentSessionId: options.parentSessionId ?? sourceMetadata.id, + metadata: options.metadata ?? sourceMetadata.metadata, + }); + for (const entry of forkedEntries) { + await storage.appendEntry(entry); + } + return toSession(storage); + } catch (error) { + await db.close(); + throw error; + } + } +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts b/packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts new file mode 100644 index 00000000..0ea60156 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts @@ -0,0 +1,57 @@ +import type { SessionTreeEntry } from "@earendil-works/pi-agent-core"; +import type { SqliteDatabase } from "../types.ts"; +import { decodeEntry, type SessionEntryRow } from "./session-entries.ts"; +import { invalidSession } from "./shared.ts"; + +export interface BranchEntryRow { + entry_id: string; + entry_seq: number; +} + +export async function getMaterializedBranchPathOrCompaction( + db: SqliteDatabase, + sessionId: string, + branchId: string, + byId: Map, +): Promise { + const branchRows = await db + .prepare( + "SELECT entry_id, entry_seq FROM branch_entries WHERE session_id = ? AND branch_id = ? ORDER BY entry_seq", + ) + .all(sessionId, branchId); + if (branchRows.length === 0) { + return []; + } + const entryIds = branchRows.map((row) => row.entry_id); + const placeholders = entryIds.map(() => "?").join(", "); + const entryRows = await db + .prepare( + `SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND id IN (${placeholders})`, + ) + .all(sessionId, ...entryIds); + const entryRowsById = new Map(entryRows.map((row) => [row.id, row])); + const entries: SessionTreeEntry[] = []; + for (const branchRow of branchRows) { + // leaf entries are navigation markers used to mark which branch became active; + // they are not part of the model/context path reconstructed from branch_entries. + const cached = byId.get(branchRow.entry_id); + if (cached) { + if (cached.type !== "leaf") { + entries.push(cached); + } + continue; + } + const entryRow = entryRowsById.get(branchRow.entry_id); + if (!entryRow) throw invalidSession(`missing entry row for branch entry ${branchRow.entry_id}`); + try { + const entry = decodeEntry(entryRow); + byId.set(entry.id, entry); + if (entry.type !== "leaf") { + entries.push(entry); + } + } catch { + throw invalidSession(`invalid entry row for branch entry ${branchRow.entry_id}`); + } + } + return entries; +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/index.ts b/packages/storage/sqlite-node/src/sqlite/storage/index.ts new file mode 100644 index 00000000..93376852 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/index.ts @@ -0,0 +1,449 @@ +import type { + LeafEntry, + SessionEntryCursorOptions, + SessionStorage, + SessionTreeEntry, +} from "@earendil-works/pi-agent-core"; +import { SessionError } from "@earendil-works/pi-agent-core"; +import { uuidv7 } from "@earendil-works/pi-ai"; +import type { SqliteDatabase, SqliteSessionMetadata } from "../types.ts"; +import { getMaterializedBranchPathOrCompaction } from "./branch-entries.ts"; +import { decodeEntry, encodeEntry, type SessionEntryRow } from "./session-entries.ts"; +import { + applyEntryToMaterializedState, + createEmptyMaterializedState, + type EntryMaterializedRow, + entryMaterializedValues, + materializedStateFromRows, + materializedStateValues, + type SessionMaterializedRow, + type SessionMaterializedState, + serializeSummary, + sessionStatsFromMaterializedState, +} from "./session-materialized.ts"; +import { advanceSequence, getNextSequence } from "./session-sequences.ts"; +import { rowToMetadata, type SessionRow } from "./sessions.ts"; +import { generateEntryId, invalidSession, leafIdAfterEntry } from "./shared.ts"; + +async function decodeEntryRows(entryRows: SessionEntryRow[]): Promise<{ + entries: SessionTreeEntry[]; + leafId: string | null; +}> { + const entries: SessionTreeEntry[] = []; + let leafId: string | null = null; + for (const entryRow of entryRows) { + try { + const entry = decodeEntry(entryRow); + entries.push(entry); + leafId = leafIdAfterEntry(entry); + } catch { + // Keep JSONL-like permissive resume behavior: skip malformed entries. + } + } + return { entries, leafId }; +} + +async function loadEntryRowsByIds( + db: SqliteDatabase, + sessionId: string, + entryIds: string[], +): Promise> { + if (entryIds.length === 0) return new Map(); + const placeholders = entryIds.map(() => "?").join(", "); + const rows = await db + .prepare( + `SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND id IN (${placeholders})`, + ) + .all(sessionId, ...entryIds); + return new Map(rows.map((row) => [row.id, row])); +} + +async function loadActiveBranchId(db: SqliteDatabase, sessionId: string): Promise { + // branch_entries includes leaf navigation entries for the active branch, so the + // newest branch_entries row identifies the branch that was most recently made active. + const row = await db + .prepare( + "SELECT branch_id FROM branch_entries WHERE session_id = ? ORDER BY entry_seq DESC, branch_id DESC LIMIT 1", + ) + .get<{ branch_id: string }>(sessionId); + return row?.branch_id ?? null; +} + +async function hasExistingChild(db: SqliteDatabase, sessionId: string, parentId: string | null): Promise { + const row = + parentId === null + ? await db + .prepare("SELECT 1 AS found FROM session_entries WHERE session_id = ? AND parent_id IS NULL LIMIT 1") + .get<{ found: number }>(sessionId) + : await db + .prepare("SELECT 1 AS found FROM session_entries WHERE session_id = ? AND parent_id = ? LIMIT 1") + .get<{ found: number }>(sessionId, parentId); + return row !== undefined; +} + +async function loadSqliteStorage( + db: SqliteDatabase, + sessionId: string, +): Promise<{ + row: SessionRow; + leafId: string | null; + activeBranchId: string | null; + materializedState: SessionMaterializedState; +}> { + const row = await db + .prepare("SELECT id, created_at, metadata, cwd, parent_session_id, active_leaf_id FROM sessions WHERE id = ?") + .get(sessionId); + if (!row) throw new SessionError("not_found", `Session not found: ${sessionId}`); + + const leafId = row.active_leaf_id; + const materializedRow = await db + .prepare("SELECT session_id, payload FROM session_materialized WHERE session_id = ?") + .get(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(sessionId); + return { + row, + leafId, + activeBranchId: await loadActiveBranchId(db, sessionId), + materializedState: materializedStateFromRows(materializedRow, entryMaterializedRows), + }; +} + +export class SqliteSessionStorage implements SessionStorage { + private readonly db: SqliteDatabase; + private readonly metadata: SqliteSessionMetadata; + private byId: Map; + private labelsById: Map; + private currentLeafId: string | null; + private activeBranchId: string | null; + private materializedState: SessionMaterializedState; + + private async getPathToRootOrCompactionEntries(leafId: string | null): Promise { + if (leafId === null) return []; + const path: SessionTreeEntry[] = []; + let stopAtEntryId: string | null = null; + let current = await this.getEntry(leafId); + if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`); + while (current) { + path.unshift(current); + if (stopAtEntryId !== null && current.id === stopAtEntryId) break; + if (current.type === "compaction") { + if (current.retainedTail) break; + stopAtEntryId = current.firstKeptEntryId ?? null; + } + if (!current.parentId) break; + const parent = await this.getEntry(current.parentId); + if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); + current = parent; + } + return path; + } + + private async materializeBranch(leafId: string | null): Promise { + const branchId = uuidv7(); + // Rebuild the branch path only when branch membership changes: branch switch + // (leaf navigation) or a new fork from a parent that already has a child. + // Linear appends stay cheap and extend the active branch incrementally. + const path = await this.getPathToRootOrCompactionEntries(leafId); + const entryRowsById = await loadEntryRowsByIds( + this.db, + this.metadata.id, + path.map((entry) => entry.id), + ); + for (const entry of path) { + const entryRow = entryRowsById.get(entry.id); + if (!entryRow) throw invalidSession(`missing entry row for session ${this.metadata.id} entry ${entry.id}`); + await this.db + .prepare("INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) VALUES (?, ?, ?, ?)") + .run(this.metadata.id, branchId, entry.id, entryRow.entry_seq); + } + this.activeBranchId = branchId; + } + + private async appendToActiveBranch(entryId: string, parentId: string | null): Promise { + if (!this.activeBranchId) { + await this.materializeBranch(parentId); + } + // After a branch is materialized/resynced, subsequent linear appends only add the + // new tip entry. We do not rebuild the full branch on every append. + if (!this.activeBranchId) { + throw invalidSession(`active branch missing for session ${this.metadata.id}`); + } + const entryRow = await this.db + .prepare("SELECT entry_seq FROM session_entries WHERE session_id = ? AND id = ?") + .get<{ entry_seq: number }>(this.metadata.id, entryId); + if (!entryRow) throw invalidSession(`missing entry row for session ${this.metadata.id} entry ${entryId}`); + await this.db + .prepare("INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) VALUES (?, ?, ?, ?)") + .run(this.metadata.id, this.activeBranchId, entryId, entryRow.entry_seq); + } + + private constructor( + db: SqliteDatabase, + metadata: SqliteSessionMetadata, + entries: SessionTreeEntry[] | null, + leafId: string | null, + activeBranchId: string | null, + materializedState: SessionMaterializedState, + ) { + this.db = db; + this.metadata = metadata; + this.byId = new Map((entries ?? []).map((entry) => [entry.id, entry])); + this.materializedState = materializedState; + this.labelsById = materializedState.labelsById; + this.currentLeafId = leafId; + this.activeBranchId = activeBranchId; + } + + static async open(db: SqliteDatabase, metadata: SqliteSessionMetadata): Promise { + const loaded = await loadSqliteStorage(db, metadata.id); + return new SqliteSessionStorage( + db, + rowToMetadata(loaded.row, metadata.path), + null, + loaded.leafId, + loaded.activeBranchId, + loaded.materializedState, + ); + } + + static async create( + db: SqliteDatabase, + path: string, + options: { + cwd: string; + sessionId: string; + parentSessionId?: string; + metadata?: Record; + }, + ): Promise { + 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 SqliteSessionStorage( + db, + { + id: options.sessionId, + createdAt, + cwd: options.cwd, + path, + parentSessionId: options.parentSessionId, + metadata: options.metadata, + }, + [], + null, + null, + createEmptyMaterializedState(), + ); + } + + async getMetadata(): Promise { + return this.metadata; + } + + async getLeafId(): Promise { + return this.currentLeafId; + } + + async setLeafId(leafId: string | null): Promise { + if (leafId !== null && !(await this.getEntry(leafId))) { + throw new SessionError("not_found", `Entry ${leafId} not found`); + } + const entry: LeafEntry = { + type: "leaf", + id: await this.createEntryId(), + parentId: this.currentLeafId, + timestamp: new Date().toISOString(), + targetId: leafId, + }; + await this.appendEntry(entry); + } + + async createEntryId(): Promise { + for (let i = 0; i < 100; i++) { + const id = generateEntryId(this.byId); + const existing = await this.db + .prepare("SELECT 1 AS found FROM session_entries WHERE session_id = ? AND id = ? LIMIT 1") + .get<{ found: number }>(this.metadata.id, id); + if (!existing) return id; + } + return uuidv7(); + } + + async appendEntry(entry: SessionTreeEntry): Promise { + const encoded = encodeEntry(entry); + const previousMaterializedState: SessionMaterializedState = { + ...this.materializedState, + labelsById: new Map(this.materializedState.labelsById), + modelThinkingConfigs: [...this.materializedState.modelThinkingConfigs], + currentModel: this.materializedState.currentModel ? { ...this.materializedState.currentModel } : null, + }; + const previousById = new Map(this.byId); + const previousLeafId = this.currentLeafId; + const previousActiveBranchId = this.activeBranchId; + try { + applyEntryToMaterializedState(this.materializedState, entry); + await this.db.transaction(async () => { + const parentHadExistingChild = await hasExistingChild(this.db, this.metadata.id, entry.parentId); + 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(this.materializedState), 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); + } + this.byId.set(entry.id, entry); + this.currentLeafId = leafIdAfterEntry(entry); + await this.db + .prepare("UPDATE sessions SET active_leaf_id = ? WHERE id = ?") + .run(this.currentLeafId, this.metadata.id); + if (entry.type === "leaf") { + this.activeBranchId = null; + await this.materializeBranch(entry.targetId); + await this.appendToActiveBranch(entry.id, entry.parentId); + } else { + if (parentHadExistingChild) { + await this.materializeBranch(entry.parentId); + } + await this.appendToActiveBranch(entry.id, entry.parentId); + } + }); + } catch (error) { + this.materializedState = previousMaterializedState; + this.labelsById = previousMaterializedState.labelsById; + this.byId = previousById; + this.currentLeafId = previousLeafId; + this.activeBranchId = previousActiveBranchId; + if (error instanceof SessionError) throw error; + throw new SessionError("storage", `Failed to append SQLite session entry ${entry.id}`); + } + } + + async getEntry(id: string): Promise { + 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(this.metadata.id, id); + if (!row) return undefined; + try { + const entry = decodeEntry(row); + this.byId.set(entry.id, entry); + return entry; + } catch { + return undefined; + } + } + + async findEntries( + type: TType, + ): Promise>> { + const rows = await this.db + .prepare( + "SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND type = ? ORDER BY entry_seq", + ) + .all(this.metadata.id, type); + const entries: Array> = []; + for (const row of rows) { + try { + const entry = decodeEntry(row) as Extract; + this.byId.set(entry.id, entry); + entries.push(entry); + } catch { + // Keep JSONL-like permissive resume behavior: skip malformed entries. + } + } + return entries; + } + + async getLabel(id: string): Promise { + return this.labelsById.get(id); + } + + async getSessionName(): Promise { + return this.materializedState.name; + } + + async getSessionStats() { + return sessionStatsFromMaterializedState(this.materializedState); + } + + async getPathToRootOrCompaction(leafId: string | null): Promise { + if (leafId === null) return []; + if (leafId === this.currentLeafId) { + if (!this.activeBranchId) { + throw invalidSession(`missing active branch for session ${this.metadata.id} leaf ${leafId}`); + } + return getMaterializedBranchPathOrCompaction(this.db, this.metadata.id, this.activeBranchId, this.byId); + } + return this.getPathToRootOrCompactionEntries(leafId); + } + + async getEntries(options?: SessionEntryCursorOptions): Promise { + const limit = options?.limit; + if (limit !== undefined) { + const beforeOrAtEntrySeq = + options?.afterEntrySeq ?? + ( + await this.db + .prepare("SELECT entry_seq FROM session_entries WHERE session_id = ? ORDER BY entry_seq DESC LIMIT 1") + .get<{ entry_seq: number }>(this.metadata.id) + )?.entry_seq; + if (beforeOrAtEntrySeq === undefined) { + return []; + } + const rows = 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 DESC LIMIT ?", + ) + .all(this.metadata.id, beforeOrAtEntrySeq, limit); + const entries = (await decodeEntryRows(rows)).entries; + for (const entry of entries) { + this.byId.set(entry.id, entry); + } + return entries.reverse(); + } + const rows = await this.db + .prepare( + "SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? ORDER BY entry_seq", + ) + .all(this.metadata.id); + const entries = (await decodeEntryRows(rows)).entries; + for (const entry of entries) { + this.byId.set(entry.id, entry); + } + return entries; + } + + async cleanup(): Promise { + await this.db.close(); + } +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/session-entries.ts b/packages/storage/sqlite-node/src/sqlite/storage/session-entries.ts new file mode 100644 index 00000000..1b1cd75a --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/session-entries.ts @@ -0,0 +1,217 @@ +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 = Omit; + +type MessagePayload = EntryPayload>; +type ThinkingLevelChangePayload = EntryPayload>; +type ModelChangePayload = EntryPayload>; +type ActiveToolsChangePayload = EntryPayload>; +type CompactionPayload = EntryPayload>; +type BranchSummaryPayload = EntryPayload>; +type CustomPayload = EntryPayload>; +type CustomMessagePayload = EntryPayload>; +type LabelPayload = EntryPayload>; +type SessionInfoPayload = EntryPayload>; +type LeafPayload = EntryPayload>; + +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(entry: TEntry): EntryPayload { + const { type: _type, id: _id, parentId: _parentId, timestamp: _timestamp, ...payload } = entry; + return payload as EntryPayload; +} + +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}`); + } +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts b/packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts new file mode 100644 index 00000000..bba1d991 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts @@ -0,0 +1,368 @@ +import type { SessionStats, 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; + 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(); + 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(), + 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(), + 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); + } + continue; + } + if (row.type !== "label") { + } + } + return state; +} + +export function sessionStatsFromMaterializedState(state: SessionMaterializedState): SessionStats { + return { + messageCount: state.messageCount, + cachedTokens: state.cachedTokens, + uncachedTokens: state.uncachedTokens, + totalTokens: state.totalTokens, + costTotal: state.costTotal, + }; +} + +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 []; + } + } +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/session-sequences.ts b/packages/storage/sqlite-node/src/sqlite/storage/session-sequences.ts new file mode 100644 index 00000000..a38bb668 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/session-sequences.ts @@ -0,0 +1,16 @@ +import type { SqliteDatabase } from "../types.ts"; +import { invalidSession } from "./shared.ts"; + +export async function getNextSequence(db: SqliteDatabase, sessionId: string): Promise { + const sequenceRow = await 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}`); + } + return sequenceRow.next_seq; +} + +export async function advanceSequence(db: SqliteDatabase, sessionId: string, nextSeq: number): Promise { + await db.prepare("UPDATE session_sequences SET next_seq = ? WHERE session_id = ?").run(nextSeq + 1, sessionId); +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/sessions.ts b/packages/storage/sqlite-node/src/sqlite/storage/sessions.ts new file mode 100644 index 00000000..2cb4d495 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/sessions.ts @@ -0,0 +1,40 @@ +import { SessionError } from "@earendil-works/pi-agent-core"; +import type { SqliteSessionMetadata } from "../types.ts"; + +export interface SessionRow { + id: string; + created_at: string; + metadata: string | null; + cwd: string; + parent_session_id: string | null; + active_leaf_id: string | null; +} + +function parseMetadata(metadata: string | null, sessionId: string): Record | undefined { + if (metadata === null) return undefined; + let parsed: unknown; + try { + parsed = JSON.parse(metadata); + } catch (error) { + throw new SessionError( + "invalid_session", + `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`); + } + return parsed as Record; +} + +export function rowToMetadata(row: SessionRow, path: string): SqliteSessionMetadata { + return { + id: row.id, + createdAt: row.created_at, + cwd: row.cwd, + path, + parentSessionId: row.parent_session_id ?? undefined, + metadata: parseMetadata(row.metadata, row.id), + }; +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/shared.ts b/packages/storage/sqlite-node/src/sqlite/storage/shared.ts new file mode 100644 index 00000000..5fc94fcb --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/shared.ts @@ -0,0 +1,29 @@ +import type { SessionTreeEntry } from "@earendil-works/pi-agent-core"; +import { SessionError } from "@earendil-works/pi-agent-core"; +import { uuidv7 } from "@earendil-works/pi-ai"; + +export function generateEntryId(byId: { has(id: string): boolean }): string { + for (let i = 0; i < 100; i++) { + // The uuidv7 prefix is timestamp-derived and nearly constant between calls, + // so short ids must come from the random tail. + const id = uuidv7().slice(0, 8); + if (!byId.has(id)) return id; + } + return uuidv7(); +} + +export function isRecord(value: unknown): value is Record { + 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; +} diff --git a/packages/storage/sqlite-node/src/sqlite/types.ts b/packages/storage/sqlite-node/src/sqlite/types.ts new file mode 100644 index 00000000..4e97ffd2 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/types.ts @@ -0,0 +1,55 @@ +import type { FileSystem, SessionCreateOptions, SessionMetadata, SessionRepo } from "@earendil-works/pi-agent-core"; + +/** Result of a prepared SQLite statement execution. */ +export interface SqliteRunResult { + /** Number of rows changed by the statement. */ + changes: number; + /** Inserted row id when the backend exposes one. */ + lastInsertRowid?: number; +} + +/** Prepared SQLite statement capability used by the SQLite session backend. */ +export interface SqliteStatement { + run(...params: unknown[]): Promise; + get(...params: unknown[]): Promise; + all(...params: unknown[]): Promise; +} + +/** SQLite database capability used by the SQLite session backend. */ +export interface SqliteDatabase { + exec(sql: string): Promise; + prepare(sql: string): SqliteStatement; + transaction(fn: () => Promise): Promise; + close(): Promise; +} + +export interface SqliteDatabaseFactory { + open(path: string): Promise; +} + +export interface SqliteSessionMetadata extends SessionMetadata { + cwd: string; + path: string; + parentSessionId?: string; + metadata?: Record; +} + +export interface SqliteSessionCreateOptions extends SessionCreateOptions { + cwd: string; + parentSessionId?: string; + metadata?: Record; +} + +export interface SqliteSessionListOptions { + cwd?: string; +} + +export interface SqliteSessionBackendOptions { + kind: "sqlite"; + databasePath: string; +} + +export interface SqliteSessionRepoApi + extends SessionRepo {} + +export type SqliteSessionRepoEnv = Pick; diff --git a/packages/storage/sqlite-node/tsconfig.build.json b/packages/storage/sqlite-node/tsconfig.build.json new file mode 100644 index 00000000..7cdb9e03 --- /dev/null +++ b/packages/storage/sqlite-node/tsconfig.build.json @@ -0,0 +1,11 @@ +{ + "extends": "../../../tsconfig.base.json", + "compilerOptions": { + "rootDir": "src", + "outDir": "dist", + "declaration": true, + "emitDeclarationOnly": false, + "noEmit": false + }, + "include": ["src/**/*.ts"] +} diff --git a/tsconfig.json b/tsconfig.json index e144626b..626d6bbb 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -12,6 +12,7 @@ "@earendil-works/pi-ai/dist/*": ["./packages/ai/src/*"], "@earendil-works/pi-agent-core": ["./packages/agent/src/index.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"], "@earendil-works/pi-coding-agent/hooks": ["./packages/coding-agent/src/core/hooks/index.ts"], "@earendil-works/pi-coding-agent/*": ["./packages/coding-agent/src/*"], @@ -24,6 +25,6 @@ "@earendil-works/pi-agent-old/*": ["./packages/agent-old/src/*"] } }, - "include": ["packages/*/src/**/*", "packages/*/test/**/*", "packages/coding-agent/examples/**/*"], + "include": ["packages/*/src/**/*", "packages/*/test/**/*", "packages/storage/*/src/**/*", "packages/storage/*/test/**/*", "packages/coding-agent/examples/**/*"], "exclude": ["**/dist/**", "packages/coding-agent/examples/extensions/gondolin/**"] }