From f0cf8a59d29e5d8da02cf608924523f0710576c9 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Thu, 23 Apr 2026 22:06:55 +0200 Subject: [PATCH] fix(coding-agent): handle stale extension contexts fixes #3606 --- packages/coding-agent/CHANGELOG.md | 6 +++++ packages/coding-agent/docs/extensions.md | 21 ++++++++++----- .../examples/extensions/handoff.ts | 13 +++++----- .../src/core/agent-session-runtime.ts | 15 +++++++++++ .../coding-agent/src/core/agent-session.ts | 2 +- .../src/core/extensions/loader.ts | 2 +- .../src/core/extensions/runner.ts | 4 ++- .../src/modes/interactive/interactive-mode.ts | 5 ++-- .../test/agent-session-runtime-events.test.ts | 26 +++++++++++++++++++ 9 files changed, 77 insertions(+), 17 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 79863b9d..3ef9415c 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -18,8 +18,14 @@ - Added searchable auth provider login flow with fuzzy filtering in the provider selector ([#3572](https://github.com/badlogic/pi-mono/pull/3572) by [@mitsuhiko](https://github.com/mitsuhiko)) - Added GPT-5.5 Codex model +### Changed + +- Improved stale extension context errors after session replacement or reload to tell extension authors to avoid captured `pi`/command `ctx` and use `withSession` for post-replacement work. + ### Fixed +- Fixed the handoff extension example to use the replacement-session context after creating a new session, avoiding stale `ctx` errors when it installs the generated prompt ([#3606](https://github.com/badlogic/pi-mono/issues/3606)) +- Fixed session replacement and `/quit` teardown ordering to run host-owned extension UI cleanup synchronously after `session_shutdown` handlers complete but before invalidating the old extension context, preventing stale extension UI from rendering against a disposed session. - Fixed crash on `/quit` when an extension registers a custom footer whose `render()` accesses `ctx`, by tearing down extension-provided UI before invalidating the extension runner during shutdown ([#3595](https://github.com/badlogic/pi-mono/issues/3595)) - Fixed auto-retry to treat Bedrock/Smithy HTTP/2 transport failures like `http2 request did not get a response` as transient errors, so the agent retries automatically instead of waiting for a manual nudge ([#3594](https://github.com/badlogic/pi-mono/issues/3594)) - Fixed the CLI/SDK tool-selection split so `--no-builtin-tools` and `createAgentSession({ noTools: "builtin" })` disable only built-in default tools while keeping extension/custom tools enabled, instead of falling through to the same "disable everything" path as `--no-tools` ([#3592](https://github.com/badlogic/pi-mono/issues/3592)) diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 56da258b..37d92896 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -992,14 +992,19 @@ Options: Fork from a specific entry, creating a new session file: ```typescript -const result = await ctx.fork("entry-id-123"); -if (!result.cancelled) { - // Now in the forked session +const result = await ctx.fork("entry-id-123", { + withSession: async (ctx) => { + // Use only the replacement-session ctx here. + ctx.ui.notify("Now in the forked session", "info"); + }, +}); +if (result.cancelled) { + // An extension cancelled the fork } const cloneResult = await ctx.fork("entry-id-456", { position: "at" }); -if (!cloneResult.cancelled) { - // New session contains the active path through entry-id-456 +if (cloneResult.cancelled) { + // An extension cancelled the clone } ``` @@ -1060,7 +1065,11 @@ pi.registerCommand("switch", { sessions.map(s => s.file), ); if (choice) { - await ctx.switchSession(choice); + await ctx.switchSession(choice, { + withSession: async (ctx) => { + ctx.ui.notify("Switched session", "info"); + }, + }); } }, }); diff --git a/packages/coding-agent/examples/extensions/handoff.ts b/packages/coding-agent/examples/extensions/handoff.ts index 3f115a92..0f897aef 100644 --- a/packages/coding-agent/examples/extensions/handoff.ts +++ b/packages/coding-agent/examples/extensions/handoff.ts @@ -135,19 +135,20 @@ export default function (pi: ExtensionAPI) { return; } - // Create new session with parent tracking + // Create new session with parent tracking. Use the replacement-session + // context for post-switch UI work; the original ctx is stale after a + // successful session replacement. const newSessionResult = await ctx.newSession({ parentSession: currentSessionFile, + withSession: async (replacementCtx) => { + replacementCtx.ui.setEditorText(editedPrompt); + replacementCtx.ui.notify("Handoff ready. Submit when ready.", "info"); + }, }); if (newSessionResult.cancelled) { ctx.ui.notify("New session cancelled", "info"); - return; } - - // Set the edited prompt in the main editor for submission - ctx.ui.setEditorText(editedPrompt); - ctx.ui.notify("Handoff ready. Submit when ready.", "info"); }, }); } diff --git a/packages/coding-agent/src/core/agent-session-runtime.ts b/packages/coding-agent/src/core/agent-session-runtime.ts index 72d937d5..b64def1b 100644 --- a/packages/coding-agent/src/core/agent-session-runtime.ts +++ b/packages/coding-agent/src/core/agent-session-runtime.ts @@ -66,6 +66,7 @@ function extractUserMessageText(content: string | Array<{ type: string; text?: s */ export class AgentSessionRuntime { private rebindSession?: (session: AgentSession) => Promise; + private beforeSessionInvalidate?: () => void; constructor( private _session: AgentSession, @@ -99,6 +100,18 @@ export class AgentSessionRuntime { this.rebindSession = rebindSession; } + /** + * Set a synchronous callback that runs after `session_shutdown` handlers finish + * but before the current session is invalidated. + * + * This is for host-owned UI teardown that must not yield to the event loop, + * such as detaching extension-provided TUI components before the old extension + * context becomes stale. + */ + setBeforeSessionInvalidate(beforeSessionInvalidate?: () => void): void { + this.beforeSessionInvalidate = beforeSessionInvalidate; + } + private async emitBeforeSwitch( reason: "new" | "resume", targetSessionFile?: string, @@ -139,6 +152,7 @@ export class AgentSessionRuntime { reason, targetSessionFile, }); + this.beforeSessionInvalidate?.(); this.session.dispose(); } @@ -354,6 +368,7 @@ export class AgentSessionRuntime { type: "session_shutdown", reason: "quit", }); + this.beforeSessionInvalidate?.(); this.session.dispose(); } } diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 378c78d0..f303786a 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -724,7 +724,7 @@ export class AgentSession { */ dispose(): void { this._extensionRunner.invalidate( - "This extension instance is stale after session replacement or reload. Use the provided replacement-session context instead.", + "This extension ctx is stale after session replacement or reload. Do not use a captured pi or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload().", ); this._disconnectFromAgent(); this._eventListeners = []; diff --git a/packages/coding-agent/src/core/extensions/loader.ts b/packages/coding-agent/src/core/extensions/loader.ts index a62f8ed1..b72e0655 100644 --- a/packages/coding-agent/src/core/extensions/loader.ts +++ b/packages/coding-agent/src/core/extensions/loader.ts @@ -164,7 +164,7 @@ export function createExtensionRuntime(): ExtensionRuntime { invalidate: (message) => { state.staleMessage ??= message ?? - "This extension instance is stale after session replacement or reload. Use the provided replacement-session context instead."; + "This extension ctx is stale after session replacement or reload. Do not use a captured pi or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload()."; }, // Pre-bind: queue registrations so bindCore() can flush them once the // model registry is available. bindCore() replaces both with direct calls. diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 14278d79..27715439 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -458,7 +458,9 @@ export class ExtensionRunner { return this.shortcutDiagnostics; } - invalidate(message = "This extension instance is stale after session replacement or reload."): void { + invalidate( + message = "This extension ctx is stale after session replacement or reload. Do not use a captured pi or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload().", + ): void { if (!this.staleMessage) { this.staleMessage = message; this.runtime.invalidate(message); diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index ae8df81c..57c21e33 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -362,6 +362,9 @@ export class InteractiveMode { private options: InteractiveModeOptions = {}, ) { this.runtimeHost = runtimeHost; + this.runtimeHost.setBeforeSessionInvalidate(() => { + this.resetExtensionUI(); + }); this.runtimeHost.setRebindSession(async () => { await this.rebindCurrentSession(); }); @@ -1594,7 +1597,6 @@ export class InteractiveMode { } private async rebindCurrentSession(): Promise { - this.resetExtensionUI(); this.unsubscribe?.(); this.unsubscribe = undefined; this.applyRuntimeSettings(); @@ -3219,7 +3221,6 @@ export class InteractiveMode { if (this.isShuttingDown) return; this.isShuttingDown = true; this.unregisterSignalHandlers(); - this.resetExtensionUI(); await this.runtimeHost.dispose(); // Wait for any pending renders to complete diff --git a/packages/coding-agent/test/agent-session-runtime-events.test.ts b/packages/coding-agent/test/agent-session-runtime-events.test.ts index f21f7997..94b770ea 100644 --- a/packages/coding-agent/test/agent-session-runtime-events.test.ts +++ b/packages/coding-agent/test/agent-session-runtime-events.test.ts @@ -157,6 +157,32 @@ describe("AgentSessionRuntime session lifecycle events", () => { expect(events).toEqual([{ type: "session_before_switch", reason: "new", targetSessionFile: undefined }]); }); + it("runs beforeSessionInvalidate after session_shutdown and before rebindSession", async () => { + const phases: string[] = []; + const { runtimeHost } = await createRuntimeHost((pi) => { + pi.on("session_shutdown", () => { + phases.push("session_shutdown"); + }); + }); + const oldSession = runtimeHost.session; + runtimeHost.setBeforeSessionInvalidate(() => { + phases.push("beforeSessionInvalidate"); + expect(oldSession.extensionRunner.createContext().cwd).toBe(oldSession.sessionManager.getCwd()); + }); + runtimeHost.setRebindSession(async () => { + phases.push("rebindSession"); + }); + + await runtimeHost.newSession(); + + expect(phases).toEqual(["session_shutdown", "beforeSessionInvalidate", "rebindSession"]); + expect(() => oldSession.extensionRunner.createContext().cwd).toThrow( + "This extension ctx is stale after session replacement or reload. Do not use a captured pi or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload().", + ); + runtimeHost.setBeforeSessionInvalidate(undefined); + runtimeHost.setRebindSession(undefined); + }); + it("emits session_before_fork and session_start and honors cancellation", async () => { const events: RecordedSessionEvent[] = []; let cancelNextFork = false;