From 2132f8f75f0d19bdd26444d58ee7e75bc6a1a9ae Mon Sep 17 00:00:00 2001 From: Christian Klotz Date: Fri, 31 Jul 2026 16:16:42 +0300 Subject: [PATCH] fix(agent): clarify session persistence ownership --- packages/agent/docs/harness.md | 53 +- .../agent/src/harness/session/jsonl-repo.ts | 226 -------- .../src/harness/session/jsonl-storage.ts | 377 ------------- .../agent/src/harness/session/jsonl-store.ts | 355 ++++++++++++ .../agent/src/harness/session/memory-repo.ts | 102 ---- .../src/harness/session/memory-storage.ts | 190 ------- .../agent/src/harness/session/memory-store.ts | 116 ++++ .../agent/src/harness/session/repo-utils.ts | 241 --------- .../agent/src/harness/session/repository.ts | 112 ++++ .../src/harness/session/search-backend.ts | 6 +- .../agent/src/harness/session/search-index.ts | 6 +- packages/agent/src/harness/session/session.ts | 413 +++++++++----- packages/agent/src/harness/types.ts | 44 +- packages/agent/src/index.ts | 15 +- .../test/harness/agent-harness-stream.test.ts | 11 +- .../agent/test/harness/agent-harness.test.ts | 124 +++-- packages/agent/test/harness/repo.test.ts | 120 ++++- .../test/harness/session-backends.test.ts | 177 ++++++ .../agent/test/harness/session-test-utils.ts | 11 +- packages/agent/test/harness/session.test.ts | 87 +-- .../test/harness/sqlite-migrations.test.ts | 126 ++++- .../agent/test/harness/sqlite-node.test.ts | 282 +++++----- packages/agent/test/harness/storage.test.ts | 504 ------------------ packages/agent/test/scratch/simple.ts | 7 +- packages/storage/sqlite-node/README.md | 8 +- packages/storage/sqlite-node/src/index.ts | 2 +- .../storage/sqlite-node/src/sqlite/index.ts | 3 +- .../storage/sqlite-node/src/sqlite/repo.ts | 262 --------- .../sqlite-node/src/sqlite/search-backend.ts | 9 + .../sqlite-node/src/sqlite/session-store.ts | 240 +++++++++ .../src/sqlite/storage/branch-entries.ts | 5 +- .../sqlite-node/src/sqlite/storage/index.ts | 51 +- scripts/browser-smoke-entry.ts | 5 +- 33 files changed, 1917 insertions(+), 2373 deletions(-) delete mode 100644 packages/agent/src/harness/session/jsonl-repo.ts delete mode 100644 packages/agent/src/harness/session/jsonl-storage.ts create mode 100644 packages/agent/src/harness/session/jsonl-store.ts delete mode 100644 packages/agent/src/harness/session/memory-repo.ts delete mode 100644 packages/agent/src/harness/session/memory-storage.ts create mode 100644 packages/agent/src/harness/session/memory-store.ts delete mode 100644 packages/agent/src/harness/session/repo-utils.ts create mode 100644 packages/agent/src/harness/session/repository.ts create mode 100644 packages/agent/test/harness/session-backends.test.ts delete mode 100644 packages/agent/test/harness/storage.test.ts delete mode 100644 packages/storage/sqlite-node/src/sqlite/repo.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/session-store.ts diff --git a/packages/agent/docs/harness.md b/packages/agent/docs/harness.md index 9136be695..57087c572 100644 --- a/packages/agent/docs/harness.md +++ b/packages/agent/docs/harness.md @@ -1994,6 +1994,49 @@ Rules, both scopes: - **Threads are refs first.** A platform thread sharing one source of truth with its channel is a ref in the same session (section 6), not a fork. Fork when a *separate* session is wanted: subagents, exports, clones. Whether a thread becomes a ref, a fork, or a fresh session with platform backlog as prompt-time context is application policy; all three are supported. ## 14. Storage backends +### Session persistence migration + +`Session` is the only opened-session object. A `SessionStore` owns serialization, +filesystem or database access, listing, deletion, forking, and shared resources. +`SessionRepository` loads snapshots and constructs stateful `Session` aggregates. + +Store ownership is explicit and repositories borrow their stores: + +```ts +// Before: released per-session storage API +const session = new Session(new InMemorySessionStorage()); + +// After +await using store = createInMemorySessionStore(); +const repository = createSessionRepository({ store }); +const session = await repository.create({}); +``` + +Recent `SessionStore` consumers should pass the store as `store`, not `backend`, +and dispose the store rather than individual sessions or repositories: + +```ts +await using store = createJsonlSessionStore({ fs, sessionsRoot }); +const repository = createSessionRepository({ store }); +const session = await repository.open(metadata); +``` + +Custom adapters implement `create()` and `load()` by returning a +`SessionSnapshot`, persist complete entries through `appendEntry()`, and +implement multi-session `list()`, `delete()`, `fork()`, and +`Symbol.asyncDispose`. Disposal must reject new operations and wait for accepted +writes before releasing resources. Search-index adapters implement `reset()` so +rebuilds remove stale sessions. Direct `Session` construction remains available +for tests and specialized owners as `new Session(store, snapshot)`. + +The owner drains session work before store disposal. Store disposal waits only +for persistence operations already submitted to the store: + +```ts +await harness.shutdown(); +await store[Symbol.asyncDispose](); +``` + Backends implement append + read + the finder queries for one session. They know nothing about operations, queues, or recovery — the harness entry payloads are opaque to them apart from the columns they index. @@ -2291,10 +2334,10 @@ Current implementation (what is being replaced or wrapped): 5. `packages/agent/src/agent-loop.ts` — monolithic loop to split into step primitives. 6. `packages/agent/src/agent.ts` — stateful wrapper: queues, continuation, abort, settlement. 7. `packages/agent/src/harness/agent-harness.ts` — the harness this design replaces. -8. `packages/agent/src/harness/types.ts` — entry union, storage contract, event/hook types. +8. `packages/agent/src/harness/types.ts` — entry union, session store contract, event/hook types. 9. `packages/agent/src/harness/session/session.ts` — Session, context build, entry creation. -10. `packages/agent/src/harness/session/jsonl-storage.ts` — JSONL v3 format and reload. -11. `packages/agent/src/harness/session/memory-storage.ts` — in-memory parity. +10. `packages/agent/src/harness/session/jsonl-store.ts` — JSONL session store and v3 document codec. +11. `packages/agent/src/harness/session/memory-store.ts` — in-memory session store. 12. `packages/agent/src/harness/messages.ts` — defaultConvertToLlm and message helpers. 12a. `packages/ai/src/utils/transform-messages.ts` — orphaned-tool-call healing; the adjacency backstop referenced in sections 5 and 13. 13. `packages/agent/src/harness/compaction/compaction.ts` — preparation, split-turn generation, retry. @@ -2308,7 +2351,7 @@ SQLite backend: 18. `packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts` — active-branch materialization. 19. `packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts` — stats/labels/config projections. 20. `packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql` and `migrations.ts` — schema and migration mechanism. -21. `packages/storage/sqlite-node/src/sqlite/repo.ts` — create/open/fork. +21. `packages/storage/sqlite-node/src/sqlite/session-store.ts` — create/open/fork. Behavioral tests (compatibility requirements): @@ -2316,5 +2359,5 @@ Behavioral tests (compatibility requirements): 23. `packages/agent/test/agent.test.ts` 24. `packages/agent/test/harness/agent-harness.test.ts` 25. `packages/agent/test/harness/session.test.ts` -26. `packages/agent/test/harness/storage.test.ts` +26. `packages/agent/test/harness/session-stores.test.ts` 27. `packages/agent/test/harness/sqlite-migrations.test.ts` diff --git a/packages/agent/src/harness/session/jsonl-repo.ts b/packages/agent/src/harness/session/jsonl-repo.ts deleted file mode 100644 index aee81a987..000000000 --- a/packages/agent/src/harness/session/jsonl-repo.ts +++ /dev/null @@ -1,226 +0,0 @@ -import type { - FileSystem, - JsonlSessionCreateOptions, - JsonlSessionListOptions, - JsonlSessionMetadata, - LeafEntry, - SessionEntryCursorOptions, - SessionSnapshot, - SessionStorage, - SessionStore, - SessionTreeEntry, -} from "../types.ts"; -import { SessionError, toError } from "../types.ts"; -import { JsonlSessionStorage, loadJsonlSessionMetadata } from "./jsonl-storage.ts"; -import { - createSessionId, - createSessionRepository, - createTimestamp, - getEntriesToFork, - getFileSystemResultOrThrow, - type SessionRepository, -} from "./repo-utils.ts"; -import { createScanningSessionSearch } from "./search-backend.ts"; - -export type JsonlSessionStoreOptions = { fs: JsonlSessionStoreFileSystem; sessionsRoot: string }; - -export type JsonlSessionStoreFileSystem = Pick< - FileSystem, - | "cwd" - | "absolutePath" - | "joinPath" - | "readTextFile" - | "readTextLines" - | "writeFile" - | "appendFile" - | "listDir" - | "exists" - | "createDir" - | "remove" ->; - -function encodeCwd(cwd: string): string { - return `--${cwd.replace(/^[/\\]/, "").replace(/[/\\:]/g, "-")}--`; -} - -class JsonlSessionStore - implements SessionStore -{ - private readonly fs: JsonlSessionStoreFileSystem; - private readonly sessionsRootInput: string; - private sessionsRoot: string | undefined; - - constructor(options: JsonlSessionStoreOptions) { - this.fs = options.fs; - this.sessionsRootInput = options.sessionsRoot; - } - - private async getSessionsRoot(): Promise { - if (!this.sessionsRoot) { - this.sessionsRoot = getFileSystemResultOrThrow( - await this.fs.absolutePath(this.sessionsRootInput), - `Failed to resolve sessions root ${this.sessionsRootInput}`, - ); - } - return this.sessionsRoot; - } - - private async getSessionDir(cwd: string): Promise { - return getFileSystemResultOrThrow( - await this.fs.joinPath([await this.getSessionsRoot(), encodeCwd(cwd)]), - `Failed to resolve session directory for ${cwd}`, - ); - } - - private async createSessionFilePath(cwd: string, sessionId: string, timestamp: string): Promise { - return getFileSystemResultOrThrow( - await this.fs.joinPath([ - await this.getSessionDir(cwd), - `${timestamp.replace(/[:.]/g, "-")}_${sessionId}.jsonl`, - ]), - `Failed to resolve session file path for ${sessionId}`, - ); - } - - async create(options: JsonlSessionCreateOptions): Promise { - const id = options.id ?? createSessionId(); - const createdAt = createTimestamp(); - const sessionDir = await this.getSessionDir(options.cwd); - getFileSystemResultOrThrow( - await this.fs.createDir(sessionDir, { recursive: true }), - `Failed to create session directory ${sessionDir}`, - ); - const filePath = await this.createSessionFilePath(options.cwd, id, createdAt); - const storage = await JsonlSessionStorage.create(this.fs, filePath, { - cwd: options.cwd, - sessionId: id, - parentSessionPath: options.parentSessionPath, - metadata: options.metadata, - }); - return await storage.getMetadata(); - } - - async open(metadata: JsonlSessionMetadata): Promise> { - if ( - !getFileSystemResultOrThrow(await this.fs.exists(metadata.path), `Failed to check session ${metadata.path}`) - ) { - throw new SessionError("not_found", `Session not found: ${metadata.path}`); - } - return await JsonlSessionStorage.open(this.fs, metadata.path); - } - - async load(metadata: JsonlSessionMetadata): Promise> { - const storage = await this.open(metadata); - return { - metadata: await storage.getMetadata(), - leafId: await storage.getLeafId(), - entries: await storage.getEntries(), - }; - } - - async list(options: JsonlSessionListOptions = {}): Promise { - const dirs = options.cwd ? [await this.getSessionDir(options.cwd)] : await this.listSessionDirs(); - const sessions: JsonlSessionMetadata[] = []; - for (const dir of dirs) { - if (!getFileSystemResultOrThrow(await this.fs.exists(dir), `Failed to check session directory ${dir}`)) { - continue; - } - const files = getFileSystemResultOrThrow( - await this.fs.listDir(dir), - `Failed to list sessions in ${dir}`, - ).filter((file) => file.kind !== "directory" && file.name.endsWith(".jsonl")); - for (const file of files) { - try { - sessions.push(await loadJsonlSessionMetadata(this.fs, file.path)); - } catch (error) { - const cause = toError(error); - if (!(cause instanceof SessionError) || cause.code !== "invalid_session") throw cause; - } - } - } - sessions.sort((a, b) => new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime()); - return sessions; - } - - async getEntries(metadata: JsonlSessionMetadata, options?: SessionEntryCursorOptions): Promise { - return await (await this.open(metadata)).getEntries(options); - } - - async createEntryId(metadata: JsonlSessionMetadata): Promise { - return (await this.open(metadata)).createEntryId(); - } - - async appendEntry(metadata: JsonlSessionMetadata, entry: SessionTreeEntry): Promise { - await (await this.open(metadata)).appendEntry(entry); - } - - async setLeafId(metadata: JsonlSessionMetadata, leafId: string | null): Promise { - return await (await this.open(metadata)).setLeafId(leafId); - } - - async delete(metadata: JsonlSessionMetadata): Promise { - getFileSystemResultOrThrow( - await this.fs.remove(metadata.path, { force: true }), - `Failed to delete session ${metadata.path}`, - ); - } - - async fork( - sourceMetadata: JsonlSessionMetadata, - options: JsonlSessionCreateOptions & { entryId?: string; position?: "before" | "at"; id?: string }, - ): Promise { - const source = await this.open(sourceMetadata); - const forkedEntries = await getEntriesToFork(source, options); - const id = options.id ?? createSessionId(); - const createdAt = createTimestamp(); - const sessionDir = await this.getSessionDir(options.cwd); - getFileSystemResultOrThrow( - await this.fs.createDir(sessionDir, { recursive: true }), - `Failed to create session directory ${sessionDir}`, - ); - const storage = await JsonlSessionStorage.create( - this.fs, - await this.createSessionFilePath(options.cwd, id, createdAt), - { - cwd: options.cwd, - sessionId: id, - parentSessionPath: options.parentSessionPath ?? sourceMetadata.path, - metadata: options.metadata ?? sourceMetadata.metadata, - }, - ); - for (const entry of forkedEntries) { - await storage.appendEntry(entry); - } - return await storage.getMetadata(); - } - - private async listSessionDirs(): Promise { - const sessionsRoot = await this.getSessionsRoot(); - if ( - !getFileSystemResultOrThrow( - await this.fs.exists(sessionsRoot), - `Failed to check sessions root ${sessionsRoot}`, - ) - ) { - return []; - } - const entries = getFileSystemResultOrThrow( - await this.fs.listDir(sessionsRoot), - `Failed to list sessions root ${sessionsRoot}`, - ); - return entries.filter((entry) => entry.kind === "directory").map((entry) => entry.path); - } -} - -export function createJsonlSessionStore( - options: JsonlSessionStoreOptions, -): SessionStore { - return new JsonlSessionStore(options); -} - -export function createJsonlSessionRepository( - options: JsonlSessionStoreOptions, -): SessionRepository { - const store = createJsonlSessionStore(options); - return createSessionRepository({ store, search: createScanningSessionSearch(store) }); -} diff --git a/packages/agent/src/harness/session/jsonl-storage.ts b/packages/agent/src/harness/session/jsonl-storage.ts deleted file mode 100644 index 943ccfb0a..000000000 --- a/packages/agent/src/harness/session/jsonl-storage.ts +++ /dev/null @@ -1,377 +0,0 @@ -import { uuidv7 } from "@earendil-works/pi-ai"; -import type { - FileSystem, - JsonlSessionMetadata, - LeafEntry, - SessionEntryCursorOptions, - SessionStorage, - SessionTreeEntry, -} from "../types.ts"; -import { SessionError, toError } from "../types.ts"; -import { getFileSystemResultOrThrow } from "./repo-utils.ts"; - -type JsonlSessionStorageFileSystem = Pick; - -interface SessionHeader { - type: "session"; - version: 3; - id: string; - timestamp: string; - cwd: string; - parentSession?: string; - metadata?: Record; -} - -function updateLabelCache(labelsById: Map, entry: SessionTreeEntry): void { - if (entry.type !== "label") return; - const label = entry.label?.trim(); - if (label) { - labelsById.set(entry.targetId, label); - } else { - labelsById.delete(entry.targetId); - } -} - -function buildLabelsById(entries: SessionTreeEntry[]): Map { - const labelsById = new Map(); - for (const entry of entries) { - updateLabelCache(labelsById, entry); - } - return labelsById; -} - -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(-8); - if (!byId.has(id)) return id; - } - return uuidv7(); -} - -function invalidSession(filePath: string, message: string, cause?: Error): SessionError { - return new SessionError("invalid_session", `Invalid JSONL session file ${filePath}: ${message}`, cause); -} - -function invalidEntry(filePath: string, lineNumber: number, message: string, cause?: Error): SessionError { - return new SessionError( - "invalid_entry", - `Invalid JSONL session file ${filePath}: line ${lineNumber} ${message}`, - cause, - ); -} - -function parseHeaderLine(line: string, filePath: string): SessionHeader { - let parsed: unknown; - try { - parsed = JSON.parse(line); - } catch (error) { - throw invalidSession(filePath, "first line is not a valid session header", toError(error)); - } - if (typeof parsed !== "object" || parsed === null) { - throw invalidSession(filePath, "first line is not a valid session header"); - } - const header = parsed as Partial; - if (header.type !== "session") throw invalidSession(filePath, "first line is not a valid session header"); - if (header.version !== 3) throw invalidSession(filePath, "unsupported session version"); - if (typeof header.id !== "string" || !header.id) throw invalidSession(filePath, "session header is missing id"); - if (typeof header.timestamp !== "string" || !header.timestamp) { - throw invalidSession(filePath, "session header is missing timestamp"); - } - if (typeof header.cwd !== "string" || !header.cwd) throw invalidSession(filePath, "session header is missing cwd"); - if (header.parentSession !== undefined && typeof header.parentSession !== "string") { - throw invalidSession(filePath, "session header parentSession must be a string"); - } - if ( - header.metadata !== undefined && - (typeof header.metadata !== "object" || header.metadata === null || Array.isArray(header.metadata)) - ) { - throw invalidSession(filePath, "session header metadata must be an object"); - } - return { - type: "session", - version: 3, - id: header.id, - timestamp: header.timestamp, - cwd: header.cwd, - parentSession: header.parentSession, - metadata: header.metadata, - }; -} - -function parseEntryLine(line: string, filePath: string, lineNumber: number): SessionTreeEntry { - let parsed: unknown; - try { - parsed = JSON.parse(line); - } catch (error) { - throw invalidEntry(filePath, lineNumber, "is not valid JSON", toError(error)); - } - if (typeof parsed !== "object" || parsed === null) { - throw invalidEntry(filePath, lineNumber, "is not a valid session entry"); - } - const entry = parsed as { - type?: unknown; - id?: unknown; - parentId?: unknown; - timestamp?: unknown; - targetId?: unknown; - }; - if (typeof entry.type !== "string") throw invalidEntry(filePath, lineNumber, "is missing entry type"); - if (typeof entry.id !== "string" || !entry.id) throw invalidEntry(filePath, lineNumber, "is missing entry id"); - if (entry.parentId !== null && typeof entry.parentId !== "string") { - throw invalidEntry(filePath, lineNumber, "has invalid parentId"); - } - if (typeof entry.timestamp !== "string" || !entry.timestamp) { - throw invalidEntry(filePath, lineNumber, "is missing timestamp"); - } - if (entry.type === "leaf" && entry.targetId !== null && typeof entry.targetId !== "string") { - throw invalidEntry(filePath, lineNumber, "has invalid targetId"); - } - return entry as SessionTreeEntry; -} - -function leafIdAfterEntry(entry: SessionTreeEntry): string | null { - return entry.type === "leaf" ? entry.targetId : entry.id; -} - -function headerToSessionMetadata(header: SessionHeader, path: string): JsonlSessionMetadata { - return { - id: header.id, - createdAt: header.timestamp, - cwd: header.cwd, - path, - parentSessionPath: header.parentSession, - metadata: header.metadata, - }; -} - -export async function loadJsonlSessionMetadata( - fs: JsonlSessionStorageFileSystem, - filePath: string, -): Promise { - const lines = getFileSystemResultOrThrow( - await fs.readTextLines(filePath, { maxLines: 1 }), - `Failed to read session header ${filePath}`, - ); - const line = lines[0]; - if (line?.trim()) return headerToSessionMetadata(parseHeaderLine(line, filePath), filePath); - throw invalidSession(filePath, "missing session header"); -} - -async function loadJsonlStorage( - fs: JsonlSessionStorageFileSystem, - filePath: string, -): Promise<{ - header: SessionHeader; - entries: SessionTreeEntry[]; - leafId: string | null; -}> { - const content = getFileSystemResultOrThrow(await fs.readTextFile(filePath), `Failed to read session ${filePath}`); - const lines = content.split("\n").filter((line) => line.trim()); - if (lines.length === 0) { - throw invalidSession(filePath, "missing session header"); - } - - const header = parseHeaderLine(lines[0]!, filePath); - const entries: SessionTreeEntry[] = []; - let leafId: string | null = null; - for (let i = 1; i < lines.length; i++) { - const entry = parseEntryLine(lines[i]!, filePath, i + 1); - entries.push(entry); - leafId = leafIdAfterEntry(entry); - } - return { header, entries, leafId }; -} - -export class JsonlSessionStorage implements SessionStorage { - private readonly fs: JsonlSessionStorageFileSystem; - private readonly filePath: string; - private readonly metadata: JsonlSessionMetadata; - private entries: SessionTreeEntry[]; - private byId: Map; - private labelsById: Map; - private currentLeafId: string | null; - - private constructor( - fs: JsonlSessionStorageFileSystem, - filePath: string, - header: SessionHeader, - entries: SessionTreeEntry[], - leafId: string | null, - ) { - this.fs = fs; - this.filePath = filePath; - this.metadata = headerToSessionMetadata(header, this.filePath); - this.entries = entries; - this.byId = new Map(entries.map((entry) => [entry.id, entry])); - this.labelsById = buildLabelsById(entries); - this.currentLeafId = leafId; - } - - static async open(fs: JsonlSessionStorageFileSystem, filePath: string): Promise { - const loaded = await loadJsonlStorage(fs, filePath); - return new JsonlSessionStorage(fs, filePath, loaded.header, loaded.entries, loaded.leafId); - } - - static async create( - fs: JsonlSessionStorageFileSystem, - filePath: string, - options: { - cwd: string; - sessionId: string; - parentSessionPath?: string; - metadata?: Record; - }, - ): Promise { - const header: SessionHeader = { - type: "session", - version: 3, - id: options.sessionId, - timestamp: new Date().toISOString(), - cwd: options.cwd, - parentSession: options.parentSessionPath, - metadata: options.metadata, - }; - getFileSystemResultOrThrow( - await fs.writeFile(filePath, `${JSON.stringify(header)}\n`), - `Failed to create session ${filePath}`, - ); - return new JsonlSessionStorage(fs, filePath, header, [], null); - } - - async getMetadata(): Promise { - return this.metadata; - } - - async getLeafId(): Promise { - if (this.currentLeafId !== null && !this.byId.has(this.currentLeafId)) { - throw new SessionError("invalid_session", `Entry ${this.currentLeafId} not found`); - } - return this.currentLeafId; - } - - async setLeafId(leafId: string | null): Promise { - if (leafId !== null && !this.byId.has(leafId)) { - throw new SessionError("not_found", `Entry ${leafId} not found`); - } - const entry: LeafEntry = { - type: "leaf", - id: generateEntryId(this.byId), - parentId: this.currentLeafId, - timestamp: new Date().toISOString(), - targetId: leafId, - }; - getFileSystemResultOrThrow( - await this.fs.appendFile(this.filePath, `${JSON.stringify(entry)}\n`), - `Failed to append session leaf ${entry.id}`, - ); - this.entries.push(entry); - this.byId.set(entry.id, entry); - this.currentLeafId = leafId; - return entry; - } - - async createEntryId(): Promise { - return generateEntryId(this.byId); - } - - async appendEntry(entry: SessionTreeEntry): Promise { - getFileSystemResultOrThrow( - await this.fs.appendFile(this.filePath, `${JSON.stringify(entry)}\n`), - `Failed to append session entry ${entry.id}`, - ); - this.entries.push(entry); - this.byId.set(entry.id, entry); - updateLabelCache(this.labelsById, entry); - this.currentLeafId = leafIdAfterEntry(entry); - } - - async getEntry(id: string): Promise { - return this.byId.get(id); - } - - async findEntries( - type: TType, - ): Promise>> { - return this.entries.filter((entry): entry is Extract => entry.type === type); - } - - async getLabel(id: string): Promise { - return this.labelsById.get(id); - } - - 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`); - current = parent; - } - return path; - } - - 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/jsonl-store.ts b/packages/agent/src/harness/session/jsonl-store.ts new file mode 100644 index 000000000..00f8aa6df --- /dev/null +++ b/packages/agent/src/harness/session/jsonl-store.ts @@ -0,0 +1,355 @@ +import type { + FileSystem, + JsonlSessionCreateOptions, + JsonlSessionListOptions, + JsonlSessionMetadata, + SessionSnapshot, + SessionStore, + SessionTreeEntry, +} from "../types.ts"; +import { SessionError, toError } from "../types.ts"; +import { createSessionId, createTimestamp, getFileSystemResultOrThrow } from "./repository.ts"; + +export type JsonlSessionStoreOptions = { fs: JsonlSessionStoreFileSystem; sessionsRoot: string }; +export type JsonlSessionStoreFileSystem = Pick< + FileSystem, + | "absolutePath" + | "joinPath" + | "readTextFile" + | "readTextLines" + | "writeFile" + | "appendFile" + | "listDir" + | "exists" + | "createDir" + | "remove" +>; + +type JsonlSessionFileSystem = Pick; + +interface SessionHeader { + type: "session"; + version: 3; + id: string; + timestamp: string; + cwd: string; + parentSession?: string; + metadata?: Record; +} + +function invalidSession(path: string, message: string, cause?: Error): SessionError { + return new SessionError("invalid_session", `Invalid JSONL session file ${path}: ${message}`, cause); +} + +function invalidEntry(path: string, line: number, message: string, cause?: Error): SessionError { + return new SessionError("invalid_entry", `Invalid JSONL session file ${path}: line ${line} ${message}`, cause); +} + +function parseHeader(line: string, path: string): SessionHeader { + let value: unknown; + try { + value = JSON.parse(line); + } catch (error) { + throw invalidSession(path, "first line is not a valid session header", toError(error)); + } + if (typeof value !== "object" || value === null) + throw invalidSession(path, "first line is not a valid session header"); + const header = value as Partial; + if (header.type !== "session" || header.version !== 3) { + throw invalidSession( + path, + header.type === "session" ? "unsupported session version" : "first line is not a valid session header", + ); + } + if (typeof header.id !== "string" || !header.id) throw invalidSession(path, "session header is missing id"); + if (typeof header.timestamp !== "string" || !header.timestamp) + throw invalidSession(path, "session header is missing timestamp"); + if (typeof header.cwd !== "string" || !header.cwd) throw invalidSession(path, "session header is missing cwd"); + if (header.parentSession !== undefined && typeof header.parentSession !== "string") { + throw invalidSession(path, "session header parentSession must be a string"); + } + if ( + header.metadata !== undefined && + (typeof header.metadata !== "object" || header.metadata === null || Array.isArray(header.metadata)) + ) { + throw invalidSession(path, "session header metadata must be an object"); + } + return { + type: "session", + version: 3, + id: header.id, + timestamp: header.timestamp, + cwd: header.cwd, + parentSession: header.parentSession, + metadata: header.metadata, + }; +} + +function parseEntry(line: string, path: string, lineNumber: number): SessionTreeEntry { + let value: unknown; + try { + value = JSON.parse(line); + } catch (error) { + throw invalidEntry(path, lineNumber, "is not valid JSON", toError(error)); + } + if (typeof value !== "object" || value === null) + throw invalidEntry(path, lineNumber, "is not a valid session entry"); + const entry = value as { type?: unknown; id?: unknown; parentId?: unknown; timestamp?: unknown; targetId?: unknown }; + if (typeof entry.type !== "string") throw invalidEntry(path, lineNumber, "is missing entry type"); + if (typeof entry.id !== "string" || !entry.id) throw invalidEntry(path, lineNumber, "is missing entry id"); + if (entry.parentId !== null && typeof entry.parentId !== "string") + throw invalidEntry(path, lineNumber, "has invalid parentId"); + if (typeof entry.timestamp !== "string" || !entry.timestamp) + throw invalidEntry(path, lineNumber, "is missing timestamp"); + if (entry.type === "leaf" && entry.targetId !== null && typeof entry.targetId !== "string") { + throw invalidEntry(path, lineNumber, "has invalid targetId"); + } + return entry as SessionTreeEntry; +} + +function metadataFromHeader(header: SessionHeader, path: string): JsonlSessionMetadata { + return { + id: header.id, + createdAt: header.timestamp, + cwd: header.cwd, + path, + parentSessionPath: header.parentSession, + metadata: header.metadata, + }; +} + +export async function loadJsonlSessionMetadata( + fs: JsonlSessionFileSystem, + path: string, +): Promise { + const lines = getFileSystemResultOrThrow( + await fs.readTextLines(path, { maxLines: 1 }), + `Failed to read session header ${path}`, + ); + if (!lines[0]?.trim()) throw invalidSession(path, "missing session header"); + return metadataFromHeader(parseHeader(lines[0], path), path); +} + +async function loadJsonlSession( + fs: JsonlSessionFileSystem, + path: string, +): Promise> { + const content = getFileSystemResultOrThrow(await fs.readTextFile(path), `Failed to read session ${path}`); + const lines = content.split("\n").filter((line) => line.trim()); + if (lines.length === 0) throw invalidSession(path, "missing session header"); + const header = parseHeader(lines[0]!, path); + const entries = lines.slice(1).map((line, index) => parseEntry(line, path, index + 2)); + const entryIds = new Set(); + for (const entry of entries) { + if (entryIds.has(entry.id)) throw invalidSession(path, `duplicate entry id ${entry.id}`); + entryIds.add(entry.id); + } + return { + metadata: metadataFromHeader(header, path), + entries, + }; +} + +function encodeCwd(cwd: string): string { + return `--${cwd.replace(/^[/\\]/, "").replace(/[/\\:]/g, "-")}--`; +} + +class SerialOperationQueue { + private tail: Promise = Promise.resolve(); + + enqueue(operation: () => Promise | T): Promise { + const result = this.tail.then(operation); + this.tail = result.then( + () => undefined, + () => undefined, + ); + return result; + } + + async drain(): Promise { + await this.tail; + } +} + +class JsonlSessionStore + implements SessionStore +{ + private readonly fs: JsonlSessionStoreFileSystem; + private readonly sessionsRootInput: string; + private sessionsRoot: string | undefined; + private readonly entryIdsBySession = new Map>(); + private readonly operations = new SerialOperationQueue(); + private disposed = false; + private disposePromise: Promise | undefined; + + constructor(options: JsonlSessionStoreOptions) { + this.fs = options.fs; + this.sessionsRootInput = options.sessionsRoot; + } + + create(options: JsonlSessionCreateOptions): Promise> { + this.assertOpen(); + return this.operations.enqueue(() => + this.createDocument(options, options.parentSessionPath, options.metadata, []), + ); + } + + load(metadata: JsonlSessionMetadata): Promise> { + this.assertOpen(); + return this.operations.enqueue(() => this.loadDocument(metadata)); + } + + private async loadDocument(metadata: JsonlSessionMetadata): Promise> { + if ( + !getFileSystemResultOrThrow(await this.fs.exists(metadata.path), `Failed to check session ${metadata.path}`) + ) { + throw new SessionError("not_found", `Session not found: ${metadata.path}`); + } + const snapshot = await loadJsonlSession(this.fs, metadata.path); + this.entryIdsBySession.set(metadata.id, new Set(snapshot.entries.map((entry) => entry.id))); + return snapshot; + } + + list(options: JsonlSessionListOptions = {}): Promise { + this.assertOpen(); + return this.operations.enqueue(() => this.listSessions(options)); + } + + private async listSessions(options: JsonlSessionListOptions): Promise { + const dirs = options.cwd ? [await this.getSessionDir(options.cwd)] : await this.listSessionDirs(); + const sessions: JsonlSessionMetadata[] = []; + for (const dir of dirs) { + if (!getFileSystemResultOrThrow(await this.fs.exists(dir), `Failed to check session directory ${dir}`)) + continue; + const files = getFileSystemResultOrThrow( + await this.fs.listDir(dir), + `Failed to list sessions in ${dir}`, + ).filter((file) => file.kind !== "directory" && file.name.endsWith(".jsonl")); + for (const file of files) sessions.push(await loadJsonlSessionMetadata(this.fs, file.path)); + } + return sessions.sort((a, b) => new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime()); + } + + appendEntry(metadata: JsonlSessionMetadata, entry: SessionTreeEntry): Promise { + this.assertOpen(); + return this.operations.enqueue(async () => { + if ( + !getFileSystemResultOrThrow(await this.fs.exists(metadata.path), `Failed to check session ${metadata.path}`) + ) { + throw new SessionError("not_found", `Session not found: ${metadata.path}`); + } + let entryIds = this.entryIdsBySession.get(metadata.id); + if (!entryIds) { + const snapshot = await this.loadDocument(metadata); + entryIds = new Set(snapshot.entries.map((existing) => existing.id)); + } + if (entryIds.has(entry.id)) throw new SessionError("invalid_entry", `Entry ${entry.id} already exists`); + getFileSystemResultOrThrow( + await this.fs.appendFile(metadata.path, `${JSON.stringify(entry)}\n`), + `Failed to append session entry ${entry.id}`, + ); + entryIds.add(entry.id); + this.entryIdsBySession.set(metadata.id, entryIds); + }); + } + + delete(metadata: JsonlSessionMetadata): Promise { + this.assertOpen(); + return this.operations.enqueue(async () => { + getFileSystemResultOrThrow( + await this.fs.remove(metadata.path, { force: true }), + `Failed to delete session ${metadata.path}`, + ); + this.entryIdsBySession.delete(metadata.id); + }); + } + + fork( + source: JsonlSessionMetadata, + options: JsonlSessionCreateOptions, + entries: readonly SessionTreeEntry[], + ): Promise> { + this.assertOpen(); + return this.operations.enqueue(() => + this.createDocument( + options, + options.parentSessionPath ?? source.path, + options.metadata ?? source.metadata, + entries, + ), + ); + } + + async [Symbol.asyncDispose](): Promise { + if (!this.disposePromise) { + this.disposed = true; + this.disposePromise = this.operations.drain(); + } + await this.disposePromise; + } + + private assertOpen(): void { + if (this.disposed) throw new SessionError("storage", "JSONL session store is disposed"); + } + + private async createDocument( + options: JsonlSessionCreateOptions, + parentSessionPath: string | undefined, + metadata: Record | undefined, + entries: readonly SessionTreeEntry[], + ): Promise> { + const id = options.id ?? createSessionId(); + const timestamp = createTimestamp(); + const dir = await this.getSessionDir(options.cwd); + getFileSystemResultOrThrow( + await this.fs.createDir(dir, { recursive: true }), + `Failed to create session directory ${dir}`, + ); + const path = getFileSystemResultOrThrow( + await this.fs.joinPath([dir, `${timestamp.replace(/[:.]/g, "-")}_${id}.jsonl`]), + `Failed to resolve session file path for ${id}`, + ); + const header: SessionHeader = { + type: "session", + version: 3, + id, + timestamp, + cwd: options.cwd, + parentSession: parentSessionPath, + metadata, + }; + const content = [JSON.stringify(header), ...entries.map((entry) => JSON.stringify(entry)), ""].join("\n"); + getFileSystemResultOrThrow(await this.fs.writeFile(path, content), `Failed to create session ${path}`); + this.entryIdsBySession.set(id, new Set(entries.map((entry) => entry.id))); + return { metadata: metadataFromHeader(header, path), entries: [...entries] }; + } + + private async getSessionsRoot(): Promise { + this.sessionsRoot ??= getFileSystemResultOrThrow( + await this.fs.absolutePath(this.sessionsRootInput), + `Failed to resolve sessions root ${this.sessionsRootInput}`, + ); + return this.sessionsRoot; + } + + private async getSessionDir(cwd: string): Promise { + return getFileSystemResultOrThrow( + await this.fs.joinPath([await this.getSessionsRoot(), encodeCwd(cwd)]), + `Failed to resolve session directory for ${cwd}`, + ); + } + + private async listSessionDirs(): Promise { + const root = await this.getSessionsRoot(); + if (!getFileSystemResultOrThrow(await this.fs.exists(root), `Failed to check sessions root ${root}`)) return []; + return getFileSystemResultOrThrow(await this.fs.listDir(root), `Failed to list sessions root ${root}`) + .filter((entry) => entry.kind === "directory") + .map((entry) => entry.path); + } +} + +export function createJsonlSessionStore( + options: JsonlSessionStoreOptions, +): SessionStore { + return new JsonlSessionStore(options); +} diff --git a/packages/agent/src/harness/session/memory-repo.ts b/packages/agent/src/harness/session/memory-repo.ts deleted file mode 100644 index 6072a6ece..000000000 --- a/packages/agent/src/harness/session/memory-repo.ts +++ /dev/null @@ -1,102 +0,0 @@ -import { - type LeafEntry, - type SessionEntryCursorOptions, - SessionError, - type SessionMetadata, - type SessionSnapshot, - type SessionStorage, - type SessionStore, - type SessionTreeEntry, -} from "../types.ts"; -import { InMemorySessionStorage } from "./memory-storage.ts"; -import { - createSessionId, - createSessionRepository, - createTimestamp, - getEntriesToFork, - type SessionRepository, -} from "./repo-utils.ts"; -import { createScanningSessionSearch } from "./search-backend.ts"; - -export type InMemorySessionCreateOptions = { id?: string }; - -class InMemorySessionStore implements SessionStore { - private sessions = new Map>(); - - async create(options: InMemorySessionCreateOptions = {}): Promise { - const metadata: SessionMetadata = { - id: options.id ?? createSessionId(), - createdAt: createTimestamp(), - }; - const storage = new InMemorySessionStorage({ metadata }); - this.sessions.set(metadata.id, storage); - return metadata; - } - - async open(metadata: SessionMetadata): Promise> { - const storage = this.sessions.get(metadata.id); - if (!storage) throw new SessionError("not_found", `Session not found: ${metadata.id}`); - return storage; - } - - async load(metadata: SessionMetadata): Promise> { - const storage = await this.open(metadata); - return { - metadata: await storage.getMetadata(), - leafId: await storage.getLeafId(), - entries: await storage.getEntries(), - }; - } - - async list(): Promise { - return Promise.all([...this.sessions.values()].map((storage) => storage.getMetadata())); - } - - async getEntries(metadata: SessionMetadata, options?: SessionEntryCursorOptions): Promise { - return await (await this.open(metadata)).getEntries(options); - } - - async createEntryId(metadata: SessionMetadata): Promise { - return (await this.open(metadata)).createEntryId(); - } - - async appendEntry(metadata: SessionMetadata, entry: SessionTreeEntry): Promise { - await (await this.open(metadata)).appendEntry(entry); - } - - async setLeafId(metadata: SessionMetadata, leafId: string | null): Promise { - return await (await this.open(metadata)).setLeafId(leafId); - } - - async delete(metadata: SessionMetadata): Promise { - this.sessions.delete(metadata.id); - } - - async fork( - sourceMetadata: SessionMetadata, - options: { entryId?: string; position?: "before" | "at"; id?: string }, - ): Promise { - const source = await this.open(sourceMetadata); - const forkedEntries = await getEntriesToFork(source, options); - const metadata: SessionMetadata = { - id: options.id ?? createSessionId(), - createdAt: createTimestamp(), - }; - const storage = new InMemorySessionStorage({ metadata, entries: forkedEntries }); - this.sessions.set(metadata.id, storage); - return metadata; - } -} - -export function createInMemorySessionStore(): SessionStore { - return new InMemorySessionStore(); -} - -export function createInMemorySessionRepository(): SessionRepository< - SessionMetadata, - InMemorySessionCreateOptions, - void -> { - const store = createInMemorySessionStore(); - return createSessionRepository({ store, search: createScanningSessionSearch(store) }); -} diff --git a/packages/agent/src/harness/session/memory-storage.ts b/packages/agent/src/harness/session/memory-storage.ts deleted file mode 100644 index d8e53be3b..000000000 --- a/packages/agent/src/harness/session/memory-storage.ts +++ /dev/null @@ -1,190 +0,0 @@ -import { uuidv7 } from "@earendil-works/pi-ai"; -import { - type LeafEntry, - type SessionEntryCursorOptions, - SessionError, - type SessionMetadata, - type SessionStorage, - type SessionTreeEntry, -} from "../types.ts"; - -function updateLabelCache(labelsById: Map, entry: SessionTreeEntry): void { - if (entry.type !== "label") return; - const label = entry.label?.trim(); - if (label) { - labelsById.set(entry.targetId, label); - } else { - labelsById.delete(entry.targetId); - } -} - -function buildLabelsById(entries: SessionTreeEntry[]): Map { - const labelsById = new Map(); - for (const entry of entries) { - updateLabelCache(labelsById, entry); - } - return labelsById; -} - -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(-8); - if (!byId.has(id)) return id; - } - return uuidv7(); -} - -function leafIdAfterEntry(entry: SessionTreeEntry): string | null { - return entry.type === "leaf" ? entry.targetId : entry.id; -} - -export class InMemorySessionStorage - implements SessionStorage -{ - private readonly metadata: TMetadata; - private entries: SessionTreeEntry[]; - private byId: Map; - private labelsById: Map; - private leafId: string | null; - - constructor(options?: { entries?: SessionTreeEntry[]; metadata?: TMetadata }) { - this.entries = options?.entries ? [...options.entries] : []; - this.byId = new Map(this.entries.map((entry) => [entry.id, entry])); - this.labelsById = buildLabelsById(this.entries); - this.leafId = null; - for (const entry of this.entries) this.leafId = leafIdAfterEntry(entry); - if (this.leafId !== null && !this.byId.has(this.leafId)) { - throw new SessionError("invalid_session", `Entry ${this.leafId} not found`); - } - this.metadata = options?.metadata ?? ({ id: uuidv7(), createdAt: new Date().toISOString() } as TMetadata); - } - - async getMetadata(): Promise { - return this.metadata; - } - - async getLeafId(): Promise { - if (this.leafId !== null && !this.byId.has(this.leafId)) { - throw new SessionError("invalid_session", `Entry ${this.leafId} not found`); - } - return this.leafId; - } - - async setLeafId(leafId: string | null): Promise { - if (leafId !== null && !this.byId.has(leafId)) { - throw new SessionError("not_found", `Entry ${leafId} not found`); - } - const entry: LeafEntry = { - type: "leaf", - id: generateEntryId(this.byId), - parentId: this.leafId, - timestamp: new Date().toISOString(), - targetId: leafId, - }; - this.entries.push(entry); - this.byId.set(entry.id, entry); - this.leafId = leafId; - return entry; - } - - async createEntryId(): Promise { - return generateEntryId(this.byId); - } - - async appendEntry(entry: SessionTreeEntry): Promise { - this.entries.push(entry); - this.byId.set(entry.id, entry); - updateLabelCache(this.labelsById, entry); - this.leafId = leafIdAfterEntry(entry); - } - - async getEntry(id: string): Promise { - return this.byId.get(id); - } - - async findEntries( - type: TType, - ): Promise>> { - return this.entries.filter((entry): entry is Extract => entry.type === type); - } - - async getLabel(id: string): Promise { - return this.labelsById.get(id); - } - - 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`); - current = parent; - } - return path; - } - - 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-store.ts b/packages/agent/src/harness/session/memory-store.ts new file mode 100644 index 000000000..d2999abee --- /dev/null +++ b/packages/agent/src/harness/session/memory-store.ts @@ -0,0 +1,116 @@ +import type { SessionMetadata, SessionSnapshot, SessionStore, SessionTreeEntry } from "../types.ts"; +import { SessionError } from "../types.ts"; +import { createSessionId, createTimestamp } from "./repository.ts"; + +export type InMemorySessionCreateOptions = { id?: string }; + +interface InMemorySessionState { + metadata: SessionMetadata; + entries: SessionTreeEntry[]; +} + +class SerialOperationQueue { + private tail: Promise = Promise.resolve(); + + enqueue(operation: () => Promise | T): Promise { + const result = this.tail.then(operation); + this.tail = result.then( + () => undefined, + () => undefined, + ); + return result; + } + + async drain(): Promise { + await this.tail; + } +} + +class InMemorySessionStore implements SessionStore { + private readonly sessions = new Map(); + private readonly operations = new SerialOperationQueue(); + private disposed = false; + private disposePromise: Promise | undefined; + + create(options: InMemorySessionCreateOptions = {}): Promise> { + this.assertOpen(); + return this.operations.enqueue(() => { + const state: InMemorySessionState = { + metadata: { id: options.id ?? createSessionId(), createdAt: createTimestamp() }, + entries: [], + }; + this.sessions.set(state.metadata.id, state); + return this.snapshot(state); + }); + } + + load(metadata: SessionMetadata): Promise> { + this.assertOpen(); + return this.operations.enqueue(() => this.snapshot(this.getState(metadata))); + } + + list(): Promise { + this.assertOpen(); + return this.operations.enqueue(() => [...this.sessions.values()].map((state) => state.metadata)); + } + + appendEntry(metadata: SessionMetadata, entry: SessionTreeEntry): Promise { + this.assertOpen(); + return this.operations.enqueue(() => { + const state = this.getState(metadata); + if (state.entries.some((existing) => existing.id === entry.id)) { + throw new SessionError("invalid_entry", `Entry ${entry.id} already exists`); + } + state.entries.push(entry); + }); + } + + delete(metadata: SessionMetadata): Promise { + this.assertOpen(); + return this.operations.enqueue(() => { + this.sessions.delete(metadata.id); + }); + } + + fork( + _source: SessionMetadata, + options: InMemorySessionCreateOptions, + entries: readonly SessionTreeEntry[], + ): Promise> { + this.assertOpen(); + return this.operations.enqueue(() => { + const state: InMemorySessionState = { + metadata: { id: options.id ?? createSessionId(), createdAt: createTimestamp() }, + entries: [...entries], + }; + this.sessions.set(state.metadata.id, state); + return this.snapshot(state); + }); + } + + async [Symbol.asyncDispose](): Promise { + if (!this.disposePromise) { + this.disposed = true; + this.disposePromise = this.operations.drain(); + } + await this.disposePromise; + } + + private assertOpen(): void { + if (this.disposed) throw new SessionError("storage", "In-memory session store is disposed"); + } + + private getState(metadata: SessionMetadata): InMemorySessionState { + const state = this.sessions.get(metadata.id); + if (!state) throw new SessionError("not_found", `Session not found: ${metadata.id}`); + return state; + } + + private snapshot(state: InMemorySessionState): SessionSnapshot { + return { metadata: state.metadata, entries: [...state.entries] }; + } +} + +export function createInMemorySessionStore(): SessionStore { + return new InMemorySessionStore(); +} diff --git a/packages/agent/src/harness/session/repo-utils.ts b/packages/agent/src/harness/session/repo-utils.ts deleted file mode 100644 index f4e3b290e..000000000 --- a/packages/agent/src/harness/session/repo-utils.ts +++ /dev/null @@ -1,241 +0,0 @@ -import { uuidv7 } from "@earendil-works/pi-ai"; -import { - type FileError, - type Result, - type SessionCreateOptions, - SessionError, - type SessionForkOptions, - type SessionMetadata, - type SessionSearch, - type SessionSearchHit, - type SessionStats, - type SessionStorage, - type SessionStore, - type SessionTreeEntry, -} from "../types.ts"; -import { Session } from "./session.ts"; - -export function createSessionId(): string { - return uuidv7(); -} - -export function createTimestamp(): string { - return new Date().toISOString(); -} - -export function toSession(storage: SessionStorage): Session { - return new Session(storage); -} - -function entriesById(entries: readonly SessionTreeEntry[]): Map { - return new Map(entries.map((entry) => [entry.id, entry])); -} - -function getLabel(entries: readonly SessionTreeEntry[], id: string): string | undefined { - let label: string | undefined; - for (const entry of entries) { - if (entry.type !== "label" || entry.targetId !== id) continue; - label = entry.label?.trim() || undefined; - } - return label; -} - -function getSessionName(entries: readonly SessionTreeEntry[]): string | undefined { - for (let i = entries.length - 1; i >= 0; i--) { - const entry = entries[i]!; - if (entry.type === "session_info") return entry.name?.trim() || undefined; - } - return undefined; -} - -function getSessionStats(entries: readonly SessionTreeEntry[]): SessionStats { - let messageCount = 0; - let cachedTokens = 0; - let uncachedTokens = 0; - let totalTokens = 0; - let costTotal = 0; - for (const entry of 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 }; -} - -function getPathToRootOrCompaction(entries: readonly SessionTreeEntry[], leafId: string | null): SessionTreeEntry[] { - if (leafId === null) return []; - const byId = entriesById(entries); - const path: SessionTreeEntry[] = []; - let stopAtEntryId: string | null = null; - let current = 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 = byId.get(current.parentId); - if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); - current = parent; - } - return path; -} - -export function toStoreSession( - store: Pick, "load" | "getEntries" | "createEntryId" | "appendEntry" | "setLeafId">, - metadata: TMetadata, -): Session { - const load = () => store.load(metadata); - const storage: SessionStorage = { - async getMetadata() { - return metadata; - }, - async getLeafId() { - return (await load()).leafId; - }, - setLeafId: (leafId) => store.setLeafId(metadata, leafId), - createEntryId: () => store.createEntryId(metadata), - appendEntry: (entry) => store.appendEntry(metadata, entry), - async getEntry(id) { - return entriesById((await load()).entries).get(id); - }, - async findEntries(type) { - return (await load()).entries.filter( - (entry): entry is Extract => entry.type === type, - ); - }, - async getLabel(id) { - return getLabel((await load()).entries, id); - }, - async getSessionName() { - return getSessionName((await load()).entries); - }, - async getSessionStats() { - return getSessionStats((await load()).entries); - }, - async getPathToRootOrCompaction(leafId) { - return getPathToRootOrCompaction((await load()).entries, leafId); - }, - getEntries: (options) => store.getEntries(metadata, options), - }; - return new Session(storage); -} - -export class SessionRepository< - TMetadata extends SessionMetadata = SessionMetadata, - TCreateOptions extends SessionCreateOptions = SessionCreateOptions, - TListOptions = void, -> { - private readonly store: SessionStore; - private readonly searchBackend: SessionSearch | null; - - constructor(options: { - store: SessionStore; - search?: SessionSearch | null; - }) { - this.store = options.store; - this.searchBackend = options.search ?? null; - } - - async create(options: TCreateOptions): Promise> { - return toStoreSession(this.store, await this.store.create(options)); - } - - async open(metadata: TMetadata): Promise> { - const snapshot = await this.store.load(metadata); - return toStoreSession(this.store, snapshot.metadata); - } - - list(options?: TListOptions): Promise { - return this.store.list(options); - } - - async delete(metadata: TMetadata): Promise { - await this.store.delete(metadata); - } - - async fork(source: TMetadata, options: SessionForkOptions & TCreateOptions): Promise> { - return toStoreSession(this.store, await this.store.fork(source, options)); - } - - async search(options: Parameters["search"]>[0]): Promise[]> { - return this.searchBackend ? await this.searchBackend.search(options) : []; - } -} - -export function createSessionRepository< - TMetadata extends SessionMetadata = SessionMetadata, - TCreateOptions extends SessionCreateOptions = SessionCreateOptions, - TListOptions = void, ->(options: { - store: SessionStore; - search?: SessionSearch | null; -}): SessionRepository { - return new SessionRepository(options); -} - -export function findSessionEntryMatches( - metadata: TMetadata, - entries: SessionTreeEntry[], - text: string, -): SessionSearchHit[] { - const normalizedText = text.trim().toLowerCase(); - if (!normalizedText) return []; - return entries.flatMap((entry) => { - const payload = JSON.stringify(entry); - if (!payload.toLowerCase().includes(normalizedText)) return []; - return [{ metadata, entryId: entry.id, timestamp: entry.timestamp, snippet: payload }]; - }); -} - -export function getFileSystemResultOrThrow(result: Result, message: string): TValue { - if (!result.ok) { - const code = result.error.code === "not_found" ? "not_found" : "storage"; - throw new SessionError(code, `${message}: ${result.error.message}`, result.error); - } - return result.value; -} - -export async function getEntriesToFork( - storage: SessionStorage, - options: { entryId?: string; position?: "before" | "at" }, -): Promise { - if (!options.entryId) return storage.getEntries(); - const target = await storage.getEntry(options.entryId); - if (!target) { - throw new SessionError("invalid_fork_target", `Entry ${options.entryId} not found`); - } - let effectiveLeafId: string | null; - if ((options.position ?? "before") === "at") { - effectiveLeafId = target.id; - } else { - if (target.type !== "message" || target.message.role !== "user") { - throw new SessionError("invalid_fork_target", `Entry ${options.entryId} is not a user message`); - } - effectiveLeafId = target.parentId; - } - return storage.getPathToRootOrCompaction(effectiveLeafId); -} diff --git a/packages/agent/src/harness/session/repository.ts b/packages/agent/src/harness/session/repository.ts new file mode 100644 index 000000000..ad48a31b6 --- /dev/null +++ b/packages/agent/src/harness/session/repository.ts @@ -0,0 +1,112 @@ +import { uuidv7 } from "@earendil-works/pi-ai"; +import { + type FileError, + type Result, + type SessionCreateOptions, + SessionError, + type SessionForkOptions, + type SessionMetadata, + type SessionSearch, + type SessionSearchHit, + type SessionStore, + type SessionTreeEntry, +} from "../types.ts"; +import { Session } from "./session.ts"; + +export function createSessionId(): string { + return uuidv7(); +} + +export function createTimestamp(): string { + return new Date().toISOString(); +} + +export class SessionRepository< + TMetadata extends SessionMetadata = SessionMetadata, + TCreateOptions extends SessionCreateOptions = SessionCreateOptions, + TListOptions = void, +> { + private readonly store: SessionStore; + private readonly searchStore: SessionSearch | null; + + constructor(options: { + store: SessionStore; + search?: SessionSearch | null; + }) { + this.store = options.store; + this.searchStore = options.search ?? null; + } + + async create(options: TCreateOptions): Promise> { + return new Session(this.store, await this.store.create(options)); + } + + async open(metadata: TMetadata): Promise> { + return new Session(this.store, await this.store.load(metadata)); + } + + async list(options?: TListOptions): Promise { + return await this.store.list(options); + } + + async delete(metadata: TMetadata): Promise { + await this.store.delete(metadata); + } + + async fork(source: TMetadata, options: SessionForkOptions & TCreateOptions): Promise> { + const sourceSession = await this.open(source); + const entries = await getEntriesToFork(sourceSession, options); + return new Session(this.store, await this.store.fork(source, options, entries)); + } + + async search(options: Parameters["search"]>[0]): Promise[]> { + return this.searchStore ? await this.searchStore.search(options) : []; + } +} + +export function createSessionRepository< + TMetadata extends SessionMetadata = SessionMetadata, + TCreateOptions extends SessionCreateOptions = SessionCreateOptions, + TListOptions = void, +>(options: { + store: SessionStore; + search?: SessionSearch | null; +}): SessionRepository { + return new SessionRepository(options); +} + +export function findSessionEntryMatches( + metadata: TMetadata, + entries: readonly SessionTreeEntry[], + text: string, +): SessionSearchHit[] { + const normalizedText = text.trim().toLowerCase(); + if (!normalizedText) return []; + return entries.flatMap((entry) => { + const payload = JSON.stringify(entry); + if (!payload.toLowerCase().includes(normalizedText)) return []; + return [{ metadata, entryId: entry.id, timestamp: entry.timestamp, snippet: payload }]; + }); +} + +export function getFileSystemResultOrThrow(result: Result, message: string): TValue { + if (!result.ok) { + const code = result.error.code === "not_found" ? "not_found" : "storage"; + throw new SessionError(code, `${message}: ${result.error.message}`, result.error); + } + return result.value; +} + +export async function getEntriesToFork( + source: Pick, + options: { entryId?: string; position?: "before" | "at" }, +): Promise { + if (!options.entryId) return source.getEntries(); + const target = await source.getEntry(options.entryId); + if (!target) throw new SessionError("invalid_fork_target", `Entry ${options.entryId} not found`); + if ((options.position ?? "before") === "at") return source.getBranch(target.id); + if (target.type !== "message" || target.message.role !== "user") { + throw new SessionError("invalid_fork_target", `Entry ${options.entryId} is not a user message`); + } + return source.getBranch(target.parentId); +} diff --git a/packages/agent/src/harness/session/search-backend.ts b/packages/agent/src/harness/session/search-backend.ts index 8eceb53de..fbad0f3ab 100644 --- a/packages/agent/src/harness/session/search-backend.ts +++ b/packages/agent/src/harness/session/search-backend.ts @@ -5,7 +5,7 @@ import type { SessionSearchOptions, SessionSnapshot, } from "../types.ts"; -import { findSessionEntryMatches } from "./repo-utils.ts"; +import { findSessionEntryMatches } from "./repository.ts"; type SessionSearchSource = { load(metadata: TMetadata): Promise>; @@ -25,8 +25,8 @@ class ScanningSessionSearch for (const metadata of await this.source.list()) { const cwd = (metadata as { cwd?: unknown }).cwd; if (options.cwd !== undefined && cwd !== options.cwd) continue; - const state = await this.source.load(metadata); - hits.push(...findSessionEntryMatches(metadata, state.entries, options.text)); + const snapshot = await this.source.load(metadata); + hits.push(...findSessionEntryMatches(snapshot.metadata, snapshot.entries, options.text)); } return hits; } diff --git a/packages/agent/src/harness/session/search-index.ts b/packages/agent/src/harness/session/search-index.ts index bef8408c9..12164b933 100644 --- a/packages/agent/src/harness/session/search-index.ts +++ b/packages/agent/src/harness/session/search-index.ts @@ -7,10 +7,12 @@ export interface SessionSearchIndexSource( source: SessionSearchIndexSource, - index: Pick, "replaceSession">, + index: Pick, "reset" | "replaceSession">, options?: TListOptions, ): Promise { + await index.reset(); for (const metadata of await source.list(options)) { - await index.replaceSession(metadata, (await source.load(metadata)).entries); + const snapshot = await source.load(metadata); + await index.replaceSession(snapshot.metadata, snapshot.entries); } } diff --git a/packages/agent/src/harness/session/session.ts b/packages/agent/src/harness/session/session.ts index 57a7d8c9e..79663fafa 100644 --- a/packages/agent/src/harness/session/session.ts +++ b/packages/agent/src/harness/session/session.ts @@ -1,4 +1,4 @@ -import type { ImageContent, TextContent, Usage } from "@earendil-works/pi-ai"; +import { type ImageContent, type TextContent, type Usage, uuidv7 } from "@earendil-works/pi-ai"; import type { AgentMessage } from "../../types.ts"; import { createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage } from "../messages.ts"; import type { @@ -15,8 +15,9 @@ import type { SessionEntryCursorOptions, SessionInfoEntry, SessionMetadata, + SessionSnapshot, SessionStats, - SessionStorage, + SessionStore, SessionTreeEntry, ThinkingLevelChangeEntry, } from "../types.ts"; @@ -148,33 +149,160 @@ export function buildSessionContext( return { ...state, messages }; } -export class Session { - private readonly storage: SessionStorage; - private readonly contextBuildOptions: SessionContextBuildOptions; +class SessionEntryIndex { + private readonly entries: SessionTreeEntry[]; + private readonly byId: Map; + private readonly labelsById = new Map(); + private leafId: string | null = null; - constructor(storage: SessionStorage, contextBuildOptions: SessionContextBuildOptions = {}) { - this.storage = storage; + constructor(entries: readonly SessionTreeEntry[]) { + this.entries = [...entries]; + this.byId = new Map(this.entries.map((entry) => [entry.id, entry])); + for (const entry of this.entries) this.applyProjection(entry); + if (this.leafId !== null && !this.byId.has(this.leafId)) { + throw new SessionError("invalid_session", `Entry ${this.leafId} not found`); + } + } + + getLeafId(): string | null { + return this.leafId; + } + getEntry(id: string): SessionTreeEntry | undefined { + return this.byId.get(id); + } + + getEntries(options?: SessionEntryCursorOptions): SessionTreeEntry[] { + if (options?.limit === undefined) return [...this.entries]; + const end = options.afterEntrySeq ?? this.entries.length; + return this.entries.slice(Math.max(0, end - options.limit), end); + } + + append(entry: SessionTreeEntry): void { + this.entries.push(entry); + this.byId.set(entry.id, entry); + this.applyProjection(entry); + } + + assertCanAppend(entry: SessionTreeEntry): void { + if (this.byId.has(entry.id)) throw new SessionError("invalid_entry", `Entry ${entry.id} already exists`); + } + + createEntryId(): string { + for (let i = 0; i < 100; i++) { + const id = uuidv7().slice(-8); + if (!this.byId.has(id)) return id; + } + return uuidv7(); + } + + getLabel(id: string): string | undefined { + return this.labelsById.get(id); + } + + getSessionName(): string | undefined { + for (let i = this.entries.length - 1; i >= 0; i--) { + const entry = this.entries[i]!; + if (entry.type === "session_info") return entry.name?.trim() || undefined; + } + return undefined; + } + + getSessionStats(): SessionStats { + 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 }; + } + + getPathToRootOrCompaction(leafId: string | null): SessionTreeEntry[] { + 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`); + current = parent; + } + return path; + } + + private applyProjection(entry: SessionTreeEntry): void { + this.leafId = entry.type === "leaf" ? entry.targetId : entry.id; + if (entry.type !== "label") return; + const label = entry.label?.trim(); + if (label) this.labelsById.set(entry.targetId, label); + else this.labelsById.delete(entry.targetId); + } +} + +export class Session { + private readonly store: Pick, "appendEntry">; + private readonly metadata: TMetadata; + private readonly index: SessionEntryIndex; + private readonly contextBuildOptions: SessionContextBuildOptions; + private appendTail: Promise = Promise.resolve(); + + constructor( + store: Pick, "appendEntry">, + snapshot: SessionSnapshot, + contextBuildOptions: SessionContextBuildOptions = {}, + ) { + this.store = store; + this.metadata = snapshot.metadata; + this.index = new SessionEntryIndex(snapshot.entries); this.contextBuildOptions = contextBuildOptions; } - getMetadata(): Promise { - return this.storage.getMetadata(); + async getMetadata(): Promise { + return this.metadata; + } + async getLeafId(): Promise { + return this.index.getLeafId(); + } + async getEntry(id: string): Promise { + return this.index.getEntry(id); + } + async getEntries(options?: SessionEntryCursorOptions): Promise { + return this.index.getEntries(options); } - getLeafId(): Promise { - return this.storage.getLeafId(); - } - - getEntry(id: string): Promise { - return this.storage.getEntry(id); - } - - getEntries(options?: SessionEntryCursorOptions): Promise { - return this.storage.getEntries(options); - } - - async getBranch(fromId?: string): Promise { - return this.storage.getPathToRootOrCompaction(fromId ?? (await this.storage.getLeafId())); + async getBranch(fromId?: string | null): Promise { + return this.index.getPathToRootOrCompaction(fromId === undefined ? this.index.getLeafId() : fromId); } async buildContextEntries(options: SessionContextBuildOptions = {}): Promise { @@ -195,74 +323,95 @@ export class Session { }; } - getLabel(id: string): Promise { - return this.storage.getLabel(id); + async getLabel(id: string): Promise { + return this.index.getLabel(id); + } + async getSessionStats(): Promise { + return this.index.getSessionStats(); + } + async getSessionName(): Promise { + return this.index.getSessionName(); } - getSessionStats(): Promise { - return this.storage.getSessionStats(); - } - - getSessionName(): Promise { - return this.storage.getSessionName(); - } - - private async appendEntry(entry: SessionTreeEntry): Promise { - await this.storage.appendEntry(entry); + private enqueueAppend( + createEntry: (base: Pick) => TEntry, + ): Promise { + const operation = this.appendTail.then(async () => { + const entry = createEntry({ + id: this.index.createEntryId(), + parentId: this.index.getLeafId(), + timestamp: new Date().toISOString(), + }); + this.index.assertCanAppend(entry); + await this.store.appendEntry(this.metadata, entry); + this.index.append(entry); + return entry; + }); + this.appendTail = operation.then( + () => undefined, + () => undefined, + ); + return operation; } private setLeafId(leafId: string | null): Promise { - return this.storage.setLeafId(leafId); + return this.enqueueAppend((base) => { + if (leafId !== null && !this.index.getEntry(leafId)) { + throw new SessionError("not_found", `Entry ${leafId} not found`); + } + return { ...base, type: "leaf", targetId: leafId }; + }); } - private async createEntryId(): Promise { - return this.storage.createEntryId(); - } - - private async appendTypedEntry(entry: TEntry): Promise { - await this.appendEntry(entry); - return entry.id; + private async appendTypedEntry( + createEntry: (base: Pick) => TEntry, + ): Promise { + return (await this.enqueueAppend(createEntry)).id; } async appendMessage(message: AgentMessage): Promise { - return this.appendTypedEntry({ - type: "message", - id: await this.createEntryId(), - parentId: await this.getLeafId(), - timestamp: new Date().toISOString(), - message, - } satisfies MessageEntry); + return this.appendTypedEntry( + (base) => + ({ + ...base, + type: "message", + message, + }) satisfies MessageEntry, + ); } async appendThinkingLevelChange(thinkingLevel: string): Promise { - return this.appendTypedEntry({ - type: "thinking_level_change", - id: await this.createEntryId(), - parentId: await this.getLeafId(), - timestamp: new Date().toISOString(), - thinkingLevel, - } satisfies ThinkingLevelChangeEntry); + return this.appendTypedEntry( + (base) => + ({ + ...base, + type: "thinking_level_change", + thinkingLevel, + }) satisfies ThinkingLevelChangeEntry, + ); } async appendModelChange(provider: string, modelId: string): Promise { - return this.appendTypedEntry({ - type: "model_change", - id: await this.createEntryId(), - parentId: await this.getLeafId(), - timestamp: new Date().toISOString(), - provider, - modelId, - } satisfies ModelChangeEntry); + return this.appendTypedEntry( + (base) => + ({ + ...base, + type: "model_change", + provider, + modelId, + }) satisfies ModelChangeEntry, + ); } async appendActiveToolsChange(activeToolNames: string[]): Promise { - return this.appendTypedEntry({ - type: "active_tools_change", - id: await this.createEntryId(), - parentId: await this.getLeafId(), - timestamp: new Date().toISOString(), - activeToolNames: [...activeToolNames], - } satisfies ActiveToolsChangeEntry); + return this.appendTypedEntry( + (base) => + ({ + ...base, + type: "active_tools_change", + activeToolNames: [...activeToolNames], + }) satisfies ActiveToolsChangeEntry, + ); } async appendCompaction( @@ -274,30 +423,32 @@ export class Session { usage?: Usage, retainedTail?: AgentMessage[], ): Promise { - return this.appendTypedEntry({ - type: "compaction", - id: await this.createEntryId(), - parentId: await this.getLeafId(), - timestamp: new Date().toISOString(), - summary, - firstKeptEntryId, - tokensBefore, - retainedTail, - details, - usage, - fromHook, - } satisfies CompactionEntry); + return this.appendTypedEntry( + (base) => + ({ + ...base, + type: "compaction", + summary, + firstKeptEntryId, + tokensBefore, + retainedTail, + details, + usage, + fromHook, + }) satisfies CompactionEntry, + ); } async appendCustomEntry(customType: string, data?: unknown): Promise { - return this.appendTypedEntry({ - type: "custom", - id: await this.createEntryId(), - parentId: await this.getLeafId(), - timestamp: new Date().toISOString(), - customType, - data, - } satisfies CustomEntry); + return this.appendTypedEntry( + (base) => + ({ + ...base, + type: "custom", + customType, + data, + }) satisfies CustomEntry, + ); } async appendCustomMessageEntry( @@ -306,41 +457,44 @@ export class Session { display: boolean, details?: T, ): Promise { - return this.appendTypedEntry({ - type: "custom_message", - id: await this.createEntryId(), - parentId: await this.getLeafId(), - timestamp: new Date().toISOString(), - customType, - content, - display, - details, - } satisfies CustomMessageEntry); + return this.appendTypedEntry( + (base) => + ({ + ...base, + type: "custom_message", + customType, + content, + display, + details, + }) satisfies CustomMessageEntry, + ); } async appendLabel(targetId: string, label: string | undefined): Promise { if (!(await this.getEntry(targetId))) { throw new SessionError("not_found", `Entry ${targetId} not found`); } - return this.appendTypedEntry({ - type: "label", - id: await this.createEntryId(), - parentId: await this.getLeafId(), - timestamp: new Date().toISOString(), - targetId, - label, - } satisfies LabelEntry); + return this.appendTypedEntry( + (base) => + ({ + ...base, + type: "label", + targetId, + label, + }) satisfies LabelEntry, + ); } async appendSessionName(name: string): Promise { const sanitizedName = name.replace(/[\r\n]+/g, " ").trim(); - return this.appendTypedEntry({ - type: "session_info", - id: await this.createEntryId(), - parentId: await this.getLeafId(), - timestamp: new Date().toISOString(), - name: sanitizedName, - } satisfies SessionInfoEntry); + return this.appendTypedEntry( + (base) => + ({ + ...base, + type: "session_info", + name: sanitizedName, + }) satisfies SessionInfoEntry, + ); } async moveTo( @@ -352,16 +506,17 @@ export class Session { } await this.setLeafId(entryId); if (!summary) return undefined; - return this.appendTypedEntry({ - type: "branch_summary", - id: await this.createEntryId(), - parentId: entryId, - timestamp: new Date().toISOString(), - fromId: entryId ?? "root", - summary: summary.summary, - details: summary.details, - usage: summary.usage, - fromHook: summary.fromHook, - } satisfies BranchSummaryEntry); + return this.appendTypedEntry( + (base) => + ({ + ...base, + type: "branch_summary", + fromId: entryId ?? "root", + summary: summary.summary, + details: summary.details, + usage: summary.usage, + fromHook: summary.fromHook, + }) satisfies BranchSummaryEntry, + ); } } diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index 8fe59793e..96db531a5 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -495,30 +495,6 @@ export interface SessionEntryCursorOptions { limit?: number; } -export interface SessionSnapshot { - metadata: TMetadata; - leafId: string | null; - entries: SessionTreeEntry[]; -} - -export interface SessionStorage { - getMetadata(): Promise; - getLeafId(): Promise; - /** Persist a leaf entry that records the active session-tree leaf. */ - setLeafId(leafId: string | null): Promise; - createEntryId(): Promise; - appendEntry(entry: SessionTreeEntry): Promise; - getEntry(id: string): Promise; - findEntries( - type: TType, - ): Promise>>; - getLabel(id: string): Promise; - getSessionName(): Promise; - getSessionStats(): Promise; - getPathToRootOrCompaction(leafId: string | null): Promise; - getEntries(options?: SessionEntryCursorOptions): Promise; -} - export type { Session } from "./session/session.ts"; export interface SessionCreateOptions { @@ -540,6 +516,7 @@ export interface SessionSearchHit { + reset(): Promise; upsertEntry(metadata: TMetadata, entry: SessionTreeEntry): Promise; replaceSession(metadata: TMetadata, entries: readonly SessionTreeEntry[]): Promise; deleteSession(metadata: TMetadata): Promise; @@ -556,20 +533,27 @@ export interface SessionForkOptions { id?: string; } +export interface SessionSnapshot { + metadata: TMetadata; + entries: SessionTreeEntry[]; +} + +/** Owns persistence and resources shared by all sessions in a repository. */ export interface SessionStore< TMetadata extends SessionMetadata = SessionMetadata, TCreateOptions extends SessionCreateOptions = SessionCreateOptions, TListOptions = void, -> { - create(options: TCreateOptions): Promise; +> extends AsyncDisposable { + create(options: TCreateOptions): Promise>; load(metadata: TMetadata): Promise>; list(options?: TListOptions): Promise; - getEntries(metadata: TMetadata, options?: SessionEntryCursorOptions): Promise; - createEntryId(metadata: TMetadata): Promise; appendEntry(metadata: TMetadata, entry: SessionTreeEntry): Promise; - setLeafId(metadata: TMetadata, leafId: string | null): Promise; delete(metadata: TMetadata): Promise; - fork(source: TMetadata, options: SessionForkOptions & TCreateOptions): Promise; + fork( + source: TMetadata, + options: SessionForkOptions & TCreateOptions, + entries: readonly SessionTreeEntry[], + ): Promise>; } export interface JsonlSessionCreateOptions extends SessionCreateOptions { diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index 9811e84d2..4444940e7 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -29,11 +29,16 @@ export { } from "./harness/compaction/compaction.ts"; export * from "./harness/messages.ts"; export * from "./harness/prompt-templates.ts"; -export * from "./harness/session/jsonl-repo.ts"; -export * from "./harness/session/jsonl-storage.ts"; -export * from "./harness/session/memory-repo.ts"; -export * from "./harness/session/memory-storage.ts"; -export * from "./harness/session/repo-utils.ts"; +export { + createJsonlSessionStore, + type JsonlSessionStoreFileSystem, + type JsonlSessionStoreOptions, +} from "./harness/session/jsonl-store.ts"; +export { + createInMemorySessionStore, + type InMemorySessionCreateOptions, +} from "./harness/session/memory-store.ts"; +export * from "./harness/session/repository.ts"; export * from "./harness/session/search-backend.ts"; export * from "./harness/session/search-index.ts"; export * from "./harness/session/session.ts"; diff --git a/packages/agent/test/harness/agent-harness-stream.test.ts b/packages/agent/test/harness/agent-harness-stream.test.ts index 3d493802c..0b1e5402f 100644 --- a/packages/agent/test/harness/agent-harness-stream.test.ts +++ b/packages/agent/test/harness/agent-harness-stream.test.ts @@ -8,10 +8,9 @@ import { } from "@earendil-works/pi-ai"; import { describe, expect, it } from "vitest"; import { AgentHarness } from "../../src/harness/agent-harness.ts"; -import { InMemorySessionStorage } from "../../src/harness/session/memory-storage.ts"; -import { Session } from "../../src/harness/session/session.ts"; import type { AgentHarnessOptions } from "../../src/harness/types.ts"; import { calculateTool } from "../utils/calculate.ts"; +import { createInMemorySession } from "./session-test-utils.ts"; /** Shared collection; each faux provider gets a unique id so coexisting fakes route correctly. */ const models = createModels(); @@ -46,7 +45,7 @@ describe("AgentHarness stream configuration", () => { }, ]); - const session = new Session(new InMemorySessionStorage({ metadata: { id: "session-1", createdAt: "now" } })); + const session = await createInMemorySession("session-1"); const harness = createHarness({ models, session, @@ -97,7 +96,7 @@ describe("AgentHarness stream configuration", () => { const harness = createHarness({ models, - session: new Session(new InMemorySessionStorage()), + session: await createInMemorySession(), model: registration.getModel(), streamOptions: { timeoutMs: 1000, @@ -154,7 +153,7 @@ describe("AgentHarness stream configuration", () => { const harness = createHarness({ models, - session: new Session(new InMemorySessionStorage()), + session: await createInMemorySession(), model: registration.getModel(), tools: [calculateTool], streamOptions: { timeoutMs: 1000, headers: { turn: "first" } }, @@ -188,7 +187,7 @@ describe("AgentHarness stream configuration", () => { const harness = createHarness({ models, - session: new Session(new InMemorySessionStorage()), + session: await createInMemorySession(), model: registration.getModel(), }); diff --git a/packages/agent/test/harness/agent-harness.test.ts b/packages/agent/test/harness/agent-harness.test.ts index 9b1ec714b..fe1ba6412 100644 --- a/packages/agent/test/harness/agent-harness.test.ts +++ b/packages/agent/test/harness/agent-harness.test.ts @@ -11,12 +11,13 @@ import { getModel } from "@earendil-works/pi-ai/compat"; import { describe, expect, it } from "vitest"; import { AgentHarness } from "../../src/harness/agent-harness.ts"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; -import { InMemorySessionStorage } from "../../src/harness/session/memory-storage.ts"; +import { createInMemorySessionStore } from "../../src/harness/session/memory-store.ts"; import { Session } from "../../src/harness/session/session.ts"; import type { AgentHarnessTool, PromptTemplate, Skill } from "../../src/harness/types.ts"; import type { AgentMessage, AgentTool } from "../../src/types.ts"; import { calculateTool, createCalculateToolWithUsage } from "../utils/calculate.ts"; import { getCurrentTimeTool } from "../utils/get-current-time.ts"; +import { createInMemorySession } from "./session-test-utils.ts"; interface AppSkill extends Skill { source: "project" | "user"; @@ -89,28 +90,37 @@ function createAssistantMessage(text: string): AgentMessage { }; } -class BlockingSessionStorage extends InMemorySessionStorage { - readonly allWritesStarted = deferred(); - readonly releaseWrites = deferred(); - private readonly expectedWrites: number; - private writesStarted = 0; - - constructor(expectedWrites: number) { - super(); - this.expectedWrites = expectedWrites; - } - - override async appendEntry(entry: Parameters[0]): Promise { - this.writesStarted++; - if (this.writesStarted === this.expectedWrites) this.allWritesStarted.resolve(); - await this.releaseWrites.promise; - await super.appendEntry(entry); - } +async function createBlockingSession(expectedWrites: number): Promise<{ + session: Session; + allWritesStarted: ReturnType; + releaseWrites: ReturnType; + getEntries(): ReturnType; +}> { + const source = createInMemorySessionStore(); + const snapshot = await source.create({}); + const allWritesStarted = deferred(); + const releaseWrites = deferred(); + let writesStarted = 0; + const blockingStore = { + async appendEntry(metadata: typeof snapshot.metadata, entry: (typeof snapshot.entries)[number]) { + writesStarted++; + if (writesStarted === expectedWrites) allWritesStarted.resolve(); + await releaseWrites.promise; + await source.appendEntry(metadata, entry); + }, + }; + const session = new Session(blockingStore, snapshot); + return { + session, + allWritesStarted, + releaseWrites, + getEntries: () => session.getEntries(), + }; } describe("AgentHarness", () => { - it("constructs directly and exposes queue modes", () => { - const session = new Session(new InMemorySessionStorage()); + it("constructs directly and exposes queue modes", async () => { + const session = await createInMemorySession(); const initialModel = getModel("anthropic", "claude-sonnet-4-5"); const harness = new AgentHarness({ models, @@ -146,7 +156,7 @@ describe("AgentHarness", () => { ]); const harness = new AgentHarness({ models, - session: new Session(new InMemorySessionStorage()), + session: await createInMemorySession(), model: registration.getModel(), }); const prompt = harness.prompt("hello"); @@ -190,7 +200,7 @@ describe("AgentHarness", () => { ]); const harness = new AgentHarness({ models, - session: new Session(new InMemorySessionStorage()), + session: await createInMemorySession(), model: registration.getModel(), }); harness.on("before_agent_start", async () => { @@ -227,7 +237,7 @@ describe("AgentHarness", () => { return fauxAssistantMessage("summary produced after shutdown"); }, ]); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); await session.appendMessage(createUserMessage("one")); await session.appendMessage(createAssistantMessage("two")); const harness = new AgentHarness({ models, session, model: registration.getModel() }); @@ -261,7 +271,7 @@ describe("AgentHarness", () => { return fauxAssistantMessage("summary produced after shutdown"); }, ]); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); const targetId = await session.appendMessage(createUserMessage("first branch")); await session.appendMessage(createAssistantMessage("first reply")); await session.appendMessage(createUserMessage("abandoned work")); @@ -285,14 +295,14 @@ describe("AgentHarness", () => { }); it("does not treat concurrent mutations as active operations", async () => { - const storage = new BlockingSessionStorage(1); + const blocking = await createBlockingSession(1); const harness = new AgentHarness({ models, - session: new Session(storage), + session: blocking.session, model: getModel("anthropic", "claude-sonnet-4-5"), }); const mutation = harness.appendMessage(createUserMessage("concurrent")); - await storage.allWritesStarted.promise; + await blocking.allWritesStarted.promise; const firstSettlement = await Promise.race([ harness.waitForIdle().then(() => "idle" as const), @@ -300,15 +310,15 @@ describe("AgentHarness", () => { ]); expect(firstSettlement).toBe("idle"); - storage.releaseWrites.resolve(); + blocking.releaseWrites.resolve(); await mutation; }); it("awaits concurrent idle session mutations before shutdown resolves", async () => { - const storage = new BlockingSessionStorage(3); + const blocking = await createBlockingSession(1); const harness = new AgentHarness({ models, - session: new Session(storage), + session: blocking.session, model: getModel("anthropic", "claude-sonnet-4-5"), }); const nextModel = getModel("anthropic", "claude-haiku-4-5"); @@ -317,7 +327,7 @@ describe("AgentHarness", () => { harness.setModel(nextModel), harness.setThinkingLevel("high"), ]; - await storage.allWritesStarted.promise; + await blocking.allWritesStarted.promise; const shutdown = harness.shutdown(); const firstSettlement = await Promise.race([ @@ -326,13 +336,13 @@ describe("AgentHarness", () => { ]); expect(firstSettlement).toBe("writes-pending"); - storage.releaseWrites.resolve(); + blocking.releaseWrites.resolve(); await Promise.all([...mutations, shutdown]); - expect(await storage.getEntries()).toHaveLength(3); + expect(await blocking.getEntries()).toHaveLength(3); }); it("shuts down an idle harness without modifying its durable session", async () => { - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); await session.appendMessage(createUserMessage("existing")); const harness = new AgentHarness({ models, @@ -368,7 +378,7 @@ describe("AgentHarness", () => { ]); const harness = new AgentHarness({ models, - session: new Session(new InMemorySessionStorage()), + session: await createInMemorySession(), model: registration.getModel(), steeringMode: "one-at-a-time", }); @@ -400,7 +410,7 @@ describe("AgentHarness", () => { return fauxAssistantMessage("ok"); }, ]); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); const harness = new AgentHarness({ models, session, @@ -443,7 +453,7 @@ describe("AgentHarness", () => { ]); const harness = new AgentHarness({ models, - session: new Session(new InMemorySessionStorage()), + session: await createInMemorySession(), model: registration.getModel(), }); const queueUpdates: Array<{ steer: number; followUp: number; nextTurn: number }> = []; @@ -495,7 +505,7 @@ describe("AgentHarness", () => { ]); const harness = new AgentHarness({ models, - session: new Session(new InMemorySessionStorage()), + session: await createInMemorySession(), model: registration.getModel(), followUpMode: "one-at-a-time", }); @@ -521,7 +531,7 @@ describe("AgentHarness", () => { it("settles thrown hook failures with persisted assistant error messages", async () => { const registration = newFaux(); registration.setResponses([() => fauxAssistantMessage("should not be used")]); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); const harness = new AgentHarness({ models, session, @@ -582,7 +592,7 @@ describe("AgentHarness", () => { ]); const harness = new AgentHarness({ models, - session: new Session(new InMemorySessionStorage()), + session: await createInMemorySession(), model: registration.getModel(), thinkingLevel: "off", resources: { @@ -615,7 +625,7 @@ describe("AgentHarness", () => { it("orders pending listener session writes after agent-emitted messages", async () => { const registration = newFaux(); registration.setResponses([() => fauxAssistantMessage("ok")]); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); const harness = new AgentHarness({ models, session, @@ -648,7 +658,7 @@ describe("AgentHarness", () => { const barrier = deferred(); const harness = new AgentHarness({ models, - session: new Session(new InMemorySessionStorage()), + session: await createInMemorySession(), model: registration.getModel(), }); let listenerFinished = false; @@ -681,7 +691,7 @@ describe("AgentHarness", () => { stopReason: "toolUse", }), ]); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); const toolUsage = createUsage(1, 2, 3, 4); const patchedToolUsage = createUsage(5, 6, 7, 8); const calculateToolWithUsage = createCalculateToolWithUsage(toolUsage); @@ -748,7 +758,7 @@ describe("AgentHarness", () => { }; const harness = new AgentHarness({ models, - session: new Session(new InMemorySessionStorage()), + session: await createInMemorySession(), model: registration.getModel(), tools: [contextTool], toolContext, @@ -785,7 +795,7 @@ describe("AgentHarness", () => { let generation = 0; const harness = new AgentHarness({ models, - session: new Session(new InMemorySessionStorage()), + session: await createInMemorySession(), model: registration.getModel(), tools: [contextTool], toolContext: async (): Promise => ({ generation: ++generation }), @@ -799,7 +809,7 @@ describe("AgentHarness", () => { it("persists generated compaction usage", async () => { const registration = newFaux(); registration.setResponses([fauxAssistantMessage("## Goal\nTest summary")]); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); await session.appendMessage(createUserMessage("one")); await session.appendMessage(createAssistantMessage("two")); const harness = new AgentHarness({ @@ -818,7 +828,7 @@ describe("AgentHarness", () => { it("persists hook-provided compaction usage", async () => { const registration = newFaux(); const usage = createUsage(5, 6, 7, 8); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); await session.appendMessage(createUserMessage("one")); await session.appendMessage(createAssistantMessage("two")); const harness = new AgentHarness({ @@ -856,7 +866,7 @@ describe("AgentHarness", () => { return fauxAssistantMessage("## Goal\nRecovered summary"); }, ]); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); await session.appendMessage(createUserMessage("one")); await session.appendMessage(createAssistantMessage("two")); const harness = new AgentHarness({ @@ -896,7 +906,7 @@ describe("AgentHarness", () => { return fauxAssistantMessage("", { stopReason: "error", errorMessage: "insufficient_quota" }); }, ]); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); await session.appendMessage(createUserMessage("one")); await session.appendMessage(createAssistantMessage("two")); const harness = new AgentHarness({ @@ -931,7 +941,7 @@ describe("AgentHarness", () => { return fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }); }), ); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); await session.appendMessage(createUserMessage("one")); await session.appendMessage(createAssistantMessage("two")); const harness = new AgentHarness({ @@ -978,7 +988,7 @@ describe("AgentHarness", () => { return fauxAssistantMessage("## Goal\nRecovered branch summary"); }, ]); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); const targetId = await session.appendMessage(createUserMessage("first branch")); await session.appendMessage(createAssistantMessage("first reply")); await session.appendMessage(createUserMessage("abandoned work")); @@ -1015,7 +1025,7 @@ describe("AgentHarness", () => { it("persists generated branch summary usage", async () => { const registration = newFaux(); registration.setResponses([fauxAssistantMessage("## Goal\nBranch summary")]); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); const targetId = await session.appendMessage(createUserMessage("first branch")); await session.appendMessage(createAssistantMessage("first reply")); await session.appendMessage(createUserMessage("abandoned work")); @@ -1034,7 +1044,7 @@ describe("AgentHarness", () => { it("persists hook-provided branch summary usage", async () => { const registration = newFaux(); const usage = createUsage(13, 14, 15, 16); - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); const targetId = await session.appendMessage(createUserMessage("first branch")); await session.appendMessage(createAssistantMessage("first reply")); await session.appendMessage(createUserMessage("abandoned work")); @@ -1054,7 +1064,7 @@ describe("AgentHarness", () => { }); it("preserves app tool types for getters and update events", async () => { - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); const model = getModel("anthropic", "claude-sonnet-4-5"); type AppTool = AgentTool & { source: "builtin" | "extension" }; const inspectTool: AppTool = { ...calculateTool, name: "inspect", source: "builtin" }; @@ -1123,8 +1133,8 @@ describe("AgentHarness", () => { expect((await session.buildContext()).activeToolNames).toEqual(["search"]); }); - it("validates constructor tool names", () => { - const session = new Session(new InMemorySessionStorage()); + it("validates constructor tool names", async () => { + const session = await createInMemorySession(); const model = getModel("anthropic", "claude-sonnet-4-5"); expect( () => new AgentHarness({ session, models, model, tools: [calculateTool], activeToolNames: ["missing"] }), @@ -1152,7 +1162,7 @@ describe("AgentHarness", () => { }); it("preserves app resource types for getters and update events", async () => { - const session = new Session(new InMemorySessionStorage()); + const session = await createInMemorySession(); const model = getModel("anthropic", "claude-sonnet-4-5"); const harness = new AgentHarness({ session, diff --git a/packages/agent/test/harness/repo.test.ts b/packages/agent/test/harness/repo.test.ts index 57a7da11e..d34f6f998 100644 --- a/packages/agent/test/harness/repo.test.ts +++ b/packages/agent/test/harness/repo.test.ts @@ -1,15 +1,43 @@ -import { existsSync } from "node:fs"; +import { existsSync, writeFileSync } from "node:fs"; import { describe, expect, it } from "vitest"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; -import { createJsonlSessionStore } from "../../src/harness/session/jsonl-repo.ts"; +import { createJsonlSessionStore } from "../../src/harness/session/jsonl-store.ts"; import { createInMemorySessionStore, type InMemorySessionCreateOptions, -} from "../../src/harness/session/memory-repo.ts"; -import { createSessionRepository } from "../../src/harness/session/repo-utils.ts"; +} from "../../src/harness/session/memory-store.ts"; +import { createSessionRepository } from "../../src/harness/session/repository.ts"; import type { SessionMetadata, SessionStore } from "../../src/harness/types.ts"; import { createAssistantMessage, createTempDir, createUserMessage } from "./session-test-utils.ts"; +class CountingReadEnv extends NodeExecutionEnv { + readCount = 0; + + override async readTextFile(path: string, abortSignal?: AbortSignal) { + this.readCount += 1; + return super.readTextFile(path, abortSignal); + } +} + +function createDeferred(): { promise: Promise; resolve: () => void } { + let resolve!: () => void; + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise; + }); + return { promise, resolve }; +} + +class BlockingAppendEnv extends NodeExecutionEnv { + readonly appendStarted = createDeferred(); + readonly releaseAppend = createDeferred(); + + override async appendFile(path: string, content: string | Uint8Array) { + this.appendStarted.resolve(); + await this.releaseAppend.promise; + return super.appendFile(path, content); + } +} + function createCountingInMemorySessionStore(): { store: SessionStore; counter: { loadCount: number }; @@ -25,12 +53,10 @@ function createCountingInMemorySessionStore(): { return source.load(metadata); }, list: (options) => source.list(options), - getEntries: (metadata, options) => source.getEntries(metadata, options), - createEntryId: (metadata) => source.createEntryId(metadata), appendEntry: (metadata, entry) => source.appendEntry(metadata, entry), - setLeafId: (metadata, leafId) => source.setLeafId(metadata, leafId), delete: (metadata) => source.delete(metadata), - fork: (metadata, options) => source.fork(metadata, options), + fork: (metadata, options, entries) => source.fork(metadata, options, entries), + [Symbol.asyncDispose]: () => source[Symbol.asyncDispose](), }, }; } @@ -53,26 +79,77 @@ describe("InMemorySessionStore", () => { await expect(repo.open(metadata)).rejects.toThrow("Session not found: session-1"); }); - it("does not repeatedly load full snapshots for scoped reads", async () => { + it("retains the opened aggregate instead of reloading for scoped reads", async () => { const { store, counter } = createCountingInMemorySessionStore(); const repo = createSessionRepository({ store }); const session = await repo.create({ id: "session-1" }); const entryId = await session.appendMessage(createUserMessage("one")); - counter.loadCount = 0; await session.getMetadata(); - expect(counter.loadCount).toBe(0); - await session.getLeafId(); - expect(counter.loadCount).toBe(1); - - counter.loadCount = 0; await session.getEntry(entryId); - expect(counter.loadCount).toBe(1); + expect(counter.loadCount).toBe(0); + }); + + it("rejects repository operations and session writes after store disposal", async () => { + const store = createInMemorySessionStore(); + const repo = createSessionRepository({ store }); + const session = await repo.create({ id: "session-1" }); + await store[Symbol.asyncDispose](); + + await expect(repo.list()).rejects.toThrow("In-memory session store is disposed"); + await expect(session.appendMessage(createUserMessage("late"))).rejects.toThrow( + "In-memory session store is disposed", + ); + }); + + it("supports lexical ownership with await using", async () => { + let listDisposedStore: (() => Promise) | undefined; + { + await using store = createInMemorySessionStore(); + const repository = createSessionRepository({ store }); + listDisposedStore = () => repository.list(); + await repository.create({ id: "session-1" }); + } + + await expect(listDisposedStore!()).rejects.toThrow("In-memory session store is disposed"); }); }); describe("JsonlSessionStore", () => { + it("waits for accepted appends before disposal and rejects later writes", async () => { + const root = createTempDir(); + const env = new BlockingAppendEnv({ cwd: root }); + const store = createJsonlSessionStore({ fs: env, sessionsRoot: root }); + const repo = createSessionRepository({ store }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const append = session.appendMessage(createUserMessage("accepted")); + await env.appendStarted.promise; + let closed = false; + const dispose = store[Symbol.asyncDispose]().then(() => { + closed = true; + }); + await Promise.resolve(); + expect(closed).toBe(false); + env.releaseAppend.resolve(); + await append; + await dispose; + await expect(session.appendMessage(createUserMessage("late"))).rejects.toThrow("JSONL session store is disposed"); + }); + + it("parses once when opened and retains state across appends", async () => { + const root = createTempDir(); + const env = new CountingReadEnv({ cwd: root }); + const repo = createSessionRepository({ store: createJsonlSessionStore({ fs: env, sessionsRoot: root }) }); + const created = await repo.create({ cwd: root, id: "session-1" }); + const metadata = await created.getMetadata(); + await created.appendMessage(createUserMessage("initial")); + env.readCount = 0; + const opened = await repo.open(metadata); + for (let i = 0; i < 10; i++) await opened.appendMessage(createUserMessage(`message ${i}`)); + expect(env.readCount).toBe(1); + }); + it("stores sessions below encoded cwd directories and lists by cwd", async () => { const root = createTempDir(); const env = new NodeExecutionEnv({ cwd: root }); @@ -92,6 +169,17 @@ describe("JsonlSessionStore", () => { ); }); + it("fails loudly when listing a malformed session file", async () => { + const root = createTempDir(); + const env = new NodeExecutionEnv({ cwd: root }); + const repo = createSessionRepository({ store: createJsonlSessionStore({ fs: env, sessionsRoot: root }) }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const metadata = await session.getMetadata(); + writeFileSync(metadata.path, "not json\n"); + + await expect(repo.list()).rejects.toMatchObject({ code: "invalid_session" }); + }); + it("opens, deletes, and forks by metadata", async () => { const root = createTempDir(); const env = new NodeExecutionEnv({ cwd: root }); diff --git a/packages/agent/test/harness/session-backends.test.ts b/packages/agent/test/harness/session-backends.test.ts new file mode 100644 index 000000000..aac9785d4 --- /dev/null +++ b/packages/agent/test/harness/session-backends.test.ts @@ -0,0 +1,177 @@ +import { existsSync, readFileSync, writeFileSync } from "node:fs"; +import { describe, expect, it } from "vitest"; +import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; +import { createJsonlSessionStore, loadJsonlSessionMetadata } from "../../src/harness/session/jsonl-store.ts"; +import { createInMemorySessionStore } from "../../src/harness/session/memory-store.ts"; +import { createSessionRepository } from "../../src/harness/session/repository.ts"; +import type { Session } from "../../src/harness/session/session.ts"; +import { createAssistantMessage, createTempDir, createUserMessage } from "./session-test-utils.ts"; + +async function appendUsageEntries(session: Session) { + const assistant = createAssistantMessage("reply"); + if (assistant.role !== "assistant") throw new Error("Expected assistant test message"); + assistant.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 }, + }; + await session.appendMessage(assistant); + await session.appendCompaction("summary", undefined, 1234, undefined, undefined, { + 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(await session.getLeafId(), { + 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 }, + }, + }); +} + +function createInMemoryRepository() { + return createSessionRepository({ store: createInMemorySessionStore() }); +} + +function createJsonlRepository(options: Parameters[0]) { + return createSessionRepository({ store: createJsonlSessionStore(options) }); +} + +describe("Session aggregate", () => { + it("owns leaf navigation, labels, names, stats, and branch traversal", async () => { + const repo = createInMemoryRepository(); + const session = await repo.create({ id: "session-1" }); + const root = await session.appendMessage(createUserMessage("root")); + const child = await session.appendMessage(createAssistantMessage("child")); + await session.appendLabel(root, "checkpoint"); + await session.appendSessionName(" review\nname "); + + expect(await session.getLeafId()).not.toBeNull(); + expect(await session.getLabel(root)).toBe("checkpoint"); + expect(await session.getSessionName()).toBe("review name"); + expect((await session.getBranch(child)).map((entry) => entry.id)).toEqual([root, child]); + + await session.moveTo(root); + expect(await session.getLeafId()).toBe(root); + expect((await session.getEntries()).at(-1)).toMatchObject({ type: "leaf", targetId: root }); + await expect(session.moveTo("missing")).rejects.toThrow("Entry missing not found"); + }); + + it("serializes concurrent appends into one parent chain", async () => { + const session = await createInMemoryRepository().create({}); + await Promise.all( + Array.from({ length: 20 }, (_, index) => session.appendMessage(createUserMessage(`message ${index}`))), + ); + const entries = await session.getEntries(); + expect(entries).toHaveLength(20); + for (let index = 0; index < entries.length; index++) { + expect(entries[index]!.parentId).toBe(index === 0 ? null : entries[index - 1]!.id); + } + }); + + it("includes assistant and summary usage in statistics", async () => { + const session = await createInMemoryRepository().create({}); + await appendUsageEntries(session); + expect(await session.getSessionStats()).toEqual({ + messageCount: 1, + cachedTokens: 40, + uncachedTokens: 68, + totalTokens: 136, + costTotal: 1.36, + }); + }); + + it("stops branch traversal at retained-tail compaction", async () => { + const session = await createInMemoryRepository().create({}); + await session.appendMessage(createUserMessage("root")); + const child = await session.appendMessage(createAssistantMessage("child")); + const compaction = await session.appendCompaction("summary", child, 1234, undefined, undefined, undefined, [ + createAssistantMessage("child"), + ]); + const tail = await session.appendMessage(createUserMessage("after")); + expect((await session.getBranch(tail)).map((entry) => entry.id)).toEqual([compaction, tail]); + }); +}); + +describe("JsonlSessionStore", () => { + it("writes headers and entries and reopens the aggregate", async () => { + const root = createTempDir(); + const repo = createJsonlRepository({ fs: new NodeExecutionEnv({ cwd: root }), sessionsRoot: root }); + const session = await repo.create({ cwd: root, id: "session-1", metadata: { profile: "reviewer" } }); + const metadata = await session.getMetadata(); + const entryId = await session.appendMessage(createUserMessage("one")); + const reopened = await repo.open(metadata); + + expect(existsSync(metadata.path)).toBe(true); + expect((await reopened.getEntries()).map((entry) => entry.id)).toEqual([entryId]); + expect((await loadJsonlSessionMetadata(new NodeExecutionEnv({ cwd: root }), metadata.path)).metadata).toEqual({ + profile: "reviewer", + }); + const lines = readFileSync(metadata.path, "utf8").trim().split("\n"); + expect(JSON.parse(lines[0]!)).toMatchObject({ type: "session", version: 3, id: "session-1" }); + expect(JSON.parse(lines[1]!)).toMatchObject({ id: entryId, type: "message" }); + }); + + it("fails loudly for malformed headers and entries", async () => { + const root = createTempDir(); + const env = new NodeExecutionEnv({ cwd: root }); + const repo = createJsonlRepository({ fs: env, sessionsRoot: root }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const metadata = await session.getMetadata(); + writeFileSync(metadata.path, "not json\n"); + await expect(repo.open(metadata)).rejects.toMatchObject({ code: "invalid_session" }); + + const header = { type: "session", version: 3, id: "session-1", timestamp: metadata.createdAt, cwd: root }; + writeFileSync(metadata.path, `${JSON.stringify(header)}\nnot json\n`); + await expect(repo.open(metadata)).rejects.toMatchObject({ code: "invalid_entry" }); + }); + + it("enforces entry uniqueness and does not recreate deleted files", async () => { + const root = createTempDir(); + const store = createJsonlSessionStore({ fs: new NodeExecutionEnv({ cwd: root }), sessionsRoot: root }); + const snapshot = await store.create({ cwd: root, id: "session-1" }); + const entry = { + type: "message" as const, + id: "entry-1", + parentId: null, + timestamp: "2026-01-01T00:00:00.000Z", + message: createUserMessage("one"), + }; + await store.appendEntry(snapshot.metadata, entry); + await expect(store.appendEntry(snapshot.metadata, entry)).rejects.toThrow("Entry entry-1 already exists"); + await store.delete(snapshot.metadata); + await expect(store.appendEntry(snapshot.metadata, { ...entry, id: "entry-2" })).rejects.toThrow( + "Session not found", + ); + expect(existsSync(snapshot.metadata.path)).toBe(false); + }); + + it("rejects non-object header metadata", async () => { + const root = createTempDir(); + const env = new NodeExecutionEnv({ cwd: root }); + const repo = createJsonlRepository({ fs: env, sessionsRoot: root }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const metadata = await session.getMetadata(); + const header = { + type: "session", + version: 3, + id: metadata.id, + timestamp: metadata.createdAt, + cwd: root, + metadata: "profile", + }; + writeFileSync(metadata.path, `${JSON.stringify(header)}\n`); + await expect(repo.open(metadata)).rejects.toThrow("session header metadata must be an object"); + }); +}); diff --git a/packages/agent/test/harness/session-test-utils.ts b/packages/agent/test/harness/session-test-utils.ts index 9a4f55a01..fa728e74f 100644 --- a/packages/agent/test/harness/session-test-utils.ts +++ b/packages/agent/test/harness/session-test-utils.ts @@ -3,6 +3,13 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import type { AgentMessage } from "@earendil-works/pi-agent-core"; import { afterEach } from "vitest"; +import { createInMemorySessionStore } from "../../src/harness/session/memory-store.ts"; +import { Session } from "../../src/harness/session/session.ts"; + +export async function createInMemorySession(id?: string): Promise { + const store = createInMemorySessionStore(); + return new Session(store, await store.create({ id })); +} export function createUserMessage(text: string): AgentMessage { return { @@ -41,10 +48,6 @@ export function createTempDir(): string { return dir; } -export function getLatestTempDir(): string { - return tempDirs[tempDirs.length - 1]!; -} - afterEach(() => { while (tempDirs.length > 0) { const dir = tempDirs.pop()!; diff --git a/packages/agent/test/harness/session.test.ts b/packages/agent/test/harness/session.test.ts index 70adbb386..485aabed7 100644 --- a/packages/agent/test/harness/session.test.ts +++ b/packages/agent/test/harness/session.test.ts @@ -2,11 +2,14 @@ import { readFileSync } from "node:fs"; import { join } from "node:path"; import { describe, expect, it } from "vitest"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; -import { JsonlSessionStorage } from "../../src/harness/session/jsonl-storage.ts"; -import { InMemorySessionStorage } from "../../src/harness/session/memory-storage.ts"; -import { type ContextEntryTransform, Session } from "../../src/harness/session/session.ts"; -import type { SessionStorage } from "../../src/harness/types.ts"; -import { createAssistantMessage, createTempDir, createUserMessage, getLatestTempDir } from "./session-test-utils.ts"; +import { createJsonlSessionStore } from "../../src/harness/session/jsonl-store.ts"; +import { createInMemorySessionStore } from "../../src/harness/session/memory-store.ts"; +import { + type ContextEntryTransform, + Session, + type SessionContextBuildOptions, +} from "../../src/harness/session/session.ts"; +import { createAssistantMessage, createTempDir, createUserMessage } from "./session-test-utils.ts"; function getTextData(data: unknown): string { if (typeof data !== "object" || data === null || !("text" in data)) { @@ -16,14 +19,15 @@ function getTextData(data: unknown): string { return typeof value === "string" ? value : ""; } -async function runSessionSuite( - name: string, - createStorage: () => SessionStorage | Promise, - inspect?: () => void, -) { +interface SessionFixture { + createSession(options?: SessionContextBuildOptions): Session; + reloadSession(options?: SessionContextBuildOptions): Promise; +} + +async function runSessionSuite(name: string, createFixture: () => Promise, inspect?: () => void) { describe(name, () => { it("appends messages and builds context in order", async () => { - const session = new Session(await createStorage()); + const session = (await createFixture()).createSession(); await session.appendMessage(createUserMessage("one")); await session.appendMessage(createAssistantMessage("two")); const context = await session.buildContext(); @@ -31,7 +35,7 @@ async function runSessionSuite( }); it("tracks model and thinking level changes", async () => { - const session = new Session(await createStorage()); + const session = (await createFixture()).createSession(); await session.appendMessage(createUserMessage("one")); await session.appendModelChange("openai", "gpt-4.1"); await session.appendThinkingLevelChange("high"); @@ -41,7 +45,7 @@ async function runSessionSuite( }); it("supports branching by moving the leaf and appending a new branch", async () => { - const session = new Session(await createStorage()); + const session = (await createFixture()).createSession(); const user1 = await session.appendMessage(createUserMessage("one")); const assistant1 = await session.appendMessage(createAssistantMessage("two")); await session.appendMessage(createUserMessage("three")); @@ -55,7 +59,7 @@ async function runSessionSuite( }); it("supports moving the leaf to root", async () => { - const session = new Session(await createStorage()); + const session = (await createFixture()).createSession(); await session.appendMessage(createUserMessage("one")); await session.moveTo(null); expect(await session.getLeafId()).toBeNull(); @@ -63,7 +67,7 @@ async function runSessionSuite( }); it("reconstructs compaction summaries in context", async () => { - const session = new Session(await createStorage()); + const session = (await createFixture()).createSession(); await session.appendMessage(createUserMessage("one")); await session.appendMessage(createAssistantMessage("two")); const user2 = await session.appendMessage(createUserMessage("three")); @@ -85,7 +89,7 @@ async function runSessionSuite( }); it("supports moving with branch summary entries in context", async () => { - const session = new Session(await createStorage()); + const session = (await createFixture()).createSession(); const user1 = await session.appendMessage(createUserMessage("one")); const summaryId = await session.moveTo(user1, { summary: "summary text" }); expect(summaryId).toBeTruthy(); @@ -96,7 +100,7 @@ async function runSessionSuite( }); it("persists compaction usage", async () => { - const session = new Session(await createStorage()); + const session = (await createFixture()).createSession(); const firstKeptEntryId = await session.appendMessage(createUserMessage("one")); const usage = { input: 1, @@ -121,7 +125,7 @@ async function runSessionSuite( }); it("persists branch summary usage", async () => { - const session = new Session(await createStorage()); + const session = (await createFixture()).createSession(); const user1 = await session.appendMessage(createUserMessage("one")); const usage = { input: 1, @@ -139,7 +143,7 @@ async function runSessionSuite( }); it("supports custom message entries in context", async () => { - const session = new Session(await createStorage()); + const session = (await createFixture()).createSession(); await session.appendMessage(createUserMessage("one")); await session.appendCustomMessageEntry("custom", "hello", true, { ok: true }); const context = await session.buildContext(); @@ -147,7 +151,7 @@ async function runSessionSuite( }); it("keeps custom entries in context entries but omits them from messages by default", async () => { - const session = new Session(await createStorage()); + const session = (await createFixture()).createSession(); await session.appendMessage(createUserMessage("one")); await session.appendCustomEntry("chat_message", { text: "hello" }); const contextEntries = await session.buildContextEntries(); @@ -157,7 +161,7 @@ async function runSessionSuite( }); it("projects custom entries with configured custom-entry projectors", async () => { - const session = new Session(await createStorage(), { + const session = (await createFixture()).createSession({ entryProjectors: { chat_message: (entry) => [createUserMessage(`chat: ${getTextData(entry.data)}`)], }, @@ -175,7 +179,7 @@ async function runSessionSuite( observedFirstEntryType = entries[0]?.type; return entries.filter((entry) => entry.type !== "compaction"); }; - const session = new Session(await createStorage(), { entryTransforms: [dropCompaction] }); + const session = (await createFixture()).createSession({ entryTransforms: [dropCompaction] }); await session.appendMessage(createUserMessage("one")); const kept = await session.appendMessage(createUserMessage("two")); await session.appendCompaction("summary", kept, 1234); @@ -186,13 +190,13 @@ async function runSessionSuite( }); it("normalizes session names", async () => { - const session = new Session(await createStorage()); + const session = (await createFixture()).createSession(); await session.appendSessionName(" hello\nworld\r\nagain "); expect(await session.getSessionName()).toBe("hello world again"); }); it("supports labels and session info entries without affecting context", async () => { - const session = new Session(await createStorage()); + const session = (await createFixture()).createSession(); const user1 = await session.appendMessage(createUserMessage("one")); await session.appendLabel(user1, "checkpoint"); await session.appendSessionName("name"); @@ -205,20 +209,20 @@ async function runSessionSuite( }); it("rejects labels for missing entries", async () => { - const session = new Session(await createStorage()); + const session = (await createFixture()).createSession(); await expect(session.appendLabel("missing", "checkpoint")).rejects.toThrow("Entry missing not found"); }); - it("persists leaf changes and appended entries via storage", async () => { - const storage = await createStorage(); - const session = new Session(storage); + it("persists leaf changes and appended entries through the store", async () => { + const fixture = await createFixture(); + const session = fixture.createSession(); const user1 = await session.appendMessage(createUserMessage("one")); await session.appendMessage(createAssistantMessage("two")); await session.appendLabel(user1, "checkpoint"); await session.appendSessionName("name"); await session.moveTo(user1); await session.appendMessage(createAssistantMessage("branched")); - const session2 = new Session(storage); + const session2 = await fixture.reloadSession(); const context = await session2.buildContext(); expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]); expect(await session2.getLabel(user1)).toBe("checkpoint"); @@ -228,19 +232,32 @@ async function runSessionSuite( }); } -runSessionSuite("Session with in-memory storage", () => new InMemorySessionStorage()); +runSessionSuite("Session with in-memory store", async () => { + const store = createInMemorySessionStore(); + const snapshot = await store.create({}); + return { + createSession: (options) => new Session(store, snapshot, options), + reloadSession: async (options) => new Session(store, await store.load(snapshot.metadata), options), + }; +}); +let jsonlSessionPath = ""; runSessionSuite( - "Session with JSONL storage", + "Session with JSONL store", async () => { const dir = createTempDir(); const env = new NodeExecutionEnv({ cwd: dir }); - return await JsonlSessionStorage.create(env, join(dir, "session.jsonl"), { cwd: dir, sessionId: "session-1" }); + const store = createJsonlSessionStore({ fs: env, sessionsRoot: join(dir, "sessions") }); + const snapshot = await store.create({ cwd: dir, id: "session-1" }); + const metadata = snapshot.metadata; + jsonlSessionPath = metadata.path; + return { + createSession: (options) => new Session(store, snapshot, options), + reloadSession: async (options) => new Session(store, await store.load(snapshot.metadata), options), + }; }, () => { - const dir = getLatestTempDir(); - const filePath = join(dir, "session.jsonl"); - const lines = readFileSync(filePath, "utf8").trim().split("\n"); + const lines = readFileSync(jsonlSessionPath, "utf8").trim().split("\n"); expect(lines.length).toBeGreaterThan(1); const header = JSON.parse(lines[0]!); expect(header.type).toBe("session"); diff --git a/packages/agent/test/harness/sqlite-migrations.test.ts b/packages/agent/test/harness/sqlite-migrations.test.ts index 17c521b2d..1cf5446b4 100644 --- a/packages/agent/test/harness/sqlite-migrations.test.ts +++ b/packages/agent/test/harness/sqlite-migrations.test.ts @@ -10,11 +10,11 @@ import { type SqliteDatabaseFactory, type SqliteRunResult, type SqliteSessionMetadata, - SqliteSessionStorage, type SqliteStatement, } from "../../../storage/sqlite-node/src/index.ts"; +import { SqliteSessionConnection } from "../../../storage/sqlite-node/src/sqlite/storage/index.ts"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; -import { createSessionRepository } from "../../src/harness/session/repo-utils.ts"; +import { createSessionRepository } from "../../src/harness/session/repository.ts"; import { createAssistantMessage, createUserMessage } from "./session-test-utils.ts"; function createTempDir(): string { @@ -96,7 +96,8 @@ describe("SQLite migrations", () => { const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); const sqlite = createNodeSqliteFactory(); - const repo = createSessionRepository({ store: createSqliteSessionStore({ env, sqlite, databasePath }) }); + const store = createSqliteSessionStore({ env, sqlite, databasePath }); + const repo = createSessionRepository({ store }); await repo.create({ cwd: root, id: "session-1" }); const db = await sqlite.open(databasePath); @@ -160,12 +161,53 @@ describe("SQLite migrations", () => { expect((await overridden.getMetadata()).metadata).toEqual({ profile: "writer" }); }); + it("rolls back the entire fork when copying an entry fails", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const store = createSqliteSessionStore({ env, sqlite, databasePath }); + const repo = createSessionRepository({ store }); + const source = await repo.create({ cwd: root, id: "source" }); + await source.appendMessage(createUserMessage("one")); + await source.appendMessage(createAssistantMessage("two")); + + const db = await sqlite.open(databasePath); + try { + await db.exec(` +CREATE TRIGGER fail_fork_entry BEFORE INSERT ON session_entries +WHEN new.session_id = 'fork' AND new.entry_seq = 2 +BEGIN + SELECT RAISE(ABORT, 'fail fork'); +END; +`); + } finally { + await db.close(); + } + + await expect(repo.fork(await source.getMetadata(), { cwd: root, id: "fork" })).rejects.toMatchObject({ + code: "storage", + }); + const inspection = await sqlite.open(databasePath); + try { + expect( + await inspection.prepare("SELECT id FROM sessions WHERE id = ?").get<{ id: string }>("fork"), + ).toBeUndefined(); + expect( + await inspection.prepare("SELECT id FROM session_entries WHERE session_id = ?").all<{ id: string }>("fork"), + ).toEqual([]); + } finally { + await inspection.close(); + } + }); + 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 = createSessionRepository({ store: createSqliteSessionStore({ env, sqlite, databasePath }) }); + const store = createSqliteSessionStore({ env, sqlite, databasePath }); + const repo = createSessionRepository({ store }); const session = await repo.create({ cwd: root, id: "session-1" }); const rootId = await session.appendMessage(createUserMessage("root")); const childId = await session.appendMessage(createAssistantMessage("child")); @@ -201,7 +243,8 @@ describe("SQLite migrations", () => { const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); const sqlite = createNodeSqliteFactory(); - const repo = createSessionRepository({ store: createSqliteSessionStore({ env, sqlite, databasePath }) }); + const store = createSqliteSessionStore({ env, sqlite, databasePath }); + const repo = createSessionRepository({ store }); 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")); @@ -280,11 +323,12 @@ describe("SQLite migrations", () => { open: async () => db, }; const env = new NodeExecutionEnv({ cwd: root }); - const repo = createSessionRepository({ - store: createSqliteSessionStore({ env, sqlite, databasePath: join(root, "sessions.sqlite") }), - }); + const store = createSqliteSessionStore({ env, sqlite, databasePath: join(root, "sessions.sqlite") }); + const repo = createSessionRepository({ store }); await expect(repo.create({ cwd: root, id: "session-1" })).rejects.toThrow("insert failed"); + expect(db.closeCount).toBe(0); + await store[Symbol.asyncDispose](); expect(db.closeCount).toBe(1); }); @@ -300,9 +344,8 @@ describe("SQLite migrations", () => { open: async () => db, }; const env = new NodeExecutionEnv({ cwd: root }); - const repo = createSessionRepository({ - store: createSqliteSessionStore({ env, sqlite, databasePath: join(root, "sessions.sqlite") }), - }); + const store = createSqliteSessionStore({ env, sqlite, databasePath: join(root, "sessions.sqlite") }); + const repo = createSessionRepository({ store }); const metadata: SqliteSessionMetadata = { id: "missing", createdAt: new Date().toISOString(), @@ -312,22 +355,66 @@ describe("SQLite migrations", () => { writeFileSync(metadata.path, ""); await expect(repo.open(metadata)).rejects.toThrow("Session not found: missing"); + expect(db.closeCount).toBe(0); + await store[Symbol.asyncDispose](); expect(db.closeCount).toBe(1); }); - it("closes the source storage after fork reads its entries", async () => { + it("retains one connection for repeated session operations", async () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); const { sqlite, counts } = createCloseCountingSqliteFactory(); - const repo = createSessionRepository({ store: createSqliteSessionStore({ env, sqlite, databasePath }) }); + const store = createSqliteSessionStore({ env, sqlite, databasePath }); + const repo = createSessionRepository({ store }); + + const session = await repo.create({ cwd: root, id: "session-1" }); + for (let i = 0; i < 10; i++) await session.appendMessage(createUserMessage(`message ${i}`)); + await session.getEntries(); + expect(counts).toEqual({ opens: 1, closes: 0 }); + await store[Symbol.asyncDispose](); + expect(counts).toEqual({ opens: 1, closes: 1 }); + await store[Symbol.asyncDispose](); + expect(counts).toEqual({ opens: 1, closes: 1 }); + }); + + it("shares one connection across source and fork until the store is disposed", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const { sqlite, counts } = createCloseCountingSqliteFactory(); + const store = createSqliteSessionStore({ env, sqlite, databasePath }); + const repo = createSessionRepository({ store }); const source = await repo.create({ cwd: root, id: "session-1" }); - counts.opens = 0; - counts.closes = 0; - await repo.fork(await source.getMetadata(), { cwd: root, id: "session-2" }); + const fork = await repo.fork(await source.getMetadata(), { cwd: root, id: "session-2" }); + await fork.appendMessage(createUserMessage("fork")); + expect(counts).toEqual({ opens: 1, closes: 0 }); + await store[Symbol.asyncDispose](); + expect(counts).toEqual({ opens: 1, closes: 1 }); + }); - expect(counts).toEqual({ opens: 2, closes: 2 }); + it("fails loudly when a stored entry cannot be decoded", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const store = createSqliteSessionStore({ env, sqlite, databasePath }); + const repo = createSessionRepository({ store }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const entryId = await session.appendMessage(createUserMessage("message")); + const metadata = await session.getMetadata(); + + const db = await sqlite.open(databasePath); + try { + await db + .prepare("UPDATE session_entries SET payload = ? WHERE session_id = ? AND id = ?") + .run("not json", metadata.id, entryId); + } finally { + await db.close(); + } + + await expect(repo.open(metadata)).rejects.toMatchObject({ code: "invalid_entry" }); }); it("restores in-memory state when appendEntry fails after mutating caches", async () => { @@ -336,7 +423,7 @@ describe("SQLite migrations", () => { const sqlite = createNodeSqliteFactory(); const db = await sqlite.open(databasePath); await applyMigrations(db); - const storage = await SqliteSessionStorage.create(db, databasePath, { + const storage = await SqliteSessionConnection.create(db, databasePath, { cwd: root, sessionId: "session-1", }); @@ -370,7 +457,8 @@ describe("SQLite migrations", () => { const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); const sqlite = createNodeSqliteFactory(); - const repo = createSessionRepository({ store: createSqliteSessionStore({ env, sqlite, databasePath }) }); + const store = createSqliteSessionStore({ env, sqlite, databasePath }); + const repo = createSessionRepository({ store }); const session = await repo.create({ cwd: root, id: "session-1" }); const userId = await session.appendMessage(createUserMessage("one")); await session.appendThinkingLevelChange("high"); diff --git a/packages/agent/test/harness/sqlite-node.test.ts b/packages/agent/test/harness/sqlite-node.test.ts index 2c51462ce..30f2f6032 100644 --- a/packages/agent/test/harness/sqlite-node.test.ts +++ b/packages/agent/test/harness/sqlite-node.test.ts @@ -1,35 +1,52 @@ import { join } from "node:path"; -import { describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it } from "vitest"; import { createNodeSqliteFactory, - createSqliteSessionRepository, createSqliteSessionSearch, createSqliteSessionStore, - type SqliteSessionCreateOptions, - type SqliteSessionListOptions, type SqliteSessionMetadata, } from "../../../storage/sqlite-node/src/index.ts"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; -import { createJsonlSessionRepository, createJsonlSessionStore } from "../../src/harness/session/jsonl-repo.ts"; -import { createSessionRepository } from "../../src/harness/session/repo-utils.ts"; +import { createJsonlSessionStore } from "../../src/harness/session/jsonl-store.ts"; +import { createSessionRepository } from "../../src/harness/session/repository.ts"; +import { createScanningSessionSearch } from "../../src/harness/session/search-backend.ts"; +import { rebuildSessionSearchIndex } from "../../src/harness/session/search-index.ts"; import type { - JsonlSessionCreateOptions, - JsonlSessionListOptions, JsonlSessionMetadata, SessionSearch, SessionSearchHit, SessionSearchIndex, SessionSearchOptions, - SessionStore, SessionTreeEntry, } from "../../src/harness/types.ts"; import { createTempDir, createUserMessage } from "./session-test-utils.ts"; +const ownedStores: AsyncDisposable[] = []; + +afterEach(async () => { + for (const store of ownedStores.splice(0)) await store[Symbol.asyncDispose](); +}); + +function createSqliteRepository(options: Parameters[0]) { + const store = createSqliteSessionStore(options); + ownedStores.push(store); + return createSessionRepository({ + store, + search: createSqliteSessionSearch({ ...options, mode: "canonical" }), + }); +} + +function createJsonlRepository(options: Parameters[0]) { + const store = createJsonlSessionStore(options); + ownedStores.push(store); + return createSessionRepository({ store, search: createScanningSessionSearch(store) }); +} + describe("JsonlSessionStore with scanning search", () => { it("searches canonical session entries by scanning", async () => { const root = createTempDir(); const env = new NodeExecutionEnv({ cwd: root }); - const repo = createJsonlSessionRepository({ fs: env, sessionsRoot: join(root, "sessions") }); + const repo = createJsonlRepository({ fs: env, sessionsRoot: join(root, "sessions") }); const included = await repo.create({ cwd: root, id: "included" }); const excluded = await repo.create({ cwd: `${root}/other`, id: "excluded" }); const entryId = await included.appendMessage(createUserMessage("Find the auth defect")); @@ -42,12 +59,12 @@ describe("JsonlSessionStore with scanning search", () => { }); describe("SqliteSessionStore with explicit SQLite FTS5 search", () => { - it("uses SQLite FTS5 when composed with its search backend", async () => { + it("uses SQLite FTS5 when composed with its search store", async () => { const root = createTempDir(); const env = new NodeExecutionEnv({ cwd: root }); const sqlite = createNodeSqliteFactory(); const databasePath = join(root, "sessions.sqlite"); - const repo = createSqliteSessionRepository({ env, sqlite, databasePath }); + const repo = createSqliteRepository({ env, sqlite, databasePath }); const included = await repo.create({ cwd: root, id: "included" }); const excluded = await repo.create({ cwd: `${root}/other`, id: "excluded" }); const metadata = await included.getMetadata(); @@ -76,10 +93,73 @@ describe("SqliteSessionStore with explicit SQLite FTS5 search", () => { await expect(repo.search({ text: "auth", cwd: root })).resolves.toEqual([]); }); + it("creates an empty canonical session without initializing FTS", async () => { + const root = createTempDir(); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const databasePath = join(root, "sessions.sqlite"); + const repo = createSqliteRepository({ env, sqlite, databasePath }); + const session = await repo.create({ cwd: root, id: "session-1" }); + + const db = await sqlite.open(databasePath); + try { + const fts = await db + .prepare("SELECT 1 AS found FROM sqlite_master WHERE type = 'table' AND name = 'session_search_fts'") + .get<{ found: number }>(); + expect(fts).toBeUndefined(); + } finally { + await db.close(); + } + await expect(session.appendMessage(createUserMessage("still writable"))).resolves.toBeTypeOf("string"); + }); + + it("rolls back canonical appends when co-located FTS trigger writes fail", async () => { + const root = createTempDir(); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const databasePath = join(root, "sessions.sqlite"); + const repo = createSqliteRepository({ env, sqlite, databasePath }); + await repo.search({ text: "initialize" }); + const session = await repo.create({ cwd: root, id: "session-1" }); + + const db = await sqlite.open(databasePath); + try { + await db.exec("DROP TABLE session_search_fts"); + } finally { + await db.close(); + } + + await expect(session.appendMessage(createUserMessage("must roll back"))).rejects.toThrow(); + await expect(session.getEntries()).resolves.toEqual([]); + }); + + it("rolls back canonical deletion when co-located FTS cleanup fails", async () => { + const root = createTempDir(); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const databasePath = join(root, "sessions.sqlite"); + const repo = createSqliteRepository({ env, sqlite, databasePath }); + await repo.search({ text: "initialize" }); + const session = await repo.create({ cwd: root, id: "session-1" }); + await session.appendMessage(createUserMessage("must remain")); + const metadata = await session.getMetadata(); + + const db = await sqlite.open(databasePath); + try { + await db.exec("DROP TABLE session_search_fts"); + } finally { + await db.close(); + } + + await expect(repo.delete(metadata)).rejects.toThrow(); + const reopened = await repo.open(metadata); + await expect(reopened.getEntries()).resolves.toHaveLength(1); + }); + it("initializes canonical storage when searched before the first session is created", async () => { const root = createTempDir(); const env = new NodeExecutionEnv({ cwd: root }); - const repo = createSqliteSessionRepository({ + const repo = createSqliteRepository({ env, sqlite: createNodeSqliteFactory(), databasePath: join(root, "sessions.sqlite"), @@ -96,51 +176,59 @@ describe("SqliteSessionStore with explicit SQLite FTS5 search", () => { }); }); -describe("JsonlSessionStore with SQLite search index", () => { - // This is not the intended production pairing; it exists to demonstrate that - // Canonical session storage and search are independently swappable/composable. - it("writes JSONL session entries into the configured SQLite search", async () => { +describe("Session search index rebuild", () => { + it("removes sessions that no longer exist in canonical storage", async () => { const root = createTempDir(); - const env = new NodeExecutionEnv({ cwd: root }); - const sqlite = createNodeSqliteFactory(); const search = createSqliteSessionSearch({ - env, - sqlite, + env: new NodeExecutionEnv({ cwd: root }), + sqlite: createNodeSqliteFactory(), databasePath: join(root, "search.sqlite"), }); - const canonical = createJsonlSessionStore({ fs: env, sessionsRoot: join(root, "sessions") }); - const store = { - load: (metadata) => canonical.load(metadata), - list: (options) => canonical.list(options), - getEntries: (metadata, options) => canonical.getEntries(metadata, options), - createEntryId: (metadata) => canonical.createEntryId(metadata), - async create(options) { - const metadata = await canonical.create(options); - await search.replaceSession(metadata, []); - return metadata; + const metadata: JsonlSessionMetadata = { + id: "deleted", + createdAt: "2026-01-01T00:00:00.000Z", + cwd: root, + path: join(root, "deleted.jsonl"), + }; + await search.upsertEntry(metadata, { + type: "message", + id: "entry-1", + parentId: null, + timestamp: metadata.createdAt, + message: createUserMessage("stale entry"), + }); + await rebuildSessionSearchIndex( + { + async list() { + return []; + }, + async load() { + throw new Error("load should not be called"); + }, }, - async appendEntry(metadata, entry) { - await canonical.appendEntry(metadata, entry); - await search.upsertEntry(metadata, entry); - }, - async setLeafId(metadata, leafId) { - const entry = await canonical.setLeafId(metadata, leafId); - await search.upsertEntry(metadata, entry); - return entry; - }, - async delete(metadata) { - await canonical.delete(metadata); - await search.deleteSession(metadata); - }, - async fork(source, options) { - const metadata = await canonical.fork(source, options); - await search.replaceSession(metadata, (await canonical.load(metadata)).entries); - return metadata; - }, - } satisfies SessionStore; + search, + ); + await expect(search.search({ text: "stale" })).resolves.toEqual([]); + }); +}); + +describe("JsonlSessionStore with SQLite search index", () => { + it("allows explicit projection into an independently configured SQLite index", async () => { + const root = createTempDir(); + const env = new NodeExecutionEnv({ cwd: root }); + const search = createSqliteSessionSearch({ + env, + sqlite: createNodeSqliteFactory(), + databasePath: join(root, "search.sqlite"), + }); + const store = createJsonlSessionStore({ fs: env, sessionsRoot: join(root, "sessions") }); const repo = createSessionRepository({ store, search }); const session = await repo.create({ cwd: root, id: "jsonl-session" }); + const metadata = await session.getMetadata(); const entryId = await session.appendMessage(createUserMessage("Find the auth defect")); + const entry = await session.getEntry(entryId); + expect(entry).toBeDefined(); + await search.upsertEntry(metadata, entry!); await expect(repo.search({ text: "auth", cwd: root })).resolves.toEqual([ expect.objectContaining({ entryId, metadata: expect.objectContaining({ id: "jsonl-session" }) }), @@ -149,7 +237,7 @@ describe("JsonlSessionStore with SQLite search index", () => { }); describe("JsonlSessionStore with multiple search indexes", () => { - it("queries one search implementation and composes index writes explicitly", async () => { + it("queries one search implementation and projects writes explicitly", async () => { const root = createTempDir(); const env = new NodeExecutionEnv({ cwd: root }); const sqlite = createNodeSqliteFactory(); @@ -163,45 +251,15 @@ describe("JsonlSessionStore with multiple search indexes", () => { sqlite, databasePath: join(root, "secondary-search.sqlite"), }); - const canonical = createJsonlSessionStore({ fs: env, sessionsRoot: join(root, "sessions") }); - const store = { - load: (metadata) => canonical.load(metadata), - list: (options) => canonical.list(options), - getEntries: (metadata, options) => canonical.getEntries(metadata, options), - createEntryId: (metadata) => canonical.createEntryId(metadata), - async create(options) { - const metadata = await canonical.create(options); - await primary.replaceSession(metadata, []); - await secondary.replaceSession(metadata, []); - return metadata; - }, - async appendEntry(metadata, entry) { - await canonical.appendEntry(metadata, entry); - await primary.upsertEntry(metadata, entry); - await secondary.upsertEntry(metadata, entry); - }, - async setLeafId(metadata, leafId) { - const entry = await canonical.setLeafId(metadata, leafId); - await primary.upsertEntry(metadata, entry); - await secondary.upsertEntry(metadata, entry); - return entry; - }, - async delete(metadata) { - await canonical.delete(metadata); - await primary.deleteSession(metadata); - await secondary.deleteSession(metadata); - }, - async fork(source, options) { - const metadata = await canonical.fork(source, options); - const entries = (await canonical.load(metadata)).entries; - await primary.replaceSession(metadata, entries); - await secondary.replaceSession(metadata, entries); - return metadata; - }, - } satisfies SessionStore; + const store = createJsonlSessionStore({ fs: env, sessionsRoot: join(root, "sessions") }); const repo = createSessionRepository({ store, search: primary }); const session = await repo.create({ cwd: root, id: "jsonl-session" }); + const metadata = await session.getMetadata(); const entryId = await session.appendMessage(createUserMessage("indexed in both places")); + const entry = await session.getEntry(entryId); + expect(entry).toBeDefined(); + await primary.upsertEntry(metadata, entry!); + await secondary.upsertEntry(metadata, entry!); await expect(repo.search({ text: "both" })).resolves.toEqual([ expect.objectContaining({ entryId, metadata: expect.objectContaining({ id: "jsonl-session" }) }), @@ -219,10 +277,14 @@ describe("SqliteSessionStore with custom search", () => { const removals: SqliteSessionMetadata[] = []; const searches: SessionSearchOptions[] = []; const index: SessionSearchIndex = { + async reset() { + upserts.length = 0; + removals.length = 0; + }, async upsertEntry(metadata, entry) { upserts.push({ metadata, entry }); }, - async replaceSession(metadata, entries: readonly SessionTreeEntry[]) { + async replaceSession(metadata, entries) { for (const entry of entries) upserts.push({ metadata, entry }); }, async deleteSession(metadata) { @@ -236,43 +298,22 @@ describe("SqliteSessionStore with custom search", () => { }, }; const env = new NodeExecutionEnv({ cwd: root }); - const sqlite = createNodeSqliteFactory(); - const databasePath = join(root, "sessions.sqlite"); - const canonical = createSqliteSessionStore({ env, sqlite, databasePath }); - const store = { - load: (metadata) => canonical.load(metadata), - list: (options) => canonical.list(options), - getEntries: (metadata, options) => canonical.getEntries(metadata, options), - createEntryId: (metadata) => canonical.createEntryId(metadata), - async create(options) { - const metadata = await canonical.create(options); - await index.replaceSession(metadata, []); - return metadata; - }, - async appendEntry(metadata, entry) { - await canonical.appendEntry(metadata, entry); - await index.upsertEntry(metadata, entry); - }, - async setLeafId(metadata, leafId) { - const entry = await canonical.setLeafId(metadata, leafId); - await index.upsertEntry(metadata, entry); - return entry; - }, - async delete(metadata) { - await canonical.delete(metadata); - await index.deleteSession(metadata); - }, - async fork(source, options) { - const metadata = await canonical.fork(source, options); - await index.replaceSession(metadata, (await canonical.load(metadata)).entries); - return metadata; - }, - } satisfies SessionStore; + const store = createSqliteSessionStore({ + env, + sqlite: createNodeSqliteFactory(), + databasePath: join(root, "sessions.sqlite"), + }); const repo = createSessionRepository({ store, search }); const session = await repo.create({ cwd: root, id: "session-1" }); const metadata = await session.getMetadata(); const entryId = await session.appendMessage(createUserMessage("indexed remotely")); + const entry = await session.getEntry(entryId); + expect(entry).toBeDefined(); + await index.upsertEntry(metadata, entry!); await session.moveTo(null); + const leafEntry = (await session.getEntries()).at(-1); + expect(leafEntry?.type).toBe("leaf"); + await index.upsertEntry(metadata, leafEntry!); expect(upserts).toContainEqual( expect.objectContaining({ metadata, entry: expect.objectContaining({ id: entryId }) }), @@ -282,6 +323,7 @@ describe("SqliteSessionStore with custom search", () => { ); await expect(repo.search({ text: "indexed remotely" })).resolves.toEqual([]); expect(searches).toEqual([{ text: "indexed remotely" }]); + await index.deleteSession(metadata); await repo.delete(metadata); expect(removals).toEqual([metadata]); }); diff --git a/packages/agent/test/harness/storage.test.ts b/packages/agent/test/harness/storage.test.ts deleted file mode 100644 index 50f4fc59a..000000000 --- a/packages/agent/test/harness/storage.test.ts +++ /dev/null @@ -1,504 +0,0 @@ -import { existsSync, readFileSync, writeFileSync } from "node:fs"; -import { join } from "node:path"; -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 BranchSummaryEntry, - type CompactionEntry, - type MessageEntry, - ok, - type SessionMetadata, -} from "../../src/harness/types.ts"; -import { createAssistantMessage, createTempDir, createUserMessage } from "./session-test-utils.ts"; - -describe("InMemorySessionStorage", () => { - it("returns configured session metadata", async () => { - const metadata: SessionMetadata = { id: "session-1", createdAt: "2026-01-01T00:00:00.000Z" }; - const storage = new InMemorySessionStorage({ metadata }); - expect(await storage.getMetadata()).toEqual(metadata); - }); - - it("copies initial entries and persists leaf changes", async () => { - const entry: MessageEntry = { - type: "message", - id: "entry-1", - parentId: null, - timestamp: "2026-01-01T00:00:00.000Z", - message: createUserMessage("one"), - }; - const initialEntries = [entry]; - const storage = new InMemorySessionStorage({ entries: initialEntries }); - initialEntries.push({ ...entry, id: "entry-2" }); - expect((await storage.getEntries()).map((storedEntry) => storedEntry.id)).toEqual(["entry-1"]); - expect(await storage.getLeafId()).toBe("entry-1"); - await storage.setLeafId(null); - expect(await storage.getLeafId()).toBeNull(); - expect((await storage.getEntries()).at(-1)).toMatchObject({ type: "leaf", targetId: null }); - }); - - it("rejects invalid leaf ids", async () => { - const storage = new InMemorySessionStorage(); - await expect(storage.setLeafId("missing")).rejects.toThrow("Entry missing not found"); - }); - - it("finds entries by type", async () => { - const entry: MessageEntry = { - type: "message", - id: "entry-1", - parentId: null, - timestamp: "2026-01-01T00:00:00.000Z", - message: createUserMessage("one"), - }; - const storage = new InMemorySessionStorage({ entries: [entry] }); - expect((await storage.findEntries("message")).map((found) => found.id)).toEqual(["entry-1"]); - expect(await storage.findEntries("session_info")).toEqual([]); - }); - - it("maintains label lookup", async () => { - const entry: MessageEntry = { - type: "message", - id: "entry-1", - parentId: null, - timestamp: "2026-01-01T00:00:00.000Z", - message: createUserMessage("one"), - }; - const storage = new InMemorySessionStorage({ entries: [entry] }); - expect(await storage.getLabel("entry-1")).toBeUndefined(); - await storage.appendEntry({ - type: "label", - id: "label-1", - parentId: "entry-1", - timestamp: "2026-01-01T00:00:01.000Z", - targetId: "entry-1", - label: "checkpoint", - }); - expect(await storage.getLabel("entry-1")).toBe("checkpoint"); - await storage.appendEntry({ - type: "label", - id: "label-2", - parentId: "label-1", - timestamp: "2026-01-01T00:00:02.000Z", - targetId: "entry-1", - label: undefined, - }); - expect(await storage.getLabel("entry-1")).toBeUndefined(); - }); - - 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", - parentId: null, - timestamp: "2026-01-01T00:00:00.000Z", - message: createUserMessage("root"), - }; - const child: MessageEntry = { - ...root, - id: "child", - parentId: "root", - message: createAssistantMessage("child"), - }; - 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([]); - }); -}); - -describe("JsonlSessionStorage", () => { - it("throws for missing files when opening", async () => { - const dir = createTempDir(); - const env = new NodeExecutionEnv({ cwd: dir }); - const filePath = join(dir, "session.jsonl"); - await expect(JsonlSessionStorage.open(env, filePath)).rejects.toMatchObject({ code: "not_found" }); - }); - - it("writes the header on create", 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" }); - expect(existsSync(filePath)).toBe(true); - expect(readFileSync(filePath, "utf8").trim().split("\n")).toHaveLength(1); - expect(await storage.getLeafId()).toBeNull(); - expect(await storage.getEntries()).toEqual([]); - await storage.appendEntry({ - type: "message", - id: "user-1", - parentId: null, - timestamp: "2026-01-01T00:00:00.000Z", - message: createUserMessage("one"), - }); - const lines = readFileSync(filePath, "utf8").trim().split("\n"); - expect(JSON.parse(lines[0]!).type).toBe("session"); - expect(JSON.parse(lines[1]!).id).toBe("user-1"); - expect(lines).toHaveLength(2); - }); - - it("throws for malformed session headers", async () => { - const dir = createTempDir(); - const env = new NodeExecutionEnv({ cwd: dir }); - const filePath = join(dir, "session.jsonl"); - writeFileSync(filePath, "not json\n"); - await expect(JsonlSessionStorage.open(env, filePath)).rejects.toThrow("first line is not a valid session header"); - }); - - it("throws for malformed entry lines", async () => { - const dir = createTempDir(); - const env = new NodeExecutionEnv({ cwd: dir }); - const filePath = join(dir, "session.jsonl"); - const header = { - type: "session", - version: 3, - id: "session-1", - timestamp: "2026-01-01T00:00:00.000Z", - cwd: dir, - }; - const entry: MessageEntry = { - type: "message", - id: "entry-1", - parentId: null, - timestamp: "2026-01-01T00:00:00.000Z", - message: createUserMessage("one"), - }; - writeFileSync(filePath, `${JSON.stringify(header)}\nnot json\n${JSON.stringify(entry)}\n`); - await expect(JsonlSessionStorage.open(env, filePath)).rejects.toMatchObject({ code: "invalid_entry" }); - }); - - it("creates and reads session metadata from the header", 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", - parentSessionPath: "/tmp/parent.jsonl", - }); - const metadata = await storage.getMetadata(); - expect(metadata).toMatchObject({ - id: "session-1", - cwd: dir, - path: filePath, - parentSessionPath: "/tmp/parent.jsonl", - }); - await storage.appendEntry({ - type: "message", - id: "user-1", - parentId: null, - timestamp: "2026-01-01T00:00:00.000Z", - message: createUserMessage("one"), - }); - expect(await loadJsonlSessionMetadata(env, filePath)).toEqual(metadata); - }); - - it("round-trips custom header metadata", 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", - metadata: { profile: "reviewer" }, - }); - expect((await storage.getMetadata()).metadata).toEqual({ profile: "reviewer" }); - const loaded = await JsonlSessionStorage.open(env, filePath); - expect((await loaded.getMetadata()).metadata).toEqual({ profile: "reviewer" }); - expect((await loadJsonlSessionMetadata(env, filePath)).metadata).toEqual({ profile: "reviewer" }); - }); - - it("omits header metadata when not provided", async () => { - const dir = createTempDir(); - const env = new NodeExecutionEnv({ cwd: dir }); - const filePath = join(dir, "session.jsonl"); - await JsonlSessionStorage.create(env, filePath, { cwd: dir, sessionId: "session-1" }); - expect(JSON.parse(readFileSync(filePath, "utf8").trim())).not.toHaveProperty("metadata"); - expect((await loadJsonlSessionMetadata(env, filePath)).metadata).toBeUndefined(); - }); - - it("throws for non-object header metadata", async () => { - const dir = createTempDir(); - const env = new NodeExecutionEnv({ cwd: dir }); - const filePath = join(dir, "session.jsonl"); - const header = { - type: "session", - version: 3, - id: "session-1", - timestamp: "2026-01-01T00:00:00.000Z", - cwd: dir, - metadata: "profile", - }; - writeFileSync(filePath, `${JSON.stringify(header)}\n`); - await expect(JsonlSessionStorage.open(env, filePath)).rejects.toThrow( - "session header metadata must be an object", - ); - }); - - it("loads existing entries and reconstructs leaf", 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" }); - const root: MessageEntry = { - type: "message", - id: "root", - parentId: null, - timestamp: "2026-01-01T00:00:00.000Z", - message: createUserMessage("root"), - }; - const child: MessageEntry = { - ...root, - id: "child", - parentId: "root", - message: createAssistantMessage("child"), - }; - await storage.appendEntry(root); - await storage.appendEntry(child); - const loaded = await JsonlSessionStorage.open(env, filePath); - expect(await loaded.getLeafId()).toBe("child"); - expect((await loaded.getEntries()).map((entry) => entry.id)).toEqual(["root", "child"]); - await loaded.setLeafId("root"); - 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.getPathToRootOrCompaction("child")).map((entry) => entry.id)).toEqual(["root", "child"]); - }); - - it("finds entries by type", 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: "entry-1", - parentId: null, - timestamp: "2026-01-01T00:00:00.000Z", - message: createUserMessage("one"), - }); - expect((await storage.findEntries("message")).map((found) => found.id)).toEqual(["entry-1"]); - expect(await storage.findEntries("session_info")).toEqual([]); - }); - - it("maintains label lookup", 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: "entry-1", - parentId: null, - timestamp: "2026-01-01T00:00:00.000Z", - message: createUserMessage("one"), - }); - expect(await storage.getLabel("entry-1")).toBeUndefined(); - await storage.appendEntry({ - type: "label", - id: "label-1", - parentId: "entry-1", - timestamp: "2026-01-01T00:00:01.000Z", - targetId: "entry-1", - label: "checkpoint", - }); - expect(await storage.getLabel("entry-1")).toBe("checkpoint"); - await storage.appendEntry({ - type: "label", - id: "label-2", - parentId: "label-1", - timestamp: "2026-01-01T00:00:02.000Z", - targetId: "entry-1", - label: undefined, - }); - expect(await storage.getLabel("entry-1")).toBeUndefined(); - const loaded = await JsonlSessionStorage.open(env, filePath); - 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"); - const header = { - type: "session", - version: 3, - id: "session-1", - timestamp: "2026-01-01T00:00:00.000Z", - cwd: dir, - }; - const metadata = await loadJsonlSessionMetadata( - { - readTextLines: async () => ok([JSON.stringify(header)]), - readTextFile: async () => { - throw new Error("readTextFile should not be called for metadata"); - }, - writeFile: async () => ok(undefined), - appendFile: async () => ok(undefined), - }, - filePath, - ); - expect(metadata).toEqual({ - id: "session-1", - createdAt: "2026-01-01T00:00:00.000Z", - cwd: dir, - path: filePath, - parentSessionPath: undefined, - }); - }); -}); diff --git a/packages/agent/test/scratch/simple.ts b/packages/agent/test/scratch/simple.ts index 87af9c79c..3727c6f28 100644 --- a/packages/agent/test/scratch/simple.ts +++ b/packages/agent/test/scratch/simple.ts @@ -4,18 +4,18 @@ import { createModels } from "@earendil-works/pi-ai"; import { cloudflareAIGatewayProvider } from "@earendil-works/pi-ai/providers/cloudflare-ai-gateway"; import { openaiProvider } from "@earendil-works/pi-ai/providers/openai"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; -import { InMemorySessionStorage } from "../../src/harness/session/memory-storage.ts"; import { AgentHarness, createBashTool, createEditTool, + createInMemorySessionStore, createReadTool, + createSessionRepository, createWriteTool, formatSkillsForSystemPrompt, loadSourcedPromptTemplates, loadSourcedSkills, type PromptTemplate, - Session, type Skill, } from "../../src/index.ts"; @@ -51,7 +51,8 @@ if (!model) { process.exit(-1); } -const session = new Session(new InMemorySessionStorage()); +await using store = createInMemorySessionStore(); +const session = await createSessionRepository({ store }).create({}); const agent = new AgentHarness({ session, models, diff --git a/packages/storage/sqlite-node/README.md b/packages/storage/sqlite-node/README.md index 8afafcf45..7f9757bef 100644 --- a/packages/storage/sqlite-node/README.md +++ b/packages/storage/sqlite-node/README.md @@ -2,4 +2,10 @@ Node sqlite storage backend for `@earendil-works/pi-agent-core` sessions. Provides the `node:sqlite` adapter (`SqliteDatabase` implementation) and the SQLite session -store/storage implementation (`createSqliteSessionStore`, migrations, materialized views). +store implementation (`createSqliteSessionStore`, migrations, materialized views). The store +lazily owns one shared database connection and implements `AsyncDisposable`. + +```ts +await using store = createSqliteSessionStore(options); +const repository = createSessionRepository({ store }); +``` diff --git a/packages/storage/sqlite-node/src/index.ts b/packages/storage/sqlite-node/src/index.ts index 399eadb61..05f822e3b 100644 --- a/packages/storage/sqlite-node/src/index.ts +++ b/packages/storage/sqlite-node/src/index.ts @@ -93,5 +93,5 @@ export function createNodeSqliteFactory(): SqliteDatabaseFactory { }; } -// Re-export the SQLite session storage backend and types so this package is a complete node-sqlite backend. +// Re-export the SQLite session 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 index 330150c7c..1a252ba46 100644 --- a/packages/storage/sqlite-node/src/sqlite/index.ts +++ b/packages/storage/sqlite-node/src/sqlite/index.ts @@ -1,5 +1,4 @@ export * from "./migrations.ts"; -export * from "./repo.ts"; export * from "./search-backend.ts"; -export * from "./storage/index.ts"; +export * from "./session-store.ts"; export * from "./types.ts"; diff --git a/packages/storage/sqlite-node/src/sqlite/repo.ts b/packages/storage/sqlite-node/src/sqlite/repo.ts deleted file mode 100644 index 21c132a37..000000000 --- a/packages/storage/sqlite-node/src/sqlite/repo.ts +++ /dev/null @@ -1,262 +0,0 @@ -import type { - LeafEntry, - SessionEntryCursorOptions, - SessionSnapshot, - SessionStorage, - SessionTreeEntry, -} from "@earendil-works/pi-agent-core"; -import { - createSessionId, - createSessionRepository, - getEntriesToFork, - getFileSystemResultOrThrow, - SessionError, - type SessionRepository, - type SessionStore, -} from "@earendil-works/pi-agent-core"; -import { applyMigrations } from "./migrations.ts"; -import { createSqliteSessionSearch } from "./search-backend.ts"; -import { SqliteSessionStorage } from "./storage/index.ts"; -import { rowToMetadata, type SessionRow } from "./storage/sessions.ts"; -import type { - SqliteDatabase, - SqliteDatabaseFactory, - SqliteSessionCreateOptions, - SqliteSessionListOptions, - SqliteSessionMetadata, - SqliteSessionStoreEnv, -} 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 type SqliteSessionStoreOptions = { - env: SqliteSessionStoreEnv; - sqlite: SqliteDatabaseFactory; - databasePath: string; -}; - -class SqliteSessionStore - implements SessionStore -{ - private readonly env: SqliteSessionStoreEnv; - private readonly sqlite: SqliteDatabaseFactory; - private readonly databasePathInput: string; - private databasePath: string | undefined; - - constructor(options: SqliteSessionStoreOptions) { - 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 await storage.getMetadata(); - } finally { - await db.close(); - } - } - - 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 { - return await SqliteSessionStorage.open(db, metadata); - } catch (error) { - await db.close(); - throw error; - } - } - - async load(metadata: SqliteSessionMetadata): Promise> { - const storage = await this.open(metadata); - try { - return { - metadata: await storage.getMetadata(), - leafId: await storage.getLeafId(), - entries: await storage.getEntries(), - }; - } finally { - await cleanupSessionStorage(storage); - } - } - - 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 getEntries(metadata: SqliteSessionMetadata, options?: SessionEntryCursorOptions): Promise { - const storage = await this.open(metadata); - try { - return await storage.getEntries(options); - } finally { - await cleanupSessionStorage(storage); - } - } - - async createEntryId(metadata: SqliteSessionMetadata): Promise { - const storage = await this.open(metadata); - try { - return await storage.createEntryId(); - } finally { - await cleanupSessionStorage(storage); - } - } - - async appendEntry(metadata: SqliteSessionMetadata, entry: SessionTreeEntry): Promise { - const storage = await this.open(metadata); - try { - await storage.appendEntry(entry); - } finally { - await cleanupSessionStorage(storage); - } - } - - async setLeafId(metadata: SqliteSessionMetadata, leafId: string | null): Promise { - const storage = await this.open(metadata); - try { - return await storage.setLeafId(leafId); - } finally { - await cleanupSessionStorage(storage); - } - } - - 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, options); - } finally { - await cleanupSessionStorage(source); - } - 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 await storage.getMetadata(); - } finally { - await db.close(); - } - } -} - -export function createSqliteSessionStore( - options: SqliteSessionStoreOptions, -): SessionStore { - return new SqliteSessionStore(options); -} - -export function createSqliteSessionRepository( - options: SqliteSessionStoreOptions, -): SessionRepository { - const store = createSqliteSessionStore(options); - return createSessionRepository({ - store, - search: createSqliteSessionSearch({ ...options, mode: "canonical" }), - }); -} diff --git a/packages/storage/sqlite-node/src/sqlite/search-backend.ts b/packages/storage/sqlite-node/src/sqlite/search-backend.ts index d87a60bce..d02f3ae61 100644 --- a/packages/storage/sqlite-node/src/sqlite/search-backend.ts +++ b/packages/storage/sqlite-node/src/sqlite/search-backend.ts @@ -128,6 +128,15 @@ class SqliteSessionSearch } } + async reset(): Promise { + const db = await this.openDatabase(); + try { + if (this.mode === "standalone") await db.prepare("DELETE FROM session_search_fts").run(); + } finally { + await db.close(); + } + } + async upsertEntry(metadata: TMetadata, entry: SessionTreeEntry): Promise { const db = await this.openDatabase(); try { diff --git a/packages/storage/sqlite-node/src/sqlite/session-store.ts b/packages/storage/sqlite-node/src/sqlite/session-store.ts new file mode 100644 index 000000000..4b0822ba6 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/session-store.ts @@ -0,0 +1,240 @@ +import type { SessionSnapshot, SessionStore, SessionTreeEntry } from "@earendil-works/pi-agent-core"; +import { createSessionId, getFileSystemResultOrThrow, SessionError } from "@earendil-works/pi-agent-core"; +import { applyMigrations } from "./migrations.ts"; +import { SqliteSessionConnection } from "./storage/index.ts"; +import { rowToMetadata, type SessionRow } from "./storage/sessions.ts"; +import type { + SqliteDatabase, + SqliteDatabaseFactory, + SqliteSessionCreateOptions, + SqliteSessionListOptions, + SqliteSessionMetadata, + SqliteSessionStoreEnv, +} 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"); +} + +export type SqliteSessionStoreOptions = { + env: SqliteSessionStoreEnv; + sqlite: SqliteDatabaseFactory; + databasePath: string; +}; + +class SerialOperationQueue { + private tail: Promise = Promise.resolve(); + + enqueue(operation: () => Promise | T): Promise { + const result = this.tail.then(operation); + this.tail = result.then( + () => undefined, + () => undefined, + ); + return result; + } + + async drain(): Promise { + await this.tail; + } +} + +class SqliteSessionStore + implements SessionStore +{ + private readonly env: SqliteSessionStoreEnv; + private readonly sqlite: SqliteDatabaseFactory; + private readonly databasePathInput: string; + private databasePath: string | undefined; + private databasePromise: Promise | undefined; + private database: SqliteDatabase | undefined; + private disposed = false; + private disposePromise: Promise | undefined; + private readonly operations = new SerialOperationQueue(); + private readonly writers = new Map(); + + constructor(options: SqliteSessionStoreOptions) { + this.env = options.env; + this.sqlite = options.sqlite; + this.databasePathInput = options.databasePath; + } + + create(options: SqliteSessionCreateOptions): Promise> { + this.assertOpen(); + return this.operations.enqueue(async () => { + const db = await this.getDatabase(); + const path = await this.getDatabasePath(); + const connection = await db.transaction(() => + SqliteSessionConnection.create(db, path, { + cwd: options.cwd, + sessionId: options.id ?? createSessionId(), + parentSessionId: options.parentSessionId, + metadata: options.metadata, + }), + ); + const metadata = await connection.getMetadata(); + this.writers.set(metadata.id, connection); + return { metadata, entries: [] }; + }); + } + + load(metadata: SqliteSessionMetadata): Promise> { + this.assertOpen(); + return this.operations.enqueue(() => this.loadSession(metadata)); + } + + private async loadSession(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 connection = + this.writers.get(metadata.id) ?? (await SqliteSessionConnection.open(await this.getDatabase(), metadata)); + this.writers.set(metadata.id, connection); + return { metadata: await connection.getMetadata(), entries: await connection.getEntries() }; + } + + list(options: SqliteSessionListOptions = {}): Promise { + this.assertOpen(); + return this.operations.enqueue(() => this.listSessions(options)); + } + + private async listSessions(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.getDatabase(); + 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)); + } + + appendEntry(metadata: SqliteSessionMetadata, entry: SessionTreeEntry): Promise { + this.assertOpen(); + return this.operations.enqueue(async () => { + const connection = + this.writers.get(metadata.id) ?? (await SqliteSessionConnection.open(await this.getDatabase(), metadata)); + this.writers.set(metadata.id, connection); + await connection.appendEntry(entry); + }); + } + + delete(metadata: SqliteSessionMetadata): Promise { + this.assertOpen(); + return this.operations.enqueue(async () => { + const db = await this.getDatabase(); + 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}`); + }); + this.writers.delete(metadata.id); + }); + } + + fork( + source: SqliteSessionMetadata, + options: SqliteSessionCreateOptions, + entries: readonly SessionTreeEntry[], + ): Promise> { + this.assertOpen(); + return this.operations.enqueue(async () => { + const db = await this.getDatabase(); + const connection = await db.transaction(async () => { + const connection = await SqliteSessionConnection.create(db, await this.getDatabasePath(), { + cwd: options.cwd, + sessionId: options.id ?? createSessionId(), + parentSessionId: options.parentSessionId ?? source.id, + metadata: options.metadata ?? source.metadata, + }); + for (const entry of entries) await connection.appendEntry(entry, { transaction: false }); + return connection; + }); + const metadata = await connection.getMetadata(); + this.writers.set(metadata.id, connection); + return { metadata, entries: [...entries] }; + }); + } + + async [Symbol.asyncDispose](): Promise { + if (!this.disposePromise) { + this.disposed = true; + this.disposePromise = this.finishDisposal(); + } + await this.disposePromise; + } + + private async finishDisposal(): Promise { + await this.operations.drain(); + const db = this.database ?? (this.databasePromise ? await this.databasePromise : undefined); + this.database = undefined; + this.databasePromise = undefined; + this.writers.clear(); + if (db) await db.close(); + } + + private assertOpen(): void { + if (this.disposed) throw new SessionError("storage", "SQLite session store is disposed"); + } + + private async getDatabasePath(): Promise { + this.databasePath ??= getFileSystemResultOrThrow( + await this.env.absolutePath(this.databasePathInput), + `Failed to resolve SQLite sessions database ${this.databasePathInput}`, + ); + return this.databasePath; + } + + private async getDatabase(): Promise { + if (!this.databasePromise) this.databasePromise = this.openDatabase(); + this.database = await this.databasePromise; + return this.database; + } + + private async openDatabase(): 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}`, + ); + const db = await this.sqlite.open(path); + try { + await configureSqliteDatabase(db); + await applyMigrations(db); + return db; + } catch (error) { + await db.close(); + throw error; + } + } +} + +export function createSqliteSessionStore( + options: SqliteSessionStoreOptions, +): SessionStore { + return new SqliteSessionStore(options); +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts b/packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts index 0ea601561..d526bdd63 100644 --- a/packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts +++ b/packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts @@ -1,4 +1,5 @@ import type { SessionTreeEntry } from "@earendil-works/pi-agent-core"; +import { toError } 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"; @@ -49,8 +50,8 @@ export async function getMaterializedBranchPathOrCompaction( if (entry.type !== "leaf") { entries.push(entry); } - } catch { - throw invalidSession(`invalid entry row for branch entry ${branchRow.entry_id}`); + } catch (error) { + throw invalidSession(`invalid entry row for branch entry ${branchRow.entry_id}`, toError(error)); } } return entries; diff --git a/packages/storage/sqlite-node/src/sqlite/storage/index.ts b/packages/storage/sqlite-node/src/sqlite/storage/index.ts index 833d5d7da..c6ddd6bc6 100644 --- a/packages/storage/sqlite-node/src/sqlite/storage/index.ts +++ b/packages/storage/sqlite-node/src/sqlite/storage/index.ts @@ -1,10 +1,5 @@ -import type { - LeafEntry, - SessionEntryCursorOptions, - SessionStorage, - SessionTreeEntry, -} from "@earendil-works/pi-agent-core"; -import { SessionError } from "@earendil-works/pi-agent-core"; +import type { LeafEntry, SessionEntryCursorOptions, SessionTreeEntry } from "@earendil-works/pi-agent-core"; +import { SessionError, toError } 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"; @@ -23,7 +18,7 @@ import { } 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"; +import { generateEntryId, invalidEntry, invalidSession, leafIdAfterEntry } from "./shared.ts"; async function decodeEntryRows(entryRows: SessionEntryRow[]): Promise<{ entries: SessionTreeEntry[]; @@ -36,8 +31,8 @@ async function decodeEntryRows(entryRows: SessionEntryRow[]): Promise<{ const entry = decodeEntry(entryRow); entries.push(entry); leafId = leafIdAfterEntry(entry); - } catch { - // Keep JSONL-like permissive resume behavior: skip malformed entries. + } catch (error) { + throw invalidEntry(`failed to decode entry ${entryRow.id}`, toError(error)); } } return { entries, leafId }; @@ -81,7 +76,7 @@ async function hasExistingChild(db: SqliteDatabase, sessionId: string, parentId: return row !== undefined; } -async function loadSqliteStorage( +async function loadSqliteSession( db: SqliteDatabase, sessionId: string, ): Promise<{ @@ -113,7 +108,7 @@ async function loadSqliteStorage( }; } -export class SqliteSessionStorage implements SessionStorage { +export class SqliteSessionConnection { private readonly db: SqliteDatabase; private readonly metadata: SqliteSessionMetadata; private byId: Map; @@ -199,9 +194,9 @@ export class SqliteSessionStorage implements SessionStorage { - const loaded = await loadSqliteStorage(db, metadata.id); - return new SqliteSessionStorage( + static async open(db: SqliteDatabase, metadata: SqliteSessionMetadata): Promise { + const loaded = await loadSqliteSession(db, metadata.id); + return new SqliteSessionConnection( db, rowToMetadata(loaded.row, metadata.path), null, @@ -220,7 +215,7 @@ export class SqliteSessionStorage implements SessionStorage; }, - ): Promise { + ): Promise { const createdAt = new Date().toISOString(); await db .prepare( @@ -238,7 +233,7 @@ export class SqliteSessionStorage implements SessionStorage { + async appendEntry(entry: SessionTreeEntry, options: { transaction?: boolean } = {}): Promise { const encoded = encodeEntry(entry); const previousMaterializedState: SessionMaterializedState = { ...this.materializedState, @@ -302,7 +297,7 @@ export class SqliteSessionStorage implements SessionStorage { + const write = async () => { const parentHadExistingChild = await hasExistingChild(this.db, this.metadata.id, entry.parentId); const nextSeq = await getNextSequence(this.db, this.metadata.id); await this.db @@ -334,7 +329,9 @@ export class SqliteSessionStorage implements SessionStorage; this.byId.set(entry.id, entry); entries.push(entry); - } catch { - // Keep JSONL-like permissive resume behavior: skip malformed entries. + } catch (error) { + throw invalidEntry(`failed to decode entry ${row.id}`, toError(error)); } } return entries; @@ -443,8 +440,4 @@ export class SqliteSessionStorage implements SessionStorage { - await this.db.close(); - } } diff --git a/scripts/browser-smoke-entry.ts b/scripts/browser-smoke-entry.ts index 64927b6e5..d8c0e175c 100644 --- a/scripts/browser-smoke-entry.ts +++ b/scripts/browser-smoke-entry.ts @@ -5,7 +5,8 @@ import { bashExecutionToText, convertToLlm, createCustomMessage, - createInMemorySessionRepository, + createInMemorySessionStore, + createSessionRepository, FileError, formatPromptTemplateInvocation, formatSkillInvocation, @@ -27,7 +28,7 @@ const stream = createAssistantMessageEventStream(); const agent = new Agent({ initialState: { model }, streamFn: streamSimple }); agent.steer({ role: "user", content: [{ type: "text", text: "queued" }], timestamp: 0 }); -const repo = createInMemorySessionRepository(); +const repo = createSessionRepository({ store: createInMemorySessionStore() }); const result = getOrThrow(ok({ value: 1 })); const customMessage = createCustomMessage("note", "hello", true, undefined, "2026-01-01T00:00:00.000Z"); const llmMessages = convertToLlm([customMessage]);