Merge #38: configurable embedding probe timeout

README.md

@@ -189,6 +189,7 @@ This ensures in-progress context survives compaction and is visible in the next
 | `PI_MEMORY_SNAPSHOT` | `stable`, `per-turn` | `stable` | `stable` snapshots memory at checkpoints for KV cache stability; `per-turn` rebuilds every turn (legacy behavior) |
 | `PI_MEMORY_QMD_UPDATE` | `background`, `manual`, `off` | `background` | Controls automatic `qmd update` + `qmd embed` after writes |
 | `PI_MEMORY_QMD_SEARCH_TIMEOUT_MS` | positive integer (milliseconds) | `60000` | Sets the timeout for explicit `memory_search` qmd queries |
+| `PI_MEMORY_EMBED_PROBE_TIMEOUT_MS` | positive integer (milliseconds) | `15000` | Sets the timeout for the `memory_status` embeddings readiness probe. Raise it on slower machines if the probe reports `unknown` |
 | `PI_MEMORY_NO_SEARCH` | `1` | unset | Disable selective injection in `per-turn` mode (no effect in `stable` mode) |
 | `PI_MEMORY_SUMMARIZE_TRANSITIONS` | `1`, `true`, `yes`, `on` | unset | Also write exit summaries during lifecycle transitions (`/reload`, `/new`, `/resume`, `/fork`). By default these transitions skip summaries for speed. |
 | `PI_MEMORY_EXIT_SUMMARY` | `0`, `off`, `false`, `no` to disable | unset (enabled) | Disable the exit summary on real quit (Ctrl+D, `/quit`, session end). Quitting then does no LLM call and no `qmd update`, so it is instant; explicit `memory_write` during sessions is unaffected. |

index.ts

@@ -974,6 +974,7 @@ let qmdAvailabilityCheckedAt = 0;
 const QMD_STATUS_CACHE_TTL_MS = 5 * 60 * 1000;
 const QMD_STATUS_NEGATIVE_CACHE_TTL_MS = 5 * 1000;
 const DEFAULT_QMD_SEARCH_TIMEOUT_MS = 60_000;
+const DEFAULT_EMBED_PROBE_TIMEOUT_MS = 15_000;
 const qmdCollectionStatusCache = new Map<string, { checkedAt: number; exists: boolean }>();
 
 function qmdStatusTtl(positive: boolean): number {
@@ -984,6 +985,11 @@ export function getQmdSearchTimeoutMs(env: NodeJS.ProcessEnv = process.env): num
 	const configured = Number(env.PI_MEMORY_QMD_SEARCH_TIMEOUT_MS);
 	return Number.isInteger(configured) && configured > 0 ? configured : DEFAULT_QMD_SEARCH_TIMEOUT_MS;
 }
+
+export function getEmbedProbeTimeoutMs(env: NodeJS.ProcessEnv = process.env): number {
+	const configured = Number(env.PI_MEMORY_EMBED_PROBE_TIMEOUT_MS);
+	return Number.isInteger(configured) && configured > 0 ? configured : DEFAULT_EMBED_PROBE_TIMEOUT_MS;
+}
 let updateTimer: ReturnType<typeof setTimeout> | null = null;
 let exitSummaryReason: ExitSummaryReason | null = null;
 let terminalInputUnsubscribe: (() => void) | null = null;
@@ -1311,10 +1317,11 @@ export function runQmdSearch(
 	mode: "keyword" | "semantic" | "deep",
 	query: string,
 	limit: number,
+	timeoutOverrideMs?: number,
 ): Promise<{ results: QmdSearchResult[]; stderr: string }> {
 	const subcommand = mode === "keyword" ? "search" : mode === "semantic" ? "vsearch" : "query";
 	const args = [subcommand, "--json", "-c", "pi-memory", "-n", String(limit), query];
-	const timeoutMs = getQmdSearchTimeoutMs();
+	const timeoutMs = timeoutOverrideMs ?? getQmdSearchTimeoutMs();
 
 	return new Promise((resolve, reject) => {
 		execFileFn("qmd", args, { timeout: timeoutMs }, (err, stdout, stderr) => {
@@ -1345,18 +1352,27 @@ export function runQmdSearch(
 
 /**
  * Best-effort check of whether vector embeddings are ready for semantic/deep
- * search. Bounded by a short timeout because the first semantic query can
- * trigger a model download. Returns "unknown" rather than blocking on it.
+ * search. Bounded by a timeout because the first semantic query can trigger a
+ * model download. Returns "unknown" rather than blocking on it.
  * "ready" means a probe query ran without qmd's "need embeddings" warning —
  * it does not prove the index has content.
+ *
+ * The bound must stay well clear of normal `qmd vsearch` latency: the probe
+ * runs an embed + rerank pass (measured ~2.4-3.6s idle, >4s while a background
+ * re-index competes for CPU and the embedding model). A tighter bound made
+ * `memory_status` report "unknown" immediately after a write, which is exactly
+ * when the index is busy. Override with PI_MEMORY_EMBED_PROBE_TIMEOUT_MS.
  */
 export async function probeEmbeddings(): Promise<"ready" | "missing" | "unknown"> {
+	const probeTimeoutMs = getEmbedProbeTimeoutMs();
 	let timer: ReturnType<typeof setTimeout> | undefined;
 	try {
 		const { stderr } = await Promise.race([
-			runQmdSearch("semantic", "memory", 1),
+			// Bound the child by the same budget so a probe we abandon does not
+			// leave a long-running LLM query behind.
+			runQmdSearch("semantic", "memory", 1, probeTimeoutMs),
 			new Promise<never>((_, reject) => {
-				timer = setTimeout(() => reject(new Error("timeout")), 4_000);
+				timer = setTimeout(() => reject(new Error("timeout")), probeTimeoutMs);
 			}),
 		]);
 		return /need embeddings/i.test(stderr ?? "") ? "missing" : "ready";
@@ -2413,7 +2429,10 @@ export default function (pi: ExtensionAPI) {
 							lines.push("  - Run `qmd embed` once to enable semantic/deep search.");
 						}
 					} else if (embeddings === "unknown") {
-						lines.push("  - Could not verify within the probe timeout; run a semantic search to confirm.");
+						lines.push(
+							`  - Could not verify within the ${getEmbedProbeTimeoutMs() / 1000}s probe timeout; run a semantic search to confirm.`,
+							"  - A background re-index can slow the probe. Raise PI_MEMORY_EMBED_PROBE_TIMEOUT_MS if it persists.",
+						);
 					}
 				} else {
 					lines.push("  - Run a `memory_search` (auto-creates it) or `qmd collection add` manually.");
@@ -2428,6 +2447,7 @@ export default function (pi: ExtensionAPI) {
 				`- PI_MEMORY_SNAPSHOT: ${getSnapshotMode()}`,
 				`- PI_MEMORY_QMD_UPDATE: ${getQmdUpdateMode()}`,
 				`- PI_MEMORY_QMD_SEARCH_TIMEOUT_MS: ${getQmdSearchTimeoutMs()}`,
+				`- PI_MEMORY_EMBED_PROBE_TIMEOUT_MS: ${getEmbedProbeTimeoutMs()}`,
 				`- PI_MEMORY_DIR: ${process.env.PI_MEMORY_DIR ? "set" : "default"}`,
 				`- PI_MEMORY_EXIT_SUMMARY: ${isExitSummaryEnabled() ? "enabled" : "disabled"}`,
 				`- PI_MEMORY_EXIT_SUMMARY_MODEL: ${process.env.PI_MEMORY_EXIT_SUMMARY_MODEL?.trim() || "session model"}`,

test/unit.test.ts

@@ -31,13 +31,15 @@ import {
 	ensureDirs,
 	ensureQmdEmbed,
 	forgetBlocks,
+	getEmbedProbeTimeoutMs,
 	getExitSummaryReasoningEffort,
 	getExitSummaryTimeoutMs,
 	getQmdSearchTimeoutMs,
 	isExitSummaryEmpty,
 	isExitSummaryEnabled,
 	nowTimestamp,
 	parseScratchpad,
+	probeEmbeddings,
 	qmdCollectionInstructions,
 	qmdInstallInstructions,
 	readFileSafe,
@@ -2437,3 +2439,85 @@ describe("memory_forget tool", () => {
 		expect(secondRestore.content[0].text).toContain("already restored");
 	});
 });
+
+// ==========================================================================
+// 13. probeEmbeddings timeout behavior
+// ==========================================================================
+
+describe("probeEmbeddings", () => {
+	const ENV_KEY = "PI_MEMORY_EMBED_PROBE_TIMEOUT_MS";
+	let previous: string | undefined;
+
+	beforeEach(() => {
+		previous = process.env[ENV_KEY];
+		delete process.env[ENV_KEY];
+	});
+
+	afterEach(() => {
+		_resetExecFileForTest();
+		if (previous === undefined) delete process.env[ENV_KEY];
+		else process.env[ENV_KEY] = previous;
+	});
+
+	test("defaults to a probe timeout with headroom over a contended qmd call", () => {
+		expect(getEmbedProbeTimeoutMs()).toBeGreaterThanOrEqual(15_000);
+	});
+
+	test("honors PI_MEMORY_EMBED_PROBE_TIMEOUT_MS override", () => {
+		process.env[ENV_KEY] = "9000";
+		expect(getEmbedProbeTimeoutMs()).toBe(9_000);
+	});
+
+	test("ignores invalid PI_MEMORY_EMBED_PROBE_TIMEOUT_MS values", () => {
+		for (const bad of ["0", "-1", "abc", "1.5"]) {
+			process.env[ENV_KEY] = bad;
+			expect(getEmbedProbeTimeoutMs()).toBeGreaterThanOrEqual(15_000);
+		}
+	});
+
+	test("reports ready when qmd answers without an embeddings warning", async () => {
+		_setExecFileForTest(((_file: string, _args: string[], _opts: any, cb: any) => {
+			cb(null, "[]", "");
+		}) as any);
+		expect(await probeEmbeddings()).toBe("ready");
+	});
+
+	test("reports missing when qmd warns that embeddings are needed", async () => {
+		_setExecFileForTest(((_file: string, _args: string[], _opts: any, cb: any) => {
+			cb(null, "[]", "warning: need embeddings for vector search");
+		}) as any);
+		expect(await probeEmbeddings()).toBe("missing");
+	});
+
+	test("survives a slow probe that would trip the old hardcoded 4s race", async () => {
+		process.env[ENV_KEY] = "20000";
+		_setExecFileForTest(((_file: string, _args: string[], _opts: any, cb: any) => {
+			setTimeout(() => cb(null, "[]", ""), 4_200);
+		}) as any);
+		expect(await probeEmbeddings()).toBe("ready");
+	}, 30_000);
+
+	test("bounds the qmd child process by the probe timeout, not the search timeout", async () => {
+		process.env[ENV_KEY] = "15000";
+		process.env.PI_MEMORY_QMD_SEARCH_TIMEOUT_MS = "60000";
+		let observedTimeout: number | undefined;
+		try {
+			_setExecFileForTest(((_file: string, _args: string[], opts: any, cb: any) => {
+				observedTimeout = opts.timeout;
+				cb(null, "[]", "");
+			}) as any);
+			await probeEmbeddings();
+			expect(observedTimeout).toBe(15_000);
+		} finally {
+			delete process.env.PI_MEMORY_QMD_SEARCH_TIMEOUT_MS;
+		}
+	});
+
+	test("still reports unknown when the probe genuinely times out", async () => {
+		process.env[ENV_KEY] = "150";
+		_setExecFileForTest(((_file: string, _args: string[], _opts: any, cb: any) => {
+			setTimeout(() => cb(null, "[]", ""), 2_000);
+		}) as any);
+		expect(await probeEmbeddings()).toBe("unknown");
+	}, 10_000);
+});