diff --git a/README.md b/README.md index 2927254..5c44bde 100644 --- a/README.md +++ b/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. | diff --git a/index.ts b/index.ts index 29007bc..b3ed39e 100644 --- a/index.ts +++ b/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(); 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 | 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 | 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((_, 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"}`, diff --git a/test/unit.test.ts b/test/unit.test.ts index 4d90f01..3e4ff28 100644 --- a/test/unit.test.ts +++ b/test/unit.test.ts @@ -31,6 +31,7 @@ import { ensureDirs, ensureQmdEmbed, forgetBlocks, + getEmbedProbeTimeoutMs, getExitSummaryReasoningEffort, getExitSummaryTimeoutMs, getQmdSearchTimeoutMs, @@ -38,6 +39,7 @@ import { 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); +});