refactor: search (#7797)

* refactor: search out of repo

* fix: apply cases, no resolve for index on every hit

* chore: tests

* fix: use iterator and no open

* refactor: public dir stuff in jsonl, source

* cleanup and sqlite time number

* fix: use storage for label search too, delete types

* refactor: abstract some more

* docs: example jsonl with elastic search

* docs: Christian's md

* refactor: simpler api surface

* feat: rebuild

* docs: consolidate

* fix: use trigger for sqlite, cleanup indexable, add example in doc

* chore: restore upstream responses shared

* fix: delete metadata from jsonl search hit

* fix: AsynIterable

* refactor: scanning for memory, reuse

* fix: dont export jsonlSessionDirectoryName

* cleanup: more shared stuff

* cleanup: why, jsonl?

packages/session-backends/sqlite-node/src/sqlite/search-backend.ts

@@ -1,5 +1,11 @@
-import type { SessionSearch, SessionSearchHit, SessionSearchOptions } from "@earendil-works/pi-agent-core";
-import { getFileSystemResultOrThrow } from "@earendil-works/pi-agent-core";
+import type {
+	FileError,
+	Result,
+	SessionSearch,
+	SessionSearchHit,
+	SessionSearchOptions,
+} from "@earendil-works/pi-agent-core";
+import { SessionError } from "@earendil-works/pi-agent-core";
 import { applyMigrations } from "./migrations.ts";
 import { sql } from "./sql.ts";
 import { decodeSessionMetadata, type SessionRow } from "./storage/sessions.ts";
@@ -10,6 +16,14 @@ import type {
 	SqliteSessionRepositoryEnv,
 } from "./types.ts";
 
+function getFileSystemResultOrThrow<TValue>(result: Result<TValue, FileError>, 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;
+}
+
 function getParentPath(path: string): string {
 	const normalized = path.replace(/[\\/]+$/, "");
 	const lastSlash = Math.max(normalized.lastIndexOf("/"), normalized.lastIndexOf("\\"));
@@ -18,6 +32,14 @@ function getParentPath(path: string): string {
 	return normalized.slice(0, lastSlash);
 }
 
+function throwIfAborted(signal: AbortSignal | undefined): void {
+	if (!signal?.aborted) return;
+	if (signal.reason instanceof Error) throw signal.reason;
+	const error = new Error("The operation was aborted");
+	error.name = "AbortError";
+	throw error;
+}
+
 function configureSqliteDatabase(db: SqliteDatabase): void {
 	sql`PRAGMA journal_mode=WAL`.exec(db);
 	sql`PRAGMA synchronous=FULL`.exec(db);
@@ -36,9 +58,15 @@ function tableExists(db: SqliteDatabase, name: string): boolean {
 	}>(db);
 }
 
+function rebuildSearchIndex(db: SqliteDatabase): void {
+	sql`INSERT INTO session_search_fts(session_search_fts) VALUES('rebuild')`.run(db);
+}
+
 function ensureSearchSchema(db: SqliteDatabase): void {
 	const ftsExists = tableExists(db, "session_search_fts");
-	sql`
+	const entriesExist = tableExists(db, "entries");
+	db.transaction(() => {
+		sql`
 CREATE VIRTUAL TABLE IF NOT EXISTS session_search_fts USING fts5(
   payload,
   content = 'entries',
@@ -56,11 +84,18 @@ CREATE TRIGGER IF NOT EXISTS session_search_fts_au AFTER UPDATE OF payload ON en
   INSERT INTO session_search_fts(rowid, payload) VALUES (new.rowid, new.payload);
 END;
 `.exec(db);
-	if (!ftsExists) sql`INSERT INTO session_search_fts(session_search_fts) VALUES('rebuild')`.exec(db);
+		if (!ftsExists && entriesExist) rebuildSearchIndex(db);
+	});
+}
+
+export interface SqliteSessionSearchHit extends SessionSearchHit {
+	readonly metadata: SqliteSessionMetadata;
+	readonly timestamp: number;
+	readonly score: number;
 }
 
 /** SQLite FTS search over a co-located canonical session database. */
-class SqliteSessionSearch implements SessionSearch<SqliteSessionMetadata> {
+class SqliteSessionSearch implements SessionSearch<SqliteSessionSearchHit> {
 	private readonly options: SqliteSessionSearchOptions;
 	private databasePath: string | undefined;
 
@@ -97,44 +132,63 @@ class SqliteSessionSearch implements SessionSearch<SqliteSessionMetadata> {
 		}
 	}
 
-	async search(options: SessionSearchOptions): Promise<SessionSearchHit<SqliteSessionMetadata>[]> {
-		const text = options.text.trim();
-		if (!text) return [];
+	async *search(text: string, options: SessionSearchOptions = {}): AsyncIterable<SqliteSessionSearchHit> {
+		const queryText = text.trim();
+		if (!queryText || (options.limit !== undefined && options.limit <= 0)) return;
+		if (options.entryTypes?.length === 0) return;
+		throwIfAborted(options.signal);
 		const db = await this.openDatabase();
 		try {
-			const query = `"${text.replaceAll('"', '""')}"`;
-			const cwd = options.cwd ?? null;
-			const rows = sql`SELECT s.id, s.created_at, s.metadata, s.cwd, s.parent_session_id,
-				name_fact.seq IS NOT NULL AS has_session_name,
-				name_fact.value AS session_name,
-				se.id AS entry_id, se.timestamp, bm25(session_search_fts) AS score
-			FROM session_search_fts
-			JOIN entries AS se ON se.rowid = session_search_fts.rowid
-			JOIN sessions AS s ON s.id = se.session_id
-			LEFT JOIN facts AS name_fact
-				ON name_fact.session_id = s.id
-				AND name_fact.kind = 'name'
-				AND name_fact.key IS NULL
-				AND name_fact.seq = (
-					SELECT MAX(f.seq)
-					FROM facts AS f
-					WHERE f.session_id = s.id AND f.kind = 'name' AND f.key IS NULL
+			const query = `"${queryText.replaceAll('"', '""')}"`;
+			const predicates = ["session_search_fts MATCH ?"];
+			const params: unknown[] = [query];
+			if (options.entryTypes !== undefined) {
+				predicates.push(`se.type IN (${options.entryTypes.map(() => "?").join(", ")})`);
+				params.push(...options.entryTypes);
+			}
+			const rows = db
+				.prepare(
+					`SELECT s.id, s.created_at, s.metadata, s.cwd, s.parent_session_id,
+						name_fact.seq IS NOT NULL AS has_session_name,
+						name_fact.value AS session_name,
+						se.id AS entry_id, se.timestamp, bm25(session_search_fts) AS score
+					FROM session_search_fts
+					JOIN entries AS se ON se.rowid = session_search_fts.rowid
+					JOIN sessions AS s ON s.id = se.session_id
+					LEFT JOIN facts AS name_fact
+						ON name_fact.session_id = s.id
+						AND name_fact.kind = 'name'
+						AND name_fact.key IS NULL
+						AND name_fact.seq = (
+							SELECT MAX(f.seq)
+							FROM facts AS f
+							WHERE f.session_id = s.id AND f.kind = 'name' AND f.key IS NULL
+						)
+					WHERE ${predicates.join(" AND ")}
+					ORDER BY score
+					LIMIT ?`,
 				)
-			WHERE session_search_fts MATCH ${query} AND (${cwd} IS NULL OR s.cwd = ${cwd})
-			ORDER BY score`.all<SessionRow & { entry_id: string; timestamp: number; score: number }>(db);
+				.iterate<SessionRow & { entry_id: string; timestamp: number; score: number }>(
+					...params,
+					options.limit ?? -1,
+				);
 			const path = await this.getDatabasePath();
-			return rows.map((row) => ({
-				metadata: decodeSessionMetadata(row, path),
-				entryId: row.entry_id,
-				timestamp: row.timestamp,
-				score: row.score,
-			}));
+			for (const row of rows) {
+				throwIfAborted(options.signal);
+				yield {
+					sessionId: row.id,
+					metadata: decodeSessionMetadata(row, path),
+					entryId: row.entry_id,
+					timestamp: row.timestamp,
+					score: row.score,
+				};
+			}
 		} finally {
 			db.close();
 		}
 	}
 }
 
-export function createSqliteSessionSearch(options: SqliteSessionSearchOptions): SessionSearch<SqliteSessionMetadata> {
+export function createSqliteSessionSearch(options: SqliteSessionSearchOptions): SessionSearch<SqliteSessionSearchHit> {
 	return new SqliteSessionSearch(options);
 }

packages/session-backends/sqlite-node/src/sqlite/sql.ts

@@ -28,6 +28,10 @@ export class SqlQuery {
 	all<TRow extends object>(db: SqliteDatabase): TRow[] {
 		return db.prepare(this.queryText).all<TRow>(...this.params);
 	}
+
+	iterate<TRow extends object>(db: SqliteDatabase): Iterable<TRow> {
+		return db.prepare(this.queryText).iterate<TRow>(...this.params);
+	}
 }
 
 /** Builds a parameterized query. Nested queries are inlined; other interpolations become `?` parameters. */

packages/session-backends/sqlite-node/src/sqlite/types.ts

@@ -13,6 +13,7 @@ export interface SqliteStatement {
 	run(...params: unknown[]): SqliteRunResult;
 	get<TRow extends object>(...params: unknown[]): TRow | undefined;
 	all<TRow extends object>(...params: unknown[]): TRow[];
+	iterate<TRow extends object>(...params: unknown[]): Iterable<TRow>;
 }
 
 /** SQLite database capability used by the SQLite session backend. */