From 2700511e4594b1066f316f50fc0005df73a093f4 Mon Sep 17 00:00:00 2001 From: Christian Klotz Date: Fri, 31 Jul 2026 09:48:24 +0300 Subject: [PATCH] fix(agent): refine session repository API --- packages/agent/docs/harness.md | 2 +- .../agent/src/harness/session/jsonl-repo.ts | 8 +- .../agent/src/harness/session/memory-repo.ts | 10 +- .../agent/src/harness/session/repo-utils.ts | 10 +- packages/agent/src/harness/session/session.ts | 144 +++--------------- packages/agent/test/harness/repo.test.ts | 37 ++++- .../test/harness/sqlite-migrations.test.ts | 22 +-- .../agent/test/harness/sqlite-node.test.ts | 35 ++++- .../storage/sqlite-node/src/sqlite/repo.ts | 11 +- .../sqlite-node/src/sqlite/search-backend.ts | 49 +++--- scripts/browser-smoke-entry.ts | 4 +- 11 files changed, 143 insertions(+), 189 deletions(-) diff --git a/packages/agent/docs/harness.md b/packages/agent/docs/harness.md index b1e5d63a4..9136be695 100644 --- a/packages/agent/docs/harness.md +++ b/packages/agent/docs/harness.md @@ -1977,7 +1977,7 @@ type SessionForkOptions = /** New: the entire tree — all session entries, every branch, leaf preserved. */ | { scope: "tree" }; -interface SessionRepo { +interface SessionRepository { ... create(options: TCreateOptions): Promise; fork(source, options: SessionForkOptions & TCreateOptions): Promise; diff --git a/packages/agent/src/harness/session/jsonl-repo.ts b/packages/agent/src/harness/session/jsonl-repo.ts index 20d399982..b612837d5 100644 --- a/packages/agent/src/harness/session/jsonl-repo.ts +++ b/packages/agent/src/harness/session/jsonl-repo.ts @@ -17,7 +17,7 @@ import { createTimestamp, getEntriesToFork, getFileSystemResultOrThrow, - SessionRepo, + SessionRepository, } from "./repo-utils.ts"; import { ScanningSessionSearch } from "./search-backend.ts"; @@ -213,9 +213,9 @@ export function createJsonlSessionStore(options: JsonlSessionStoreOptions): Json return new JsonlSessionStore(options); } -export function createJsonlSessionRepo( +export function createJsonlSessionRepository( options: JsonlSessionStoreOptions, -): SessionRepo { +): SessionRepository { const store = createJsonlSessionStore(options); - return new SessionRepo({ store, search: new ScanningSessionSearch(store) }); + return new SessionRepository({ store, search: new ScanningSessionSearch(store) }); } diff --git a/packages/agent/src/harness/session/memory-repo.ts b/packages/agent/src/harness/session/memory-repo.ts index 21799ba7f..79035f185 100644 --- a/packages/agent/src/harness/session/memory-repo.ts +++ b/packages/agent/src/harness/session/memory-repo.ts @@ -9,7 +9,7 @@ import { type SessionTreeEntry, } from "../types.ts"; import { InMemorySessionStorage } from "./memory-storage.ts"; -import { createSessionId, createTimestamp, getEntriesToFork, SessionRepo } from "./repo-utils.ts"; +import { createSessionId, createTimestamp, getEntriesToFork, SessionRepository } from "./repo-utils.ts"; import { ScanningSessionSearch } from "./search-backend.ts"; export type InMemorySessionCreateOptions = { id?: string }; @@ -86,7 +86,11 @@ export function createInMemorySessionStore(): InMemorySessionStore { return new InMemorySessionStore(); } -export function createInMemorySessionRepo(): SessionRepo { +export function createInMemorySessionRepository(): SessionRepository< + SessionMetadata, + InMemorySessionCreateOptions, + void +> { const store = createInMemorySessionStore(); - return new SessionRepo({ store, search: new ScanningSessionSearch(store) }); + return new SessionRepository({ store, search: new ScanningSessionSearch(store) }); } diff --git a/packages/agent/src/harness/session/repo-utils.ts b/packages/agent/src/harness/session/repo-utils.ts index 7fca35d02..f4e3b290e 100644 --- a/packages/agent/src/harness/session/repo-utils.ts +++ b/packages/agent/src/harness/session/repo-utils.ts @@ -111,7 +111,7 @@ export function toStoreSession( const load = () => store.load(metadata); const storage: SessionStorage = { async getMetadata() { - return (await load()).metadata; + return metadata; }, async getLeafId() { return (await load()).leafId; @@ -144,7 +144,7 @@ export function toStoreSession( return new Session(storage); } -export class SessionRepo< +export class SessionRepository< TMetadata extends SessionMetadata = SessionMetadata, TCreateOptions extends SessionCreateOptions = SessionCreateOptions, TListOptions = void, @@ -186,15 +186,15 @@ export class SessionRepo< } } -export function createSessionRepo< +export function createSessionRepository< TMetadata extends SessionMetadata = SessionMetadata, TCreateOptions extends SessionCreateOptions = SessionCreateOptions, TListOptions = void, >(options: { store: SessionStore; search?: SessionSearch | null; -}): SessionRepo { - return new SessionRepo(options); +}): SessionRepository { + return new SessionRepository(options); } export function findSessionEntryMatches( diff --git a/packages/agent/src/harness/session/session.ts b/packages/agent/src/harness/session/session.ts index 51123be8c..57a7d8c9e 100644 --- a/packages/agent/src/harness/session/session.ts +++ b/packages/agent/src/harness/session/session.ts @@ -15,7 +15,6 @@ import type { SessionEntryCursorOptions, SessionInfoEntry, SessionMetadata, - SessionSnapshot, SessionStats, SessionStorage, SessionTreeEntry, @@ -149,138 +148,33 @@ export function buildSessionContext( return { ...state, messages }; } -interface SessionDependencies { - load(): Promise>; - getEntries(options?: SessionEntryCursorOptions): Promise; - createEntryId(): Promise; - appendEntry(entry: SessionTreeEntry): Promise; - setLeafId(leafId: string | null): Promise; -} - -function entriesById(entries: readonly SessionTreeEntry[]): Map { - return new Map(entries.map((entry) => [entry.id, entry])); -} - -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; -} - -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; - const trimmed = entry.label?.trim(); - label = trimmed || 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 storageToDependencies( - storage: SessionStorage, -): SessionDependencies { - return { - async load() { - return { - metadata: await storage.getMetadata(), - leafId: await storage.getLeafId(), - entries: await storage.getEntries(), - }; - }, - getEntries: (options) => storage.getEntries(options), - createEntryId: () => storage.createEntryId(), - appendEntry: (entry) => storage.appendEntry(entry), - setLeafId: (leafId) => storage.setLeafId(leafId), - }; -} - export class Session { - private readonly dependencies: SessionDependencies; + private readonly storage: SessionStorage; private readonly contextBuildOptions: SessionContextBuildOptions; constructor(storage: SessionStorage, contextBuildOptions: SessionContextBuildOptions = {}) { - this.dependencies = storageToDependencies(storage); + this.storage = storage; this.contextBuildOptions = contextBuildOptions; } - async getMetadata(): Promise { - return (await this.dependencies.load()).metadata; + getMetadata(): Promise { + return this.storage.getMetadata(); } - async getLeafId(): Promise { - return (await this.dependencies.load()).leafId; + getLeafId(): Promise { + return this.storage.getLeafId(); } - async getEntry(id: string): Promise { - return entriesById((await this.dependencies.load()).entries).get(id); + getEntry(id: string): Promise { + return this.storage.getEntry(id); } getEntries(options?: SessionEntryCursorOptions): Promise { - return this.dependencies.getEntries(options); + return this.storage.getEntries(options); } async getBranch(fromId?: string): Promise { - const state = await this.dependencies.load(); - return getPathToRootOrCompaction(state.entries, fromId ?? state.leafId); + return this.storage.getPathToRootOrCompaction(fromId ?? (await this.storage.getLeafId())); } async buildContextEntries(options: SessionContextBuildOptions = {}): Promise { @@ -301,28 +195,28 @@ export class Session { }; } - async getLabel(id: string): Promise { - return getLabel((await this.dependencies.load()).entries, id); + getLabel(id: string): Promise { + return this.storage.getLabel(id); } - async getSessionStats(): Promise { - return getSessionStats((await this.dependencies.load()).entries); + getSessionStats(): Promise { + return this.storage.getSessionStats(); } - async getSessionName(): Promise { - return getSessionName((await this.dependencies.load()).entries); + getSessionName(): Promise { + return this.storage.getSessionName(); } private async appendEntry(entry: SessionTreeEntry): Promise { - await this.dependencies.appendEntry(entry); + await this.storage.appendEntry(entry); } private setLeafId(leafId: string | null): Promise { - return this.dependencies.setLeafId(leafId); + return this.storage.setLeafId(leafId); } private async createEntryId(): Promise { - return this.dependencies.createEntryId(); + return this.storage.createEntryId(); } private async appendTypedEntry(entry: TEntry): Promise { diff --git a/packages/agent/test/harness/repo.test.ts b/packages/agent/test/harness/repo.test.ts index f1323c912..fc8e06853 100644 --- a/packages/agent/test/harness/repo.test.ts +++ b/packages/agent/test/harness/repo.test.ts @@ -3,12 +3,21 @@ import { describe, expect, it } from "vitest"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; import { JsonlSessionStore } from "../../src/harness/session/jsonl-repo.ts"; import { InMemorySessionStore } from "../../src/harness/session/memory-repo.ts"; -import { SessionRepo } from "../../src/harness/session/repo-utils.ts"; +import { SessionRepository } from "../../src/harness/session/repo-utils.ts"; import { createAssistantMessage, createTempDir, createUserMessage } from "./session-test-utils.ts"; +class CountingInMemorySessionStore extends InMemorySessionStore { + loadCount = 0; + + override async load(...args: Parameters) { + this.loadCount += 1; + return super.load(...args); + } +} + describe("InMemorySessionStore", () => { it("opens, deletes, and forks by metadata", async () => { - const repo = new SessionRepo({ store: new InMemorySessionStore() }); + const repo = new SessionRepository({ store: new InMemorySessionStore() }); const session = await repo.create({ id: "session-1" }); const metadata = await session.getMetadata(); const user1 = await session.appendMessage(createUserMessage("one")); @@ -23,6 +32,24 @@ describe("InMemorySessionStore", () => { await repo.delete(metadata); await expect(repo.open(metadata)).rejects.toThrow("Session not found: session-1"); }); + + it("does not repeatedly load full snapshots for scoped reads", async () => { + const store = new CountingInMemorySessionStore(); + const repo = new SessionRepository({ store }); + const session = await repo.create({ id: "session-1" }); + const entryId = await session.appendMessage(createUserMessage("one")); + + store.loadCount = 0; + await session.getMetadata(); + expect(store.loadCount).toBe(0); + + await session.getLeafId(); + expect(store.loadCount).toBe(1); + + store.loadCount = 0; + await session.getEntry(entryId); + expect(store.loadCount).toBe(1); + }); }); describe("JsonlSessionStore", () => { @@ -31,7 +58,7 @@ describe("JsonlSessionStore", () => { const env = new NodeExecutionEnv({ cwd: root }); const cwd = "/tmp/my-project"; const otherCwd = "/tmp/other-project"; - const repo = new SessionRepo({ store: new JsonlSessionStore({ fs: env, sessionsRoot: root }) }); + const repo = new SessionRepository({ store: new JsonlSessionStore({ fs: env, sessionsRoot: root }) }); const session = await repo.create({ cwd, id: "019de8c2-de29-73e9-ae0c-e134db34c447" }); const otherSession = await repo.create({ cwd: otherCwd, id: "other-session" }); const metadata = await session.getMetadata(); @@ -48,7 +75,7 @@ describe("JsonlSessionStore", () => { it("opens, deletes, and forks by metadata", async () => { const root = createTempDir(); const env = new NodeExecutionEnv({ cwd: root }); - const repo = new SessionRepo({ store: new JsonlSessionStore({ fs: env, sessionsRoot: root }) }); + const repo = new SessionRepository({ store: new JsonlSessionStore({ fs: env, sessionsRoot: root }) }); const source = await repo.create({ cwd: "/tmp/source", id: "source-session" }); const sourceMetadata = await source.getMetadata(); const user1 = await source.appendMessage(createUserMessage("one")); @@ -70,7 +97,7 @@ describe("JsonlSessionStore", () => { it("persists header metadata through create, list, and fork", async () => { const root = createTempDir(); const env = new NodeExecutionEnv({ cwd: root }); - const repo = new SessionRepo({ store: new JsonlSessionStore({ fs: env, sessionsRoot: root }) }); + const repo = new SessionRepository({ store: new JsonlSessionStore({ fs: env, sessionsRoot: root }) }); const source = await repo.create({ cwd: "/tmp/source", id: "source-session", diff --git a/packages/agent/test/harness/sqlite-migrations.test.ts b/packages/agent/test/harness/sqlite-migrations.test.ts index f442a804b..62069db92 100644 --- a/packages/agent/test/harness/sqlite-migrations.test.ts +++ b/packages/agent/test/harness/sqlite-migrations.test.ts @@ -14,7 +14,7 @@ import { type SqliteStatement, } from "../../../storage/sqlite-node/src/index.ts"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; -import { SessionRepo } from "../../src/harness/session/repo-utils.ts"; +import { SessionRepository } from "../../src/harness/session/repo-utils.ts"; import { createAssistantMessage, createUserMessage } from "./session-test-utils.ts"; function createTempDir(): string { @@ -70,7 +70,7 @@ describe("SQLite migrations", () => { const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); const sqlite = createNodeSqliteFactory(); - const repo = new SessionRepo({ store: new SqliteSessionStore({ env, sqlite, databasePath }) }); + const repo = new SessionRepository({ store: new SqliteSessionStore({ env, sqlite, databasePath }) }); await repo.create({ cwd: root, id: "session-1" }); const db = await sqlite.open(databasePath); @@ -112,7 +112,7 @@ describe("SQLite migrations", () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); - const repo = new SessionRepo({ + const repo = new SessionRepository({ store: new SqliteSessionStore({ env, sqlite: createNodeSqliteFactory(), databasePath }), }); const source = await repo.create({ @@ -139,7 +139,7 @@ describe("SQLite migrations", () => { const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); const sqlite = createNodeSqliteFactory(); - const repo = new SessionRepo({ store: new SqliteSessionStore({ env, sqlite, databasePath }) }); + const repo = new SessionRepository({ store: new SqliteSessionStore({ env, sqlite, databasePath }) }); const session = await repo.create({ cwd: root, id: "session-1" }); const rootId = await session.appendMessage(createUserMessage("root")); const childId = await session.appendMessage(createAssistantMessage("child")); @@ -175,7 +175,7 @@ describe("SQLite migrations", () => { const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); const sqlite = createNodeSqliteFactory(); - const repo = new SessionRepo({ store: new SqliteSessionStore({ env, sqlite, databasePath }) }); + const repo = new SessionRepository({ store: new SqliteSessionStore({ env, sqlite, databasePath }) }); const session = await repo.create({ cwd: root, id: "session-1" }); const rootId = await session.appendMessage(createUserMessage("root")); const firstChildId = await session.appendMessage(createAssistantMessage("first child")); @@ -203,7 +203,7 @@ describe("SQLite migrations", () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); - const repo = new SessionRepo({ + const repo = new SessionRepository({ store: new SqliteSessionStore({ env, sqlite: createNodeSqliteFactory(), databasePath }), }); const session = await repo.create({ cwd: root, id: "session-1" }); @@ -225,7 +225,7 @@ describe("SQLite migrations", () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); - const repo = new SessionRepo({ + const repo = new SessionRepository({ store: new SqliteSessionStore({ env, sqlite: createNodeSqliteFactory(), databasePath }), }); const session = await repo.create({ cwd: root, id: "session-1" }); @@ -254,7 +254,7 @@ describe("SQLite migrations", () => { open: async () => db, }; const env = new NodeExecutionEnv({ cwd: root }); - const repo = new SessionRepo({ + const repo = new SessionRepository({ store: new SqliteSessionStore({ env, sqlite, databasePath: join(root, "sessions.sqlite") }), }); @@ -274,7 +274,7 @@ describe("SQLite migrations", () => { open: async () => db, }; const env = new NodeExecutionEnv({ cwd: root }); - const repo = new SessionRepo({ + const repo = new SessionRepository({ store: new SqliteSessionStore({ env, sqlite, databasePath: join(root, "sessions.sqlite") }), }); const metadata: SqliteSessionMetadata = { @@ -294,7 +294,7 @@ describe("SQLite migrations", () => { const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); const storage = new SqliteSessionStore({ env, sqlite: createNodeSqliteFactory(), databasePath }); - const repo = new SessionRepo({ store: storage }); + const repo = new SessionRepository({ store: storage }); let cleanupCount = 0; const sourceStorage = { async getEntries() { @@ -367,7 +367,7 @@ describe("SQLite migrations", () => { const databasePath = join(root, "sessions.sqlite"); const env = new NodeExecutionEnv({ cwd: root }); const sqlite = createNodeSqliteFactory(); - const repo = new SessionRepo({ store: new SqliteSessionStore({ env, sqlite, databasePath }) }); + const repo = new SessionRepository({ store: new SqliteSessionStore({ env, sqlite, databasePath }) }); 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 ba8c6c5b7..24b3a372e 100644 --- a/packages/agent/test/harness/sqlite-node.test.ts +++ b/packages/agent/test/harness/sqlite-node.test.ts @@ -2,15 +2,15 @@ import { join } from "node:path"; import { describe, expect, it } from "vitest"; import { createNodeSqliteFactory, - createSqliteSessionRepo, + createSqliteSessionRepository, type SqliteSessionMetadata, SqliteSessionSearch, SqliteSessionStore, type SqliteSessionStoreApi, } from "../../../storage/sqlite-node/src/index.ts"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; -import { createJsonlSessionRepo, JsonlSessionStore } from "../../src/harness/session/jsonl-repo.ts"; -import { SessionRepo } from "../../src/harness/session/repo-utils.ts"; +import { createJsonlSessionRepository, JsonlSessionStore } from "../../src/harness/session/jsonl-repo.ts"; +import { SessionRepository } from "../../src/harness/session/repo-utils.ts"; import type { JsonlSessionMetadata, JsonlSessionStoreApi, @@ -26,7 +26,7 @@ describe("JsonlSessionStore with scanning search", () => { it("searches canonical session entries by scanning", async () => { const root = createTempDir(); const env = new NodeExecutionEnv({ cwd: root }); - const repo = createJsonlSessionRepo({ fs: env, sessionsRoot: join(root, "sessions") }); + const repo = createJsonlSessionRepository({ 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")); @@ -44,7 +44,7 @@ describe("SqliteSessionStore with explicit SQLite FTS5 search", () => { const env = new NodeExecutionEnv({ cwd: root }); const sqlite = createNodeSqliteFactory(); const databasePath = join(root, "sessions.sqlite"); - const repo = createSqliteSessionRepo({ env, sqlite, databasePath }); + const repo = createSqliteSessionRepository({ 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(); @@ -72,6 +72,25 @@ describe("SqliteSessionStore with explicit SQLite FTS5 search", () => { await repo.delete(metadata); await expect(repo.search({ text: "auth", cwd: root })).resolves.toEqual([]); }); + + 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({ + env, + sqlite: createNodeSqliteFactory(), + databasePath: join(root, "sessions.sqlite"), + }); + + await expect(repo.search({ text: "auth" })).resolves.toEqual([]); + const session = await repo.create({ cwd: root, id: "session-1" }); + const entryId = await session.appendMessage(createUserMessage("Find the auth defect")); + + await expect(repo.search({ text: "auth" })).resolves.toEqual([ + expect.objectContaining({ entryId, metadata: expect.objectContaining({ id: "session-1" }) }), + ]); + await expect(session.appendMessage(createUserMessage("Still writable"))).resolves.toBeTypeOf("string"); + }); }); describe("JsonlSessionStore with SQLite search index", () => { @@ -116,7 +135,7 @@ describe("JsonlSessionStore with SQLite search index", () => { return metadata; }, } satisfies JsonlSessionStoreApi; - const repo = new SessionRepo({ store, search }); + const repo = new SessionRepository({ store, search }); const session = await repo.create({ cwd: root, id: "jsonl-session" }); const entryId = await session.appendMessage(createUserMessage("Find the auth defect")); @@ -177,7 +196,7 @@ describe("JsonlSessionStore with multiple search indexes", () => { return metadata; }, } satisfies JsonlSessionStoreApi; - const repo = new SessionRepo({ store, search: primary }); + const repo = new SessionRepository({ store, search: primary }); const session = await repo.create({ cwd: root, id: "jsonl-session" }); const entryId = await session.appendMessage(createUserMessage("indexed in both places")); @@ -246,7 +265,7 @@ describe("SqliteSessionStore with custom search", () => { return metadata; }, } satisfies SqliteSessionStoreApi; - const repo = new SessionRepo({ store, search }); + const repo = new SessionRepository({ 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")); diff --git a/packages/storage/sqlite-node/src/sqlite/repo.ts b/packages/storage/sqlite-node/src/sqlite/repo.ts index dce903e0b..cf54400c9 100644 --- a/packages/storage/sqlite-node/src/sqlite/repo.ts +++ b/packages/storage/sqlite-node/src/sqlite/repo.ts @@ -10,7 +10,7 @@ import { getEntriesToFork, getFileSystemResultOrThrow, SessionError, - SessionRepo, + SessionRepository, } from "@earendil-works/pi-agent-core"; import { applyMigrations } from "./migrations.ts"; import { SqliteSessionSearch } from "./search-backend.ts"; @@ -246,9 +246,12 @@ export function createSqliteSessionStore(options: SqliteSessionStoreOptions): Sq return new SqliteSessionStore(options); } -export function createSqliteSessionRepo( +export function createSqliteSessionRepository( options: SqliteSessionStoreOptions, -): SessionRepo { +): SessionRepository { const store = createSqliteSessionStore(options); - return new SessionRepo({ store, search: new SqliteSessionSearch(options) }); + return new SessionRepository({ + store, + search: new SqliteSessionSearch({ ...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 8c7a10cdd..015d2be27 100644 --- a/packages/storage/sqlite-node/src/sqlite/search-backend.ts +++ b/packages/storage/sqlite-node/src/sqlite/search-backend.ts @@ -7,6 +7,7 @@ import type { SessionTreeEntry, } from "@earendil-works/pi-agent-core"; import { getFileSystemResultOrThrow } from "@earendil-works/pi-agent-core"; +import { applyMigrations } from "./migrations.ts"; import { rowToMetadata, type SessionRow } from "./storage/sessions.ts"; import type { SqliteDatabase, SqliteDatabaseFactory, SqliteSessionStoreEnv } from "./types.ts"; @@ -24,7 +25,15 @@ async function configureSqliteDatabase(db: SqliteDatabase): Promise { await db.exec("PRAGMA busy_timeout=5000"); } -type SearchSchemaMode = "standalone" | "canonical"; +export type SqliteSessionSearchMode = "standalone" | "canonical"; + +export interface SqliteSessionSearchOptions { + env: Pick; + sqlite: SqliteDatabaseFactory; + databasePath: string; + /** Defaults to a standalone index. Canonical mode shares and initializes a session-store database. */ + mode?: SqliteSessionSearchMode; +} async function tableExists(db: SqliteDatabase, name: string): Promise { return !!(await db @@ -32,8 +41,8 @@ async function tableExists(db: SqliteDatabase, name: string): Promise { .get<{ found: number }>(name)); } -async function ensureSearchSchema(db: SqliteDatabase): Promise { - if (await tableExists(db, "session_entries")) { +async function ensureSearchSchema(db: SqliteDatabase, mode: SqliteSessionSearchMode): Promise { + if (mode === "canonical") { const ftsExists = await tableExists(db, "session_search_fts"); await db.exec(` CREATE VIRTUAL TABLE IF NOT EXISTS session_search_fts USING fts5( @@ -54,7 +63,7 @@ CREATE TRIGGER IF NOT EXISTS session_search_fts_au AFTER UPDATE OF payload ON se END; `); if (!ftsExists) await db.exec("INSERT INTO session_search_fts(session_search_fts) VALUES('rebuild')"); - return "canonical"; + return; } await db.exec(` @@ -68,7 +77,6 @@ CREATE VIRTUAL TABLE IF NOT EXISTS session_search_fts USING fts5( tokenize = 'trigram remove_diacritics 1' ); `); - return "standalone"; } /** @@ -83,14 +91,12 @@ export class SqliteSessionSearch; - sqlite: SqliteDatabaseFactory; - databasePath: string; - }) { + constructor(options: SqliteSessionSearchOptions) { this.options = options; + this.mode = options.mode ?? "standalone"; } private async getDatabasePath(): Promise { @@ -103,7 +109,7 @@ export class SqliteSessionSearch { + private async openDatabase(): Promise { const path = await this.getDatabasePath(); const directory = getParentPath(path); getFileSystemResultOrThrow( @@ -113,8 +119,9 @@ export class SqliteSessionSearch { - const { db, schemaMode } = await this.openDatabase(); + const db = await this.openDatabase(); try { - if (schemaMode === "canonical") return; + if (this.mode === "canonical") return; const cwd = (metadata as { cwd?: unknown }).cwd; const sessionId = metadata.id; const entryId = entry.id; @@ -147,9 +154,9 @@ export class SqliteSessionSearch { - const { db, schemaMode } = await this.openDatabase(); + const db = await this.openDatabase(); try { - if (schemaMode === "canonical") return; + if (this.mode === "canonical") return; const cwd = (metadata as { cwd?: unknown }).cwd; const sessionId = metadata.id; const metadataJson = JSON.stringify(metadata); @@ -176,9 +183,9 @@ export class SqliteSessionSearch { - const { db, schemaMode } = await this.openDatabase(); + const db = await this.openDatabase(); try { - if (schemaMode === "canonical") return; + if (this.mode === "canonical") return; await db.prepare("DELETE FROM session_search_fts WHERE session_id = ?").run(metadata.id); } finally { await db.close(); @@ -188,10 +195,10 @@ export class SqliteSessionSearch[]> { const text = options.text.trim(); if (!text) return []; - const { db, schemaMode } = await this.openDatabase(); + const db = await this.openDatabase(); try { const query = `"${text.replaceAll('"', '""')}"`; - if (schemaMode === "canonical") { + if (this.mode === "canonical") { const rows = await db .prepare( "SELECT s.id, s.created_at, s.metadata, s.cwd, s.parent_session_id, s.active_leaf_id, se.id AS entry_id, se.timestamp, bm25(session_search_fts) AS score FROM session_search_fts JOIN session_entries se ON se.rowid = session_search_fts.rowid JOIN sessions s ON s.id = se.session_id WHERE session_search_fts MATCH ? AND (? IS NULL OR s.cwd = ?) ORDER BY score", diff --git a/scripts/browser-smoke-entry.ts b/scripts/browser-smoke-entry.ts index 7de806cfa..08856e178 100644 --- a/scripts/browser-smoke-entry.ts +++ b/scripts/browser-smoke-entry.ts @@ -13,7 +13,7 @@ import { InMemorySessionStore, ok, parseCommandArgs, - SessionRepo, + SessionRepository, streamProxy, toError, truncateHead, @@ -28,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 = new SessionRepo({ store: new InMemorySessionStore() }); +const repo = new SessionRepository({ store: new InMemorySessionStore() }); const result = getOrThrow(ok({ value: 1 })); const customMessage = createCustomMessage("note", "hello", true, undefined, "2026-01-01T00:00:00.000Z"); const llmMessages = convertToLlm([customMessage]);