From c0ce827c012b71efd0d2edfc590857f7384f6bc6 Mon Sep 17 00:00:00 2001 From: NQH Date: Sun, 26 Apr 2026 18:52:32 -0500 Subject: [PATCH] fix: guard autoresearch shutdown context Commit-Queue-Session: cq_20260426235232_ecf71f8eb10a211a19af3f6b Coding-Agent: pi Coding-Agent-Session: pi-019dcbcc-b152-70a0-b167-aeaaaf7a9b32 --- extensions/pi-autoresearch/index.ts | 114 ++++++++++++++++++++-------- tests/stale-context-source.test.mjs | 21 +++++ 2 files changed, 102 insertions(+), 33 deletions(-) create mode 100644 tests/stale-context-source.test.mjs diff --git a/extensions/pi-autoresearch/index.ts b/extensions/pi-autoresearch/index.ts index 326698bd..8e9df772 100644 --- a/extensions/pi-autoresearch/index.ts +++ b/extensions/pi-autoresearch/index.ts @@ -1033,6 +1033,28 @@ export default function autoresearchExtension(pi: ExtensionAPI) { // Outlasts pi's internal retry (setTimeout 0) and compaction-continue // (setTimeout 100); see badlogic/pi-mono#2023, #2110. const SETTLED_WINDOW_MS = 800; + const STALE_EXTENSION_CONTEXT_MARKER = "extension ctx is stale after session replacement or reload"; + + const isStaleExtensionContextError = (error: unknown): boolean => + error instanceof Error && error.message.includes(STALE_EXTENSION_CONTEXT_MARKER); + + const runIgnoringStaleExtensionContext = (run: () => void): void => { + try { + run(); + } catch (error) { + if (isStaleExtensionContextError(error)) return; + throw error; + } + }; + + const runIgnoringStaleExtensionContextAsync = async (run: () => Promise): Promise => { + try { + await run(); + } catch (error) { + if (isStaleExtensionContextError(error)) return; + throw error; + } + }; const runtimeStore = createRuntimeStore(); const getSessionKey = (ctx: ExtensionContext) => ctx.sessionManager.getSessionId(); @@ -1086,7 +1108,9 @@ export default function autoresearchExtension(pi: ExtensionAPI) { pausePendingResume(runtime); runtime.pendingResumeMessage = message; runtime.pendingResumeTimer = setTimeout( - () => sendPendingResumeIfReady(ctx, runtime), + () => { + runIgnoringStaleExtensionContext(() => sendPendingResumeIfReady(ctx, runtime)); + }, SETTLED_WINDOW_MS, ); }; @@ -1113,10 +1137,12 @@ export default function autoresearchExtension(pi: ExtensionAPI) { runtime.autoResumeTurns >= MAX_AUTORESUME_TURNS; const notifyAutoResumeLimitReached = (ctx: ExtensionContext): void => { - ctx.ui.notify( - `Autoresearch auto-resume limit reached (${MAX_AUTORESUME_TURNS} turns)`, - "info", - ); + runIgnoringStaleExtensionContext(() => { + ctx.ui.notify( + `Autoresearch auto-resume limit reached (${MAX_AUTORESUME_TURNS} turns)`, + "info", + ); + }); }; const hasIdeasFile = (ctx: ExtensionContext): boolean => @@ -1191,9 +1217,11 @@ export default function autoresearchExtension(pi: ExtensionAPI) { const clearSessionUi = (ctx: ExtensionContext) => { clearOverlay(); - if (ctx.hasUI) { - ctx.ui.setWidget("autoresearch", undefined); - } + runIgnoringStaleExtensionContext(() => { + if (ctx.hasUI) { + ctx.ui.setWidget("autoresearch", undefined); + } + }); }; const autoresearchHelp = () => @@ -1299,7 +1327,11 @@ export default function autoresearchExtension(pi: ExtensionAPI) { }; const updateWidget = (ctx: ExtensionContext) => { - if (!ctx.hasUI) return; + let hasUI = false; + runIgnoringStaleExtensionContext(() => { + hasUI = ctx.hasUI; + }); + if (!hasUI) return; const runtime = getRuntime(ctx); const state = runtime.state; @@ -1460,48 +1492,64 @@ export default function autoresearchExtension(pi: ExtensionAPI) { } }; - pi.on("session_start", async (_e, ctx) => reconstructState(ctx)); - pi.on("session_tree", async (_e, ctx) => reconstructState(ctx)); + pi.on("session_start", async (_e, ctx) => { + await runIgnoringStaleExtensionContextAsync(async () => reconstructState(ctx)); + }); + pi.on("session_tree", async (_e, ctx) => { + await runIgnoringStaleExtensionContextAsync(async () => reconstructState(ctx)); + }); pi.on("session_before_switch", async () => { clearOverlay(); }); pi.on("session_shutdown", async (_e, ctx) => { - clearSessionUi(ctx); - cancelPendingResume(getRuntime(ctx)); - runtimeStore.clear(getSessionKey(ctx)); - stopDashboardServer(); + runIgnoringStaleExtensionContext(() => { + const sessionKey = getSessionKey(ctx); + const runtime = getRuntime(ctx); + cancelPendingResume(runtime); + clearSessionUi(ctx); + runtimeStore.clear(sessionKey); + stopDashboardServer(); + }); }); pi.on("agent_start", async (_event, ctx) => { - const runtime = getRuntime(ctx); - runtime.experimentsThisSession = 0; - pausePendingResume(runtime); + await runIgnoringStaleExtensionContextAsync(async () => { + const runtime = getRuntime(ctx); + runtime.experimentsThisSession = 0; + pausePendingResume(runtime); + }); }); pi.on("session_before_compact", async (_event, ctx) => { - pausePendingResume(getRuntime(ctx)); + await runIgnoringStaleExtensionContextAsync(async () => { + pausePendingResume(getRuntime(ctx)); + }); }); pi.on("session_compact", async (_event, ctx) => { - reschedulePendingResume(ctx, getRuntime(ctx)); + await runIgnoringStaleExtensionContextAsync(async () => { + reschedulePendingResume(ctx, getRuntime(ctx)); + }); }); pi.on("agent_end", async (_event, ctx) => { - const runtime = getRuntime(ctx); - runtime.runningExperiment = null; - if (overlayTui) overlayTui.requestRender(); + await runIgnoringStaleExtensionContextAsync(async () => { + const runtime = getRuntime(ctx); + runtime.runningExperiment = null; + if (overlayTui) overlayTui.requestRender(); - if (hasPendingResume(runtime)) { - reschedulePendingResume(ctx, runtime); - return; - } - if (!shouldAutoResume(runtime)) return; - if (hasReachedAutoResumeLimit(runtime)) { - notifyAutoResumeLimitReached(ctx); - return; - } + if (hasPendingResume(runtime)) { + reschedulePendingResume(ctx, runtime); + return; + } + if (!shouldAutoResume(runtime)) return; + if (hasReachedAutoResumeLimit(runtime)) { + notifyAutoResumeLimitReached(ctx); + return; + } - schedulePendingResume(ctx, runtime, composeResumeMessage(ctx)); + schedulePendingResume(ctx, runtime, composeResumeMessage(ctx)); + }); }); // When in autoresearch mode, add a static note to the system prompt. diff --git a/tests/stale-context-source.test.mjs b/tests/stale-context-source.test.mjs new file mode 100644 index 00000000..33e43c2e --- /dev/null +++ b/tests/stale-context-source.test.mjs @@ -0,0 +1,21 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { readFileSync } from "node:fs"; + +const source = readFileSync(new URL("../extensions/pi-autoresearch/index.ts", import.meta.url), "utf-8"); + +test("autoresearch guards UI cleanup against stale Pi contexts", () => { + assert.match(source, /function runIgnoringStaleExtensionContext|const runIgnoringStaleExtensionContext/, "expected stale-context guard helper"); + assert.match(source, /runIgnoringStaleExtensionContext\(\(\) => \{\s*if \(ctx\.hasUI\)/s, "expected ctx.hasUI UI cleanup to be wrapped by stale-context guard"); +}); + +test("autoresearch cancels pending resume timers before session shutdown UI work", () => { + const shutdownIndex = source.indexOf('pi.on("session_shutdown"'); + assert.notEqual(shutdownIndex, -1, "expected session_shutdown handler"); + const shutdownBlock = source.slice(shutdownIndex, shutdownIndex + 500); + assert.match(shutdownBlock, /cancelPendingResume/, "expected shutdown to cancel pending auto-resume timers"); + assert.ok( + shutdownBlock.indexOf("cancelPendingResume") < shutdownBlock.indexOf("clearSessionUi"), + "expected auto-resume cancellation before UI cleanup can touch stale ctx", + ); +});