membot 0.7.0 → 0.8.0

package/src/ingest/embedder.ts

@@ -68,9 +68,40 @@ async function getPipeline(model: string): Promise<FeatureExtractionPipeline> {
  * with `(done, total)` chunk counts so callers can drive a spinner / progress
  * bar — ONNX WASM holds the JS thread for hundreds of ms per batch and would
  * otherwise leave nanospinner's setInterval starved between updates.
+ *
+ * `directOnly` bypasses any registered EmbedderPool and runs the embed call
+ * inline in the current process. Use it for query-time single-text embedding
+ * where IPC overhead would dominate.
  */
 export interface EmbedOptions {
 	onProgress?: (done: number, total: number) => void;
+	directOnly?: boolean;
+}
+
+/**
+ * The minimal surface the embedder needs from a worker pool. Defined as an
+ * interface (not an `import type`) so we don't take a hard dependency on
+ * `embedder-pool.ts` from this hot path — the pool is plugged in via
+ * `setEmbedderPool()` from outside.
+ */
+export interface PooledEmbedder {
+	embed(texts: string[], model?: string, opts?: EmbedOptions): Promise<number[][]>;
+}
+
+let pool: PooledEmbedder | null = null;
+
+/**
+ * Register a worker pool to handle bulk embed calls. After this is set, every
+ * `embed()` call (without `directOnly`) is dispatched through the pool.
+ * Called once during `buildContext()` when `config.embedding.workers > 1`.
+ */
+export function setEmbedderPool(p: PooledEmbedder | null): void {
+	pool = p;
+}
+
+/** Read the currently registered pool, or `null` when running single-process. */
+export function getEmbedderPool(): PooledEmbedder | null {
+	return pool;
 }
 
 /**
@@ -92,6 +123,9 @@ export async function embed(
 	opts: EmbedOptions = {},
 ): Promise<number[][]> {
 	if (texts.length === 0) return [];
+	if (pool && !opts.directOnly) {
+		return pool.embed(texts, model, opts);
+	}
 	const extractor = await getPipeline(model);
 	const out: number[][] = [];
 	for (let i = 0; i < texts.length; i += EMBEDDING_BATCH_SIZE) {
@@ -114,9 +148,13 @@ export async function embed(
 	return out;
 }
 
-/** Embed a single text — convenience wrapper for query-time embedding. */
+/**
+ * Embed a single text — convenience wrapper for query-time embedding. Always
+ * runs in-process (`directOnly: true`) so search latency isn't paying the IPC
+ * round-trip through the worker pool for one vector.
+ */
 export async function embedSingle(text: string, model: string = EMBEDDING_MODEL): Promise<number[]> {
-	const all = await embed([text], model);
+	const all = await embed([text], model, { directOnly: true });
 	const vec = all[0];
 	if (!vec) {
 		throw new HelpfulError({

package/src/operations/add.ts

@@ -1,4 +1,6 @@
 import { z } from "zod";
+import { resolveEmbeddingWorkers } from "../context.ts";
+import { withEmbedderPool } from "../ingest/embedder-pool.ts";
 import {
 	countResolvedEntries,
 	type IngestCallbacks,
@@ -105,101 +107,107 @@ Pass \`logical_path\` to override. For a multi-source / directory / glob walk it
 		return `${lines.join("\n")}\n${parts.join(", ")}`;
 	},
 	handler: async (input, ctx) => {
-		const { sources, ...rest } = input;
-		const followSymlinks = rest.follow_symlinks ?? true;
+		// Spin up an ephemeral embedder pool for the whole `add` command —
+		// `withEmbedderPool` handles the workers=1 short-circuit and disposes
+		// the children when the closure returns (see embedder-pool.ts).
+		const workers = resolveEmbeddingWorkers(ctx.config.embedding.workers);
+		return withEmbedderPool(workers, ctx.config.embedding_model, async () => {
+			const { sources, ...rest } = input;
+			const followSymlinks = rest.follow_symlinks ?? true;
 
-		// Phase 1: resolve every source upfront so the shared progress bar
-		// knows its total. A resolve failure (bad path, glob with no base) is
-		// captured per-source so one bad arg doesn't abort the whole batch.
-		type ResolveOutcome = { source: string; resolved: ResolvedSource } | { source: string; error: Error };
-		const outcomes: ResolveOutcome[] = [];
-		for (const source of sources) {
-			try {
-				const resolved = await resolveSource(source, {
-					include: rest.include,
-					exclude: rest.exclude,
-					followSymlinks,
-				});
-				outcomes.push({ source, resolved });
-			} catch (err) {
-				outcomes.push({ source, error: err instanceof Error ? err : new Error(String(err)) });
+			// Phase 1: resolve every source upfront so the shared progress bar
+			// knows its total. A resolve failure (bad path, glob with no base) is
+			// captured per-source so one bad arg doesn't abort the whole batch.
+			type ResolveOutcome = { source: string; resolved: ResolvedSource } | { source: string; error: Error };
+			const outcomes: ResolveOutcome[] = [];
+			for (const source of sources) {
+				try {
+					const resolved = await resolveSource(source, {
+						include: rest.include,
+						exclude: rest.exclude,
+						followSymlinks,
+					});
+					outcomes.push({ source, resolved });
+				} catch (err) {
+					outcomes.push({ source, error: err instanceof Error ? err : new Error(String(err)) });
+				}
 			}
-		}
 
-		const total = outcomes.reduce((n, o) => ("error" in o ? n + 1 : n + countResolvedEntries(o.resolved)), 0);
+			const total = outcomes.reduce((n, o) => ("error" in o ? n + 1 : n + countResolvedEntries(o.resolved)), 0);
 
-		const aggregated: IngestResult = {
-			ingested: [],
-			total: 0,
-			ok: 0,
-			unchanged: 0,
-			failed: 0,
-		};
+			const aggregated: IngestResult = {
+				ingested: [],
+				total: 0,
+				ok: 0,
+				unchanged: 0,
+				failed: 0,
+			};
 
-		ctx.progress.start(total, "ingest");
-		const callbacks: IngestCallbacks = {
-			onEntryStart: (label) => ctx.progress.tick(label),
-			onEntryComplete: (entry) => ctx.progress.entry(formatEntryLine(entry)),
-			onEntryProgress: (_label, sublabel) => ctx.progress.update(sublabel),
-		};
+			ctx.progress.start(total, "ingest");
+			const callbacks: IngestCallbacks = {
+				onEntryStart: (label) => ctx.progress.tick(label),
+				onEntryComplete: (entry) => ctx.progress.entry(formatEntryLine(entry)),
+				onEntryProgress: (_label, sublabel) => ctx.progress.update(sublabel),
+			};
 
-		for (const outcome of outcomes) {
-			if ("error" in outcome) {
-				const failed: IngestEntryResult = {
-					source_path: outcome.source,
-					logical_path: outcome.source,
-					version_id: null,
-					status: "failed",
-					error: outcome.error.message,
-					mime_type: null,
-					size_bytes: 0,
-					fetcher: "local",
-					source_sha256: "",
-				};
-				callbacks.onEntryStart?.(outcome.source);
-				callbacks.onEntryComplete?.(failed);
-				aggregated.ingested.push(failed);
-				aggregated.total += 1;
-				aggregated.failed += 1;
-				continue;
-			}
+			for (const outcome of outcomes) {
+				if ("error" in outcome) {
+					const failed: IngestEntryResult = {
+						source_path: outcome.source,
+						logical_path: outcome.source,
+						version_id: null,
+						status: "failed",
+						error: outcome.error.message,
+						mime_type: null,
+						size_bytes: 0,
+						fetcher: "local",
+						source_sha256: "",
+					};
+					callbacks.onEntryStart?.(outcome.source);
+					callbacks.onEntryComplete?.(failed);
+					aggregated.ingested.push(failed);
+					aggregated.total += 1;
+					aggregated.failed += 1;
+					continue;
+				}
 
-			try {
-				const r = await ingestResolved(outcome.resolved, { ...rest, source: outcome.source }, ctx, callbacks);
-				aggregated.ingested.push(...r.ingested);
-				aggregated.total += r.total;
-				aggregated.ok += r.ok;
-				aggregated.unchanged += r.unchanged;
-				aggregated.failed += r.failed;
-			} catch (err) {
-				const message = err instanceof Error ? err.message : String(err);
-				const failed: IngestEntryResult = {
-					source_path: outcome.source,
-					logical_path: outcome.source,
-					version_id: null,
-					status: "failed",
-					error: message,
-					mime_type: null,
-					size_bytes: 0,
-					fetcher: "local",
-					source_sha256: "",
-				};
-				callbacks.onEntryStart?.(outcome.source);
-				callbacks.onEntryComplete?.(failed);
-				aggregated.ingested.push(failed);
-				aggregated.total += 1;
-				aggregated.failed += 1;
-			} finally {
-				// Release the DB lock between sources so other consumers (a
-				// concurrent CLI call, the daemon, or a separate MCP server)
-				// can wedge in. The next source's first DB call reopens.
-				await ctx.db.release();
+				try {
+					const r = await ingestResolved(outcome.resolved, { ...rest, source: outcome.source }, ctx, callbacks);
+					aggregated.ingested.push(...r.ingested);
+					aggregated.total += r.total;
+					aggregated.ok += r.ok;
+					aggregated.unchanged += r.unchanged;
+					aggregated.failed += r.failed;
+				} catch (err) {
+					const message = err instanceof Error ? err.message : String(err);
+					const failed: IngestEntryResult = {
+						source_path: outcome.source,
+						logical_path: outcome.source,
+						version_id: null,
+						status: "failed",
+						error: message,
+						mime_type: null,
+						size_bytes: 0,
+						fetcher: "local",
+						source_sha256: "",
+					};
+					callbacks.onEntryStart?.(outcome.source);
+					callbacks.onEntryComplete?.(failed);
+					aggregated.ingested.push(failed);
+					aggregated.total += 1;
+					aggregated.failed += 1;
+				} finally {
+					// Release the DB lock between sources so other consumers (a
+					// concurrent CLI call, the daemon, or a separate MCP server)
+					// can wedge in. The next source's first DB call reopens.
+					await ctx.db.release();
+				}
 			}
-		}
 
-		const summary = formatSummary(aggregated);
-		ctx.progress.done(summary);
-		return aggregated;
+			const summary = formatSummary(aggregated);
+			ctx.progress.done(summary);
+			return aggregated;
+		});
 	},
 });
 

package/src/operations/index.ts

@@ -8,6 +8,7 @@ import { readOperation } from "./read.ts";
 import { refreshOperation } from "./refresh.ts";
 import { removeOperation } from "./remove.ts";
 import { searchOperation } from "./search.ts";
+import { statsOperation } from "./stats.ts";
 import { treeOperation } from "./tree.ts";
 import type { Operation } from "./types.ts";
 import { versionsOperation } from "./versions.ts";
@@ -28,6 +29,7 @@ export const OPERATIONS: Operation<any, any>[] = [
 	readOperation,
 	searchOperation,
 	infoOperation,
+	statsOperation,
 	versionsOperation,
 	diffOperation,
 	writeOperation,

package/src/operations/refresh.ts

@@ -1,5 +1,7 @@
 import { z } from "zod";
+import { resolveEmbeddingWorkers } from "../context.ts";
 import { listDueRefreshes } from "../db/files.ts";
+import { withEmbedderPool } from "../ingest/embedder-pool.ts";
 import { colors } from "../output/formatter.ts";
 import { refreshOne } from "../refresh/runner.ts";
 import { defineOperation } from "./types.ts";
@@ -47,26 +49,32 @@ export const refreshOperation = defineOperation({
 		return `${lines.join("\n")}\n${parts.join(", ")}`;
 	},
 	handler: async (input, ctx) => {
-		const targets = input.logical_path
-			? [input.logical_path]
-			: (await listDueRefreshes(ctx.db)).map((r) => r.logical_path);
-		const out: Array<{
-			logical_path: string;
-			status: "ok" | "unchanged" | "failed";
-			new_version_id?: string;
-			error?: string;
-		}> = [];
-		ctx.progress.start(targets.length, "refresh");
-		for (const path of targets) {
-			ctx.progress.tick(path);
-			try {
-				const r = await refreshOne(ctx, path, input.force, (sublabel) => ctx.progress.update(sublabel));
-				out.push(r);
-			} catch (err) {
-				out.push({ logical_path: path, status: "failed", error: err instanceof Error ? err.message : String(err) });
+		// Per-command embedder pool: workers come up at the start of the
+		// refresh sweep and are killed before we return, so a manual
+		// `membot refresh` doesn't leave subprocesses around.
+		const workers = resolveEmbeddingWorkers(ctx.config.embedding.workers);
+		return withEmbedderPool(workers, ctx.config.embedding_model, async () => {
+			const targets = input.logical_path
+				? [input.logical_path]
+				: (await listDueRefreshes(ctx.db)).map((r) => r.logical_path);
+			const out: Array<{
+				logical_path: string;
+				status: "ok" | "unchanged" | "failed";
+				new_version_id?: string;
+				error?: string;
+			}> = [];
+			ctx.progress.start(targets.length, "refresh");
+			for (const path of targets) {
+				ctx.progress.tick(path);
+				try {
+					const r = await refreshOne(ctx, path, input.force, (sublabel) => ctx.progress.update(sublabel));
+					out.push(r);
+				} catch (err) {
+					out.push({ logical_path: path, status: "failed", error: err instanceof Error ? err.message : String(err) });
+				}
 			}
-		}
-		ctx.progress.done(`refresh: ${out.filter((r) => r.status === "ok").length}/${out.length} updated`);
-		return { processed: out, count: out.length };
+			ctx.progress.done(`refresh: ${out.filter((r) => r.status === "ok").length}/${out.length} updated`);
+			return { processed: out, count: out.length };
+		});
 	},
 });