From 0c9944e8d31d65fad3d18bbb35c70b5376da053e Mon Sep 17 00:00:00 2001 From: aserper Date: Tue, 28 Apr 2026 21:27:33 -0400 Subject: [PATCH] fix: speed up lifecycle transitions in pi-memory Signed-off-by: aserper --- README.md | 1 + index.ts | 85 +++++++++++++++++++++++++++++++++++------------ test/qmd-cache.ts | 76 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 140 insertions(+), 22 deletions(-) create mode 100644 test/qmd-cache.ts diff --git a/README.md b/README.md index 09d8df2..00afaf2 100644 --- a/README.md +++ b/README.md @@ -136,6 +136,7 @@ This ensures in-progress context survives compaction and is visible in the next |----------|--------|---------|-------------| | `PI_MEMORY_QMD_UPDATE` | `background`, `manual`, `off` | `background` | Controls automatic `qmd update` after writes | | `PI_MEMORY_NO_SEARCH` | `1` | unset | Disable selective injection (for A/B testing) | +| `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. | ## Running tests diff --git a/index.ts b/index.ts index 64910d7..623921c 100644 --- a/index.ts +++ b/index.ts @@ -443,6 +443,17 @@ function getQmdUpdateMode(): "background" | "manual" | "off" { return "background"; } +export function shouldSummarizeLifecycleTransitions(): boolean { + const value = (process.env.PI_MEMORY_SUMMARIZE_TRANSITIONS ?? "").toLowerCase(); + return value === "1" || value === "true" || value === "yes" || value === "on"; +} + +export function shouldSkipExitSummaryForReason(reason: string | undefined): boolean { + if (!reason) return false; + if (shouldSummarizeLifecycleTransitions()) return false; + return ["reload", "new", "resume", "fork"].includes(reason); +} + async function ensureQmdAvailableForUpdate(): Promise { if (qmdAvailable) return true; if (getQmdUpdateMode() !== "background") return false; @@ -596,6 +607,9 @@ type ExecFileFn = typeof execFile; let execFileFn: ExecFileFn = execFile; let qmdAvailable = false; +let qmdAvailabilityCheckedAt = 0; +const QMD_STATUS_CACHE_TTL_MS = 5 * 60 * 1000; +const qmdCollectionStatusCache = new Map(); let updateTimer: ReturnType | null = null; let exitSummaryReason: ExitSummaryReason | null = null; let terminalInputUnsubscribe: (() => void) | null = null; @@ -613,6 +627,7 @@ export function _resetExecFileForTest() { /** Set qmd availability flag (for testing). */ export function _setQmdAvailable(value: boolean) { qmdAvailable = value; + qmdAvailabilityCheckedAt = Date.now(); } /** Get current qmd availability flag (for testing). */ @@ -633,6 +648,12 @@ export function _clearUpdateTimer() { } } +/** Clear qmd status caches (for testing). */ +export function _clearQmdStatusCaches() { + qmdAvailabilityCheckedAt = 0; + qmdCollectionStatusCache.clear(); +} + const QMD_REPO_URL = "https://github.com/tobi/qmd"; export function qmdInstallInstructions(): string { @@ -692,41 +713,53 @@ export async function setupQmdCollection(): Promise { } export function detectQmd(): Promise { + const now = Date.now(); + if (qmdAvailabilityCheckedAt && now - qmdAvailabilityCheckedAt < QMD_STATUS_CACHE_TTL_MS) { + return Promise.resolve(qmdAvailable); + } + return new Promise((resolve) => { // qmd doesn't reliably support --version; use a fast command that exits 0 when available. execFileFn("qmd", ["status"], { timeout: 5_000 }, (err) => { - resolve(!err); + qmdAvailable = !err; + qmdAvailabilityCheckedAt = Date.now(); + resolve(qmdAvailable); }); }); } export function checkCollection(name: string): Promise { + const cached = qmdCollectionStatusCache.get(name); + const now = Date.now(); + if (cached && now - cached.checkedAt < QMD_STATUS_CACHE_TTL_MS) { + return Promise.resolve(cached.exists); + } + return new Promise((resolve) => { execFileFn("qmd", ["collection", "list", "--json"], { timeout: 10_000 }, (err, stdout) => { - if (err) { - resolve(false); - return; - } - try { - const collections = JSON.parse(stdout); - if (Array.isArray(collections)) { - resolve( - collections.some((entry) => { + let exists = false; + if (!err) { + try { + const collections = JSON.parse(stdout); + if (Array.isArray(collections)) { + exists = collections.some((entry) => { if (typeof entry === "string") return entry === name; if (entry && typeof entry === "object" && "name" in entry) { return (entry as { name?: string }).name === name; } return false; - }), - ); - } else { - // qmd may output an object with a collections array or similar - resolve(stdout.includes(name)); + }); + } else { + // qmd may output an object with a collections array or similar + exists = stdout.includes(name); + } + } catch { + // Fallback: just check if the name appears in the output + exists = stdout.includes(name); } - } catch { - // Fallback: just check if the name appears in the output - resolve(stdout.includes(name)); } + qmdCollectionStatusCache.set(name, { checkedAt: Date.now(), exists }); + resolve(exists); }); }); } @@ -907,10 +940,18 @@ export default function (pi: ExtensionAPI) { terminalInputUnsubscribe = null; } - // /reload emits session_shutdown with reason "reload" before rebuilding the - // runtime. Generating an exit summary here would make every /reload block - // for several seconds on a live LLM call. Skip it — the session continues. - if (event.reason === "reload") { + // Lifecycle transitions are usually not final session exits. By default, + // avoid generating LLM summaries and running qmd updates during /reload, + // /new, /resume, and /fork because that makes those transitions slow. + // Users who prefer the old behavior can opt in with + // PI_MEMORY_SUMMARIZE_TRANSITIONS=1. + const shutdownReason = (event as { reason?: string }).reason; + if (shouldSkipExitSummaryForReason(shutdownReason)) { + exitSummaryReason = null; + if (updateTimer) { + clearTimeout(updateTimer); + updateTimer = null; + } return; } diff --git a/test/qmd-cache.ts b/test/qmd-cache.ts new file mode 100644 index 0000000..e847969 --- /dev/null +++ b/test/qmd-cache.ts @@ -0,0 +1,76 @@ +import { strict as assert } from "node:assert"; +import type { execFile } from "node:child_process"; +import { + _clearQmdStatusCaches, + _resetExecFileForTest, + _setExecFileForTest, + _setQmdAvailable, + checkCollection, + detectQmd, + shouldSkipExitSummaryForReason, +} from "../index.ts"; + +type ExecFileFn = typeof execFile; + +type ExecCallback = (error: Error | null, stdout: string, stderr: string) => void; + +function mockExecFile(handler: (cmd: string, args: readonly string[]) => { error?: Error; stdout?: string }) { + let calls = 0; + const fn: ExecFileFn = ((cmd: string, args: readonly string[], _options: unknown, callback: ExecCallback) => { + calls++; + const result = handler(cmd, args); + queueMicrotask(() => callback(result.error ?? null, result.stdout ?? "", "")); + }) as ExecFileFn; + _setExecFileForTest(fn); + return () => calls; +} + +try { + _clearQmdStatusCaches(); + const qmdCalls = mockExecFile((cmd, args) => { + assert.equal(cmd, "qmd"); + assert.deepEqual(args, ["status"]); + return {}; + }); + + assert.equal(await detectQmd(), true); + assert.equal(await detectQmd(), true); + assert.equal(qmdCalls(), 1, "detectQmd should cache qmd status within the TTL"); + + _clearQmdStatusCaches(); + const collectionCalls = mockExecFile((cmd, args) => { + assert.equal(cmd, "qmd"); + assert.deepEqual(args, ["collection", "list", "--json"]); + return { stdout: JSON.stringify([{ name: "pi-memory" }]) }; + }); + + assert.equal(await checkCollection("pi-memory"), true); + assert.equal(await checkCollection("pi-memory"), true); + assert.equal(collectionCalls(), 1, "checkCollection should cache collection lookup within the TTL"); + + _setQmdAvailable(false); + assert.equal(await detectQmd(), false, "_setQmdAvailable should seed the cached status"); + + const originalSummarizeTransitions = process.env.PI_MEMORY_SUMMARIZE_TRANSITIONS; + try { + delete process.env.PI_MEMORY_SUMMARIZE_TRANSITIONS; + assert.equal(shouldSkipExitSummaryForReason("reload"), true); + assert.equal(shouldSkipExitSummaryForReason("new"), true); + assert.equal(shouldSkipExitSummaryForReason("session-end"), false); + + process.env.PI_MEMORY_SUMMARIZE_TRANSITIONS = "1"; + assert.equal(shouldSkipExitSummaryForReason("reload"), false); + assert.equal(shouldSkipExitSummaryForReason("new"), false); + } finally { + if (originalSummarizeTransitions === undefined) { + delete process.env.PI_MEMORY_SUMMARIZE_TRANSITIONS; + } else { + process.env.PI_MEMORY_SUMMARIZE_TRANSITIONS = originalSummarizeTransitions; + } + } + + console.log("qmd cache tests passed"); +} finally { + _resetExecFileForTest(); + _clearQmdStatusCaches(); +}