diff --git a/packages/agent/src/harness/session/repo-utils.ts b/packages/agent/src/harness/session/repo-utils.ts index e830e7678..546c8449d 100644 --- a/packages/agent/src/harness/session/repo-utils.ts +++ b/packages/agent/src/harness/session/repo-utils.ts @@ -5,7 +5,7 @@ import { SessionError, type SessionMetadata, type SessionSearchHit, - type SessionSearchIndexSink, + type SessionSearchIndexWriter, type SessionStorage, type SessionTreeEntry, } from "../types.ts"; @@ -21,9 +21,9 @@ export function createTimestamp(): string { export function toSession( storage: SessionStorage, - searchIndexSink?: SessionSearchIndexSink, + searchIndexWriter?: SessionSearchIndexWriter, ): Session { - return new Session(storage, {}, searchIndexSink); + return new Session(storage, {}, searchIndexWriter); } export function findSessionEntryMatches( diff --git a/packages/agent/src/harness/session/search-index-sink.ts b/packages/agent/src/harness/session/search-index-sink.ts deleted file mode 100644 index 7745e99cd..000000000 --- a/packages/agent/src/harness/session/search-index-sink.ts +++ /dev/null @@ -1,8 +0,0 @@ -import type { SessionMetadata, SessionSearchIndexRecord, SessionTreeEntry } from "../types.ts"; - -export function toSessionSearchIndexRecord( - metadata: TMetadata, - entry: SessionTreeEntry, -): SessionSearchIndexRecord { - return { metadata, entry }; -} diff --git a/packages/agent/src/harness/session/search-index.ts b/packages/agent/src/harness/session/search-index.ts new file mode 100644 index 000000000..14e1ba4e6 --- /dev/null +++ b/packages/agent/src/harness/session/search-index.ts @@ -0,0 +1,32 @@ +import type { + SessionMetadata, + SessionSearchIndexRecord, + SessionSearchIndexWriter, + SessionTreeEntry, +} from "../types.ts"; + +export function toSessionSearchIndexRecord( + metadata: TMetadata, + entry: SessionTreeEntry, +): SessionSearchIndexRecord { + return { metadata, entry }; +} + +/** Delivers each lifecycle event to its targets, like Go's io.MultiWriter. */ +export class MultiSessionSearchIndexWriter + implements SessionSearchIndexWriter +{ + private readonly writers: readonly SessionSearchIndexWriter[]; + + constructor(...writers: SessionSearchIndexWriter[]) { + this.writers = writers; + } + + async write(record: SessionSearchIndexRecord): Promise { + for (const writer of this.writers) await writer.write(record); + } + + async remove(record: SessionSearchIndexRecord): Promise { + for (const writer of this.writers) await writer.remove(record); + } +} diff --git a/packages/agent/src/harness/session/session.ts b/packages/agent/src/harness/session/session.ts index 6e380f519..330b8166f 100644 --- a/packages/agent/src/harness/session/session.ts +++ b/packages/agent/src/harness/session/session.ts @@ -14,14 +14,14 @@ import type { SessionEntryCursorOptions, SessionInfoEntry, SessionMetadata, - SessionSearchIndexSink, + SessionSearchIndexWriter, SessionStats, SessionStorage, SessionTreeEntry, ThinkingLevelChangeEntry, } from "../types.ts"; import { SessionError } from "../types.ts"; -import { toSessionSearchIndexRecord } from "./search-index-sink.ts"; +import { toSessionSearchIndexRecord } from "./search-index.ts"; export type ContextEntryTransform = (entries: readonly SessionTreeEntry[]) => readonly SessionTreeEntry[]; @@ -152,16 +152,16 @@ export function buildSessionContext( export class Session { private storage: SessionStorage; private contextBuildOptions: SessionContextBuildOptions; - private searchIndexSink: SessionSearchIndexSink | undefined; + private searchIndexWriter: SessionSearchIndexWriter | undefined; constructor( storage: SessionStorage, contextBuildOptions: SessionContextBuildOptions = {}, - searchIndexSink?: SessionSearchIndexSink, + searchIndexWriter?: SessionSearchIndexWriter, ) { this.storage = storage; this.contextBuildOptions = contextBuildOptions; - this.searchIndexSink = searchIndexSink; + this.searchIndexWriter = searchIndexWriter; } getMetadata(): Promise { @@ -220,8 +220,8 @@ export class Session { } private async writeSearchIndexRecord(entry: SessionTreeEntry): Promise { - if (this.searchIndexSink) { - await this.searchIndexSink.write(toSessionSearchIndexRecord(await this.storage.getMetadata(), entry)); + if (this.searchIndexWriter) { + await this.searchIndexWriter.write(toSessionSearchIndexRecord(await this.storage.getMetadata(), entry)); } } diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index 76743791e..b4c3162e9 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -538,9 +538,16 @@ export interface SessionSearchIndexRecord; +/** Receives persisted-entry lifecycle events for search indexing. */ +export interface SessionSearchIndexWriter { + write(record: SessionSearchIndexRecord): Promise; + remove(record: SessionSearchIndexRecord): Promise; +} + +/** A search index that can both receive session entries and query them. */ +export interface SessionSearchIndex + extends SessionSearchIndexWriter { + search(options: SessionSearchOptions): Promise[]>; } export interface SessionForkOptions { diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index c0583af3d..d6bd648e8 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -34,7 +34,7 @@ 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 * from "./harness/session/search-index-sink.ts"; +export * from "./harness/session/search-index.ts"; export * from "./harness/session/session.ts"; export * from "./harness/skills.ts"; export * from "./harness/system-prompt.ts"; diff --git a/packages/agent/test/harness/sqlite-node.test.ts b/packages/agent/test/harness/sqlite-node.test.ts index a4248aa96..a94ad61b5 100644 --- a/packages/agent/test/harness/sqlite-node.test.ts +++ b/packages/agent/test/harness/sqlite-node.test.ts @@ -1,7 +1,17 @@ import { join } from "node:path"; import { describe, expect, it } from "vitest"; -import { createNodeSqliteFactory, SqliteSessionRepo } from "../../../storage/sqlite-node/src/index.ts"; +import { + createNodeSqliteFactory, + type SqliteSessionMetadata, + SqliteSessionRepo, +} from "../../../storage/sqlite-node/src/index.ts"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; +import type { + SessionSearchHit, + SessionSearchIndex, + SessionSearchIndexRecord, + SessionSearchOptions, +} from "../../src/harness/types.ts"; import { createTempDir, createUserMessage } from "./session-test-utils.ts"; describe("sqlite-node adapter", () => { @@ -20,6 +30,42 @@ describe("sqlite-node adapter", () => { } }); + it("delegates lifecycle events and queries to an injected search index", async () => { + const root = createTempDir(); + const writes: SessionSearchIndexRecord[] = []; + const removals: SessionSearchIndexRecord[] = []; + const searches: SessionSearchOptions[] = []; + const index: SessionSearchIndex = { + async write(record) { + writes.push(record); + }, + async remove(record) { + removals.push(record); + }, + async search(options): Promise[]> { + searches.push(options); + return []; + }, + }; + const repo = new SqliteSessionRepo({ + env: new NodeExecutionEnv({ cwd: root }), + sqlite: createNodeSqliteFactory(), + databasePath: join(root, "sessions.sqlite"), + searchIndexFactory: () => index, + }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const metadata = await session.getMetadata(); + const entryId = await session.appendMessage(createUserMessage("indexed remotely")); + + expect(writes).toEqual([expect.objectContaining({ metadata, entry: expect.objectContaining({ id: entryId }) })]); + await expect(repo.search({ text: "indexed remotely" })).resolves.toEqual([]); + expect(searches).toEqual([{ text: "indexed remotely" }]); + await repo.delete(metadata); + expect(removals).toEqual([ + expect.objectContaining({ metadata, entry: expect.objectContaining({ id: entryId }) }), + ]); + }); + it("searches and deletes transactionally indexed canonical entries", async () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite"); diff --git a/packages/storage/sqlite-node/src/sqlite/index.ts b/packages/storage/sqlite-node/src/sqlite/index.ts index 279a6083d..583bb49ef 100644 --- a/packages/storage/sqlite-node/src/sqlite/index.ts +++ b/packages/storage/sqlite-node/src/sqlite/index.ts @@ -1,4 +1,5 @@ export * from "./migrations.ts"; export * from "./repo.ts"; +export * from "./search-index.ts"; export * from "./storage/index.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 index 4efc7e6ff..4a2140ec7 100644 --- a/packages/storage/sqlite-node/src/sqlite/repo.ts +++ b/packages/storage/sqlite-node/src/sqlite/repo.ts @@ -14,7 +14,7 @@ import { toSessionSearchIndexRecord, } from "@earendil-works/pi-agent-core"; import { applyMigrations } from "./migrations.ts"; -import { SqliteSessionSearchIndex } from "./search-index-sink.ts"; +import { SqliteSessionSearchIndex } from "./search-index.ts"; import { SqliteSessionStorage } from "./storage/index.ts"; import { decodeEntry, type SessionEntryRow } from "./storage/session-entries.ts"; import { rowToMetadata, type SessionRow } from "./storage/sessions.ts"; @@ -26,6 +26,7 @@ import type { SqliteSessionMetadata, SqliteSessionRepoApi, SqliteSessionRepoEnv, + SqliteSessionSearchIndexFactory, } from "./types.ts"; function getParentPath(path: string): string { @@ -53,12 +54,20 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi { private readonly env: SqliteSessionRepoEnv; private readonly sqlite: SqliteDatabaseFactory; private readonly databasePathInput: string; + private readonly searchIndexFactory: SqliteSessionSearchIndexFactory; private databasePath: string | undefined; - constructor(options: { env: SqliteSessionRepoEnv; sqlite: SqliteDatabaseFactory; databasePath: string }) { + constructor(options: { + env: SqliteSessionRepoEnv; + sqlite: SqliteDatabaseFactory; + databasePath: string; + searchIndexFactory?: SqliteSessionSearchIndexFactory; + }) { this.env = options.env; this.sqlite = options.sqlite; this.databasePathInput = options.databasePath; + this.searchIndexFactory = + options.searchIndexFactory ?? (({ db, databasePath }) => new SqliteSessionSearchIndex(db, databasePath)); } private async getDatabasePath(): Promise { @@ -71,6 +80,10 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi { return this.databasePath; } + private async createSearchIndex(db: SqliteDatabase) { + return this.searchIndexFactory({ db, databasePath: await this.getDatabasePath() }); + } + private async ensureDatabaseDir(): Promise { const path = await this.getDatabasePath(); const directory = getParentPath(path); @@ -103,7 +116,7 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi { parentSessionId: options.parentSessionId, metadata: options.metadata, }); - return toSession(storage, new SqliteSessionSearchIndex(db, await this.getDatabasePath())); + return toSession(storage, await this.createSearchIndex(db)); } catch (error) { await db.close(); throw error; @@ -119,7 +132,7 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi { const db = await this.openDatabase(); try { const storage = await SqliteSessionStorage.open(db, metadata); - return toSession(storage, new SqliteSessionSearchIndex(db, await this.getDatabasePath())); + return toSession(storage, await this.createSearchIndex(db)); } catch (error) { await db.close(); throw error; @@ -155,7 +168,7 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi { if (!getFileSystemResultOrThrow(await this.env.exists(path), `Failed to check database ${path}`)) return []; const db = await this.openDatabase(); try { - return await new SqliteSessionSearchIndex(db, path).search(options); + return await (await this.createSearchIndex(db)).search(options); } finally { await db.close(); } @@ -165,14 +178,14 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi { const db = await this.openDatabase(); try { await db.transaction(async () => { - const searchIndexSink = new SqliteSessionSearchIndex(db, await this.getDatabasePath()); + const searchIndex = await this.createSearchIndex(db); const entries = await db .prepare( "SELECT rowid, session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ?", ) .all(metadata.id); for (const row of entries) { - await searchIndexSink.remove(toSessionSearchIndexRecord(metadata, decodeEntry(row))); + await searchIndex.remove(toSessionSearchIndexRecord(metadata, decodeEntry(row))); } 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); @@ -209,13 +222,13 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi { parentSessionId: options.parentSessionId ?? sourceMetadata.id, metadata: options.metadata ?? sourceMetadata.metadata, }); - const searchIndexSink = new SqliteSessionSearchIndex(db, await this.getDatabasePath()); + const searchIndex = await this.createSearchIndex(db); const metadata = await storage.getMetadata(); for (const entry of forkedEntries) { await storage.appendEntry(entry); - await searchIndexSink.write(toSessionSearchIndexRecord(metadata, entry)); + await searchIndex.write(toSessionSearchIndexRecord(metadata, entry)); } - return toSession(storage, searchIndexSink); + return toSession(storage, searchIndex); } catch (error) { await db.close(); throw error; diff --git a/packages/storage/sqlite-node/src/sqlite/search-index-sink.ts b/packages/storage/sqlite-node/src/sqlite/search-index.ts similarity index 89% rename from packages/storage/sqlite-node/src/sqlite/search-index-sink.ts rename to packages/storage/sqlite-node/src/sqlite/search-index.ts index 0f271a51c..3a78c2eb7 100644 --- a/packages/storage/sqlite-node/src/sqlite/search-index-sink.ts +++ b/packages/storage/sqlite-node/src/sqlite/search-index.ts @@ -1,15 +1,15 @@ import type { SessionSearchHit, + SessionSearchIndex, SessionSearchIndexRecord, - SessionSearchIndexSink, SessionSearchOptions, } from "@earendil-works/pi-agent-core"; import { SessionError } from "@earendil-works/pi-agent-core"; import { rowToMetadata, type SessionRow } from "./storage/sessions.ts"; import type { SqliteDatabase, SqliteSessionMetadata } from "./types.ts"; -/** SQLite-backed search index and sink for persisted session entries. */ -export class SqliteSessionSearchIndex implements SessionSearchIndexSink { +/** SQLite-backed search index for persisted session entries. */ +export class SqliteSessionSearchIndex implements SessionSearchIndex { private readonly db: SqliteDatabase; private readonly databasePath: string; @@ -18,7 +18,7 @@ export class SqliteSessionSearchIndex implements SessionSearchIndexSink { this.databasePath = databasePath; } - async write(record: SessionSearchIndexRecord): Promise { + async write(record: SessionSearchIndexRecord): Promise { const row = await this.db .prepare("SELECT rowid FROM session_entries WHERE session_id = ? AND id = ?") .get<{ rowid: number }>(record.metadata.id, record.entry.id); @@ -30,7 +30,7 @@ export class SqliteSessionSearchIndex implements SessionSearchIndexSink { .run(row.rowid, JSON.stringify(record.entry)); } - async remove(record: SessionSearchIndexRecord): Promise { + async remove(record: SessionSearchIndexRecord): Promise { const row = await this.db .prepare("SELECT rowid FROM session_entries WHERE session_id = ? AND id = ?") .get<{ rowid: number }>(record.metadata.id, record.entry.id); diff --git a/packages/storage/sqlite-node/src/sqlite/types.ts b/packages/storage/sqlite-node/src/sqlite/types.ts index 4e97ffd2f..1cf56ba1b 100644 --- a/packages/storage/sqlite-node/src/sqlite/types.ts +++ b/packages/storage/sqlite-node/src/sqlite/types.ts @@ -1,4 +1,10 @@ -import type { FileSystem, SessionCreateOptions, SessionMetadata, SessionRepo } from "@earendil-works/pi-agent-core"; +import type { + FileSystem, + SessionCreateOptions, + SessionMetadata, + SessionRepo, + SessionSearchIndex, +} from "@earendil-works/pi-agent-core"; /** Result of a prepared SQLite statement execution. */ export interface SqliteRunResult { @@ -53,3 +59,8 @@ export interface SqliteSessionRepoApi extends SessionRepo {} export type SqliteSessionRepoEnv = Pick; + +export type SqliteSessionSearchIndexFactory = (options: { + db: SqliteDatabase; + databasePath: string; +}) => SessionSearchIndex;