Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
85 changes: 63 additions & 22 deletions index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<boolean> {
if (qmdAvailable) return true;
if (getQmdUpdateMode() !== "background") return false;
Expand Down Expand Up @@ -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<string, { checkedAt: number; exists: boolean }>();
let updateTimer: ReturnType<typeof setTimeout> | null = null;
let exitSummaryReason: ExitSummaryReason | null = null;
let terminalInputUnsubscribe: (() => void) | null = null;
Expand All @@ -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). */
Expand All @@ -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 {
Expand Down Expand Up @@ -692,41 +713,53 @@ export async function setupQmdCollection(): Promise<boolean> {
}

export function detectQmd(): Promise<boolean> {
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<boolean> {
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);
});
});
}
Expand Down Expand Up @@ -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;
}

Expand Down
76 changes: 76 additions & 0 deletions test/qmd-cache.ts
Original file line number Diff line number Diff line change
@@ -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();
}