feat: search index factory

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<TMetadata extends SessionMetadata>(
 	storage: SessionStorage<TMetadata>,
-	searchIndexSink?: SessionSearchIndexSink,
+	searchIndexWriter?: SessionSearchIndexWriter<TMetadata>,
 ): Session<TMetadata> {
-	return new Session(storage, {}, searchIndexSink);
+	return new Session(storage, {}, searchIndexWriter);
 }
 
 export function findSessionEntryMatches<TMetadata extends SessionMetadata>(

packages/agent/src/harness/session/search-index-sink.ts

@@ -1,8 +0,0 @@
-import type { SessionMetadata, SessionSearchIndexRecord, SessionTreeEntry } from "../types.ts";
-
-export function toSessionSearchIndexRecord<TMetadata extends SessionMetadata>(
-	metadata: TMetadata,
-	entry: SessionTreeEntry,
-): SessionSearchIndexRecord<TMetadata> {
-	return { metadata, entry };
-}

packages/agent/src/harness/session/search-index.ts

@@ -0,0 +1,32 @@
+import type {
+	SessionMetadata,
+	SessionSearchIndexRecord,
+	SessionSearchIndexWriter,
+	SessionTreeEntry,
+} from "../types.ts";
+
+export function toSessionSearchIndexRecord<TMetadata extends SessionMetadata>(
+	metadata: TMetadata,
+	entry: SessionTreeEntry,
+): SessionSearchIndexRecord<TMetadata> {
+	return { metadata, entry };
+}
+
+/** Delivers each lifecycle event to its targets, like Go's io.MultiWriter. */
+export class MultiSessionSearchIndexWriter<TMetadata extends SessionMetadata>
+	implements SessionSearchIndexWriter<TMetadata>
+{
+	private readonly writers: readonly SessionSearchIndexWriter<TMetadata>[];
+
+	constructor(...writers: SessionSearchIndexWriter<TMetadata>[]) {
+		this.writers = writers;
+	}
+
+	async write(record: SessionSearchIndexRecord<TMetadata>): Promise<void> {
+		for (const writer of this.writers) await writer.write(record);
+	}
+
+	async remove(record: SessionSearchIndexRecord<TMetadata>): Promise<void> {
+		for (const writer of this.writers) await writer.remove(record);
+	}
+}

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<TMetadata extends SessionMetadata = SessionMetadata> {
 	private storage: SessionStorage<TMetadata>;
 	private contextBuildOptions: SessionContextBuildOptions;
-	private searchIndexSink: SessionSearchIndexSink | undefined;
+	private searchIndexWriter: SessionSearchIndexWriter<TMetadata> | undefined;
 
 	constructor(
 		storage: SessionStorage<TMetadata>,
 		contextBuildOptions: SessionContextBuildOptions = {},
-		searchIndexSink?: SessionSearchIndexSink,
+		searchIndexWriter?: SessionSearchIndexWriter<TMetadata>,
 	) {
 		this.storage = storage;
 		this.contextBuildOptions = contextBuildOptions;
-		this.searchIndexSink = searchIndexSink;
+		this.searchIndexWriter = searchIndexWriter;
 	}
 
 	getMetadata(): Promise<TMetadata> {
@@ -220,8 +220,8 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
 	}
 
 	private async writeSearchIndexRecord(entry: SessionTreeEntry): Promise<void> {
-		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));
 		}
 	}
 

packages/agent/src/harness/types.ts

@@ -538,9 +538,16 @@ export interface SessionSearchIndexRecord<TMetadata extends SessionMetadata = Se
 	entry: SessionTreeEntry;
 }
 
-/** Receives entries after their session storage has persisted them for search indexing. */
-export interface SessionSearchIndexSink {
-	write(record: SessionSearchIndexRecord): Promise<void>;
+/** Receives persisted-entry lifecycle events for search indexing. */
+export interface SessionSearchIndexWriter<TMetadata extends SessionMetadata = SessionMetadata> {
+	write(record: SessionSearchIndexRecord<TMetadata>): Promise<void>;
+	remove(record: SessionSearchIndexRecord<TMetadata>): Promise<void>;
+}
+
+/** A search index that can both receive session entries and query them. */
+export interface SessionSearchIndex<TMetadata extends SessionMetadata = SessionMetadata>
+	extends SessionSearchIndexWriter<TMetadata> {
+	search(options: SessionSearchOptions): Promise<SessionSearchHit<TMetadata>[]>;
 }
 
 export interface SessionForkOptions {

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";

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<SqliteSessionMetadata>[] = [];
+		const removals: SessionSearchIndexRecord<SqliteSessionMetadata>[] = [];
+		const searches: SessionSearchOptions[] = [];
+		const index: SessionSearchIndex<SqliteSessionMetadata> = {
+			async write(record) {
+				writes.push(record);
+			},
+			async remove(record) {
+				removals.push(record);
+			},
+			async search(options): Promise<SessionSearchHit<SqliteSessionMetadata>[]> {
+				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");

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";

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<string> {
@@ -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<void> {
 		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<SessionEntryRow & { rowid: number }>(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;

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<SqliteSessionMetadata> {
 	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<void> {
+	async write(record: SessionSearchIndexRecord<SqliteSessionMetadata>): Promise<void> {
 		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<void> {
+	async remove(record: SessionSearchIndexRecord<SqliteSessionMetadata>): Promise<void> {
 		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);

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<SqliteSessionMetadata, SqliteSessionCreateOptions, SqliteSessionListOptions> {}
 
 export type SqliteSessionRepoEnv = Pick<FileSystem, "absolutePath" | "createDir" | "exists">;
+
+export type SqliteSessionSearchIndexFactory = (options: {
+	db: SqliteDatabase;
+	databasePath: string;
+}) => SessionSearchIndex<SqliteSessionMetadata>;