From 7d3bcb6d1393de6b17b3d8637dd205c96ba07e7c Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Sun, 14 Jun 2026 22:30:11 +0200 Subject: [PATCH] fix(coding-agent): defer extension reload requests safely --- packages/coding-agent/docs/extensions.md | 100 ++--- .../examples/extensions/README.md | 2 +- .../examples/extensions/reload-runtime.ts | 25 +- .../src/core/agent-session-runtime.ts | 45 +- .../coding-agent/src/core/agent-session.ts | 188 +++++++-- .../src/core/extensions/runner.ts | 12 +- .../coding-agent/src/core/extensions/types.ts | 8 +- .../src/modes/interactive/interactive-mode.ts | 45 +- .../test/agent-session-reload-request.test.ts | 387 ++++++++++++++++++ .../test/extensions-runner.test.ts | 12 + .../test/trigger-compact-extension.test.ts | 1 + 11 files changed, 680 insertions(+), 145 deletions(-) create mode 100644 packages/coding-agent/test/agent-session-reload-request.test.ts diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 7038351d1..e00761f86 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -994,6 +994,48 @@ ctx.compact({ }); ``` +### ctx.reload() + +Request the same runtime reload flow as `/reload`. This is available in `ExtensionContext`, so tools, event handlers, shortcuts, and command handlers can call it. + +`ctx.reload()` records a reload request and returns a promise that resolves once the request is accepted. The reload runs when the session reaches a safe boundary. In command handlers this may happen immediately; in tools and event handlers it is deferred so reload does not re-enter while hook results are still being applied. + +Reload behavior: +- Reload emits `session_shutdown` for the current extension runtime +- It then reloads resources and emits `session_start` with `reason: "reload"` and `resources_discover` with reason `"reload"` +- The currently running handler continues in the old call frame +- Code after `await ctx.reload()` must not assume old in-memory extension state will remain valid +- After the handler returns and reload completes, future commands/events/tool calls use the new extension version +- A reload request made while a reload is already running queues one follow-up reload when the session is idle + +For predictable command behavior, treat reload as terminal for that handler (`await ctx.reload(); return;`). + +Example tool the LLM can call to request reload: + +```typescript +import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; +import { Type } from "typebox"; + +export default function (pi: ExtensionAPI) { + pi.registerTool({ + name: "reload_runtime", + label: "Reload Runtime", + description: "Reload extensions, skills, prompts, and themes", + parameters: Type.Object({}), + async execute(_toolCallId, _params, _signal, _onUpdate, ctx) { + await ctx.reload(); + return { + content: [{ type: "text", text: "Reload requested. It will run after this turn finishes." }], + details: {}, + terminate: true, + }; + }, + }); +} +``` + +Return `terminate: true` when reload should happen as soon as the current tool batch finishes; otherwise the agent may continue with another LLM turn before it becomes idle. + ### ctx.getSystemPrompt() Returns Pi's current system prompt string. @@ -1204,62 +1246,6 @@ pi.registerCommand("handoff", { }); ``` -### ctx.reload() - -Run the same reload flow as `/reload`. - -```typescript -pi.registerCommand("reload-runtime", { - description: "Reload extensions, skills, prompts, and themes", - handler: async (_args, ctx) => { - await ctx.reload(); - return; - }, -}); -``` - -Important behavior: -- `await ctx.reload()` emits `session_shutdown` for the current extension runtime -- It then reloads resources and emits `session_start` with `reason: "reload"` and `resources_discover` with reason `"reload"` -- The currently running command handler still continues in the old call frame -- Code after `await ctx.reload()` still runs from the pre-reload version -- Code after `await ctx.reload()` must not assume old in-memory extension state is still valid -- After the handler returns, future commands/events/tool calls use the new extension version - -For predictable behavior, treat reload as terminal for that handler (`await ctx.reload(); return;`). - -Tools run with `ExtensionContext`, so they cannot call `ctx.reload()` directly. Use a command as the reload entrypoint, then expose a tool that queues that command as a follow-up user message. - -Example tool the LLM can call to trigger reload: - -```typescript -import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; -import { Type } from "typebox"; - -export default function (pi: ExtensionAPI) { - pi.registerCommand("reload-runtime", { - description: "Reload extensions, skills, prompts, and themes", - handler: async (_args, ctx) => { - await ctx.reload(); - return; - }, - }); - - pi.registerTool({ - name: "reload_runtime", - label: "Reload Runtime", - description: "Reload extensions, skills, prompts, and themes", - parameters: Type.Object({}), - async execute() { - pi.sendUserMessage("/reload-runtime", { deliverAs: "followUp" }); - return { - content: [{ type: "text", text: "Queued /reload-runtime as a follow-up command." }], - }; - }, - }); -} -``` - ## ExtensionAPI Methods ### pi.on(event, handler) @@ -2597,7 +2583,7 @@ All examples in [examples/extensions/](../examples/extensions/). | `handoff.ts` | Cross-provider model handoff | `registerCommand`, `ui.editor`, `ui.custom` | | `qna.ts` | Q&A with custom UI | `registerCommand`, `ui.custom`, `setEditorText` | | `send-user-message.ts` | Inject user messages | `registerCommand`, `sendUserMessage` | -| `reload-runtime.ts` | Reload command and LLM tool handoff | `registerCommand`, `ctx.reload()`, `sendUserMessage` | +| `reload-runtime.ts` | LLM tool reload request | `registerTool`, `ctx.reload()` | | `shutdown-command.ts` | Graceful shutdown command | `registerCommand`, `shutdown()` | | **Events & Gates** ||| | `permission-gate.ts` | Block dangerous commands | `on("tool_call")`, `ui.confirm` | diff --git a/packages/coding-agent/examples/extensions/README.md b/packages/coding-agent/examples/extensions/README.md index d4f4f2ccf..3d5751afd 100644 --- a/packages/coding-agent/examples/extensions/README.md +++ b/packages/coding-agent/examples/extensions/README.md @@ -74,7 +74,7 @@ cp permission-gate.ts ~/.pi/agent/extensions/ | `overlay-qa-tests.ts` | Comprehensive overlay QA tests: anchors, margins, stacking, overflow, animation | | `doom-overlay/` | DOOM game running as an overlay at 35 FPS (demonstrates real-time game rendering) | | `shutdown-command.ts` | Adds `/quit` command demonstrating `ctx.shutdown()` | -| `reload-runtime.ts` | Adds `/reload-runtime` and `reload_runtime` tool showing safe reload flow | +| `reload-runtime.ts` | Adds a `reload_runtime` tool showing safe direct reload requests | | `interactive-shell.ts` | Run interactive commands (vim, htop) with full terminal via `user_bash` hook | | `inline-bash.ts` | Expands `!{command}` patterns in prompts via `input` event transformation | | `input-transform-streaming.ts` | Skips expensive input preprocessing for mid-stream steering via `streamingBehavior` | diff --git a/packages/coding-agent/examples/extensions/reload-runtime.ts b/packages/coding-agent/examples/extensions/reload-runtime.ts index be46410f2..7d8f64998 100644 --- a/packages/coding-agent/examples/extensions/reload-runtime.ts +++ b/packages/coding-agent/examples/extensions/reload-runtime.ts @@ -1,36 +1,27 @@ /** * Reload Runtime Extension * - * Demonstrates ctx.reload() from ExtensionCommandContext and an LLM-callable - * tool that queues a follow-up command to trigger reload. + * Demonstrates ctx.reload() from an LLM-callable tool. Tool-triggered + * reloads are deferred until the agent turn is idle. */ import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; import { Type } from "typebox"; export default function (pi: ExtensionAPI) { - // Command entrypoint for reload. - // Treat reload as terminal for this handler. - pi.registerCommand("reload-runtime", { - description: "Reload extensions, skills, prompts, and themes", - handler: async (_args, ctx) => { - await ctx.reload(); - return; - }, - }); - - // LLM-callable tool. Tools get ExtensionContext, so they cannot call ctx.reload() directly. - // Instead, queue a follow-up user command that executes the command above. + // During a tool call, ctx.reload() requests a reload that + // runs after the current agent turn is fully idle. pi.registerTool({ name: "reload_runtime", label: "Reload Runtime", description: "Reload extensions, skills, prompts, and themes", parameters: Type.Object({}), - async execute() { - pi.sendUserMessage("/reload-runtime", { deliverAs: "followUp" }); + async execute(_toolCallId, _params, _signal, _onUpdate, ctx) { + await ctx.reload(); return { - content: [{ type: "text", text: "Queued /reload-runtime as a follow-up command." }], + content: [{ type: "text", text: "Reload requested. It will run after this turn finishes." }], details: {}, + terminate: true, }; }, }); diff --git a/packages/coding-agent/src/core/agent-session-runtime.ts b/packages/coding-agent/src/core/agent-session-runtime.ts index 7f29275aa..705a4146e 100644 --- a/packages/coding-agent/src/core/agent-session-runtime.ts +++ b/packages/coding-agent/src/core/agent-session-runtime.ts @@ -139,11 +139,13 @@ export class AgentSessionRuntime { return { cancelled: false }; } - const result = await runner.emit({ - type: "session_before_switch", - reason, - targetSessionFile, - }); + const result = await this.session.withReloadDeferred(() => + runner.emit({ + type: "session_before_switch", + reason, + targetSessionFile, + }), + ); return { cancelled: result?.cancel === true }; } @@ -156,20 +158,26 @@ export class AgentSessionRuntime { return { cancelled: false }; } - const result = await runner.emit({ - type: "session_before_fork", - entryId, - ...options, - }); + const result = await this.session.withReloadDeferred(() => + runner.emit({ + type: "session_before_fork", + entryId, + ...options, + }), + ); return { cancelled: result?.cancel === true }; } private async teardownCurrent(reason: SessionShutdownEvent["reason"], targetSessionFile?: string): Promise { - await emitSessionShutdownEvent(this.session.extensionRunner, { - type: "session_shutdown", - reason, - targetSessionFile, - }); + await this.session.withReloadDeferred( + () => + emitSessionShutdownEvent(this.session.extensionRunner, { + type: "session_shutdown", + reason, + targetSessionFile, + }), + { flush: false }, + ); this.beforeSessionInvalidate?.(); this.session.dispose(); } @@ -388,12 +396,7 @@ export class AgentSessionRuntime { } async dispose(): Promise { - await emitSessionShutdownEvent(this.session.extensionRunner, { - type: "session_shutdown", - reason: "quit", - }); - this.beforeSessionInvalidate?.(); - this.session.dispose(); + await this.teardownCurrent("quit"); } } diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index af9d32d7b..0b19a22fe 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -76,6 +76,8 @@ import { type TreePreparation, type TurnEndEvent, type TurnStartEvent, + type UserBashEvent, + type UserBashEventResult, wrapRegisteredTools, } from "./extensions/index.ts"; import { emitSessionShutdownEvent } from "./extensions/runner.ts"; @@ -190,6 +192,7 @@ export interface ExtensionBindings { uiContext?: ExtensionUIContext; mode?: ExtensionMode; commandContextActions?: ExtensionCommandContextActions; + reloadHandler?: RuntimeReloadHandler; abortHandler?: () => void; shutdownHandler?: ShutdownHandler; onError?: ExtensionErrorListener; @@ -242,6 +245,8 @@ interface ToolDefinitionEntry { sourceInfo: SourceInfo; } +type RuntimeReloadHandler = (reloadCore: () => Promise) => Promise; + // ============================================================================ // Constants // ============================================================================ @@ -304,6 +309,8 @@ export class AgentSession { private _extensionUIContext?: ExtensionUIContext; private _extensionMode: ExtensionMode = "print"; private _extensionCommandContextActions?: ExtensionCommandContextActions; + private _runtimeReloadHandler?: RuntimeReloadHandler; + private _extensionsBound = false; private _extensionAbortHandler?: () => void; private _extensionShutdownHandler?: ShutdownHandler; private _extensionErrorListener?: ExtensionErrorListener; @@ -322,6 +329,11 @@ export class AgentSession { private _baseSystemPrompt = ""; private _baseSystemPromptOptions!: BuildSystemPromptOptions; + // Deferred runtime reload state + private _reloadDeferralDepth = 0; + private _reloadRequested = false; + private _reloadInProgress = false; + constructor(config: AgentSessionConfig) { this.agent = config.agent; this.sessionManager = config.sessionManager; @@ -725,6 +737,7 @@ export class AgentSession { // Dispose must succeed even if an abort hook throws. } + this._reloadRequested = false; this._extensionRunner.invalidate( "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().", ); @@ -934,14 +947,16 @@ export class AgentSession { // ========================================================================= private async _runAgentPrompt(messages: AgentMessage | AgentMessage[]): Promise { - try { - await this.agent.prompt(messages); - while (await this._handlePostAgentRun()) { - await this.agent.continue(); + await this.withReloadDeferred(async () => { + try { + await this.agent.prompt(messages); + while (await this._handlePostAgentRun()) { + await this.agent.continue(); + } + } finally { + this._flushPendingBashMessages(); } - } finally { - this._flushPendingBashMessages(); - } + }); } private async _handlePostAgentRun(): Promise { @@ -986,6 +1001,8 @@ export class AgentSession { async prompt(text: string, options?: PromptOptions): Promise { const expandPromptTemplates = options?.expandPromptTemplates ?? true; const preflightResult = options?.preflightResult; + let preflightAccepted = false; + let reloadDeferred = false; let messages: AgentMessage[] | undefined; try { @@ -996,10 +1013,14 @@ export class AgentSession { if (handled) { // Extension command executed, no prompt to send preflightResult?.(true); + preflightAccepted = true; return; } } + this._reloadDeferralDepth++; + reloadDeferred = true; + // Emit input event for extension interception (before skill/template expansion) let currentText = text; let currentImages = options?.images; @@ -1012,6 +1033,7 @@ export class AgentSession { ); if (inputResult.action === "handled") { preflightResult?.(true); + preflightAccepted = true; return; } if (inputResult.action === "transform") { @@ -1040,6 +1062,7 @@ export class AgentSession { await this._queueSteer(expandedText, currentImages); } preflightResult?.(true); + preflightAccepted = true; return; } @@ -1123,17 +1146,27 @@ export class AgentSession { // Ensure we're using the base prompt (in case previous turn had modifications) this.agent.state.systemPrompt = this._baseSystemPrompt; } + + if (!messages) { + return; + } + + preflightResult?.(true); + preflightAccepted = true; + await this._runAgentPrompt(messages); } catch (error) { - preflightResult?.(false); + if (!preflightAccepted) { + preflightResult?.(false); + } throw error; + } finally { + if (reloadDeferred) { + this._reloadDeferralDepth--; + if (this._reloadDeferralDepth === 0) { + await this._flushRequestedReload(); + } + } } - - if (!messages) { - return; - } - - preflightResult?.(true); - await this._runAgentPrompt(messages); } /** @@ -1416,6 +1449,71 @@ export class AgentSession { await this.agent.waitForIdle(); } + async withReloadDeferred(callback: () => Promise, options: { flush?: boolean } = {}): Promise { + this._reloadDeferralDepth++; + try { + return await callback(); + } finally { + this._reloadDeferralDepth--; + if (this._reloadDeferralDepth === 0 && options.flush !== false) { + await this._flushRequestedReload(); + } + } + } + + private _shouldDeferReload(): boolean { + return this._reloadDeferralDepth > 0 || this.isStreaming || this.isCompacting; + } + + private async _runReloadHandler(): Promise { + if (this._reloadInProgress) { + this._reloadRequested = true; + return; + } + + this._reloadInProgress = true; + try { + do { + this._reloadRequested = false; + const reloadCore = () => this._reloadCore(); + await (this._runtimeReloadHandler ? this._runtimeReloadHandler(reloadCore) : reloadCore()); + } while (this._reloadRequested && !this._shouldDeferReload()); + } finally { + this._reloadInProgress = false; + } + } + + private async _flushRequestedReload(): Promise { + if (!this._reloadRequested || this._shouldDeferReload()) { + return; + } + + try { + await this._runReloadHandler(); + } catch (error) { + this._extensionRunner.emitError({ + extensionPath: "", + event: "reload", + error: error instanceof Error ? error.message : String(error), + stack: error instanceof Error ? error.stack : undefined, + }); + throw error; + } + } + + /** + * Request a runtime reload. If the agent or session is busy, reload is deferred + * until the current prompt, compaction, or branch summary operation finishes. + */ + async requestReload(): Promise { + this._reloadRequested = true; + if (this._shouldDeferReload() || this._reloadInProgress) { + return; + } + + await this._flushRequestedReload(); + } + // ========================================================================= // Model Management // ========================================================================= @@ -1426,11 +1524,13 @@ export class AgentSession { source: "set" | "cycle" | "restore", ): Promise { if (modelsAreEqual(previousModel, nextModel)) return; - await this._extensionRunner.emit({ - type: "model_select", - model: nextModel, - previousModel, - source, + await this.withReloadDeferred(async () => { + await this._extensionRunner.emit({ + type: "model_select", + model: nextModel, + previousModel, + source, + }); }); } @@ -1548,10 +1648,15 @@ export class AgentSession { this.settingsManager.setDefaultThinkingLevel(effectiveLevel); } this._emit({ type: "thinking_level_changed", level: effectiveLevel }); - void this._extensionRunner.emit({ - type: "thinking_level_select", - level: effectiveLevel, - previousLevel, + void this.withReloadDeferred(async () => { + await this._extensionRunner.emit({ + type: "thinking_level_select", + level: effectiveLevel, + previousLevel, + }); + }).catch(() => { + // _flushRequestedReload already emitted the extension error. setThinkingLevel + // is synchronous, so there is no caller to reject. }); } } @@ -1766,6 +1871,7 @@ export class AgentSession { } finally { this._compactionAbortController = undefined; this._reconnectToAgent(); + await this._flushRequestedReload(); } } @@ -2053,6 +2159,7 @@ export class AgentSession { return false; } finally { this._autoCompactionAbortController = undefined; + await this._flushRequestedReload(); } } @@ -2069,6 +2176,7 @@ export class AgentSession { } async bindExtensions(bindings: ExtensionBindings): Promise { + this._extensionsBound = true; if (bindings.uiContext !== undefined) { this._extensionUIContext = bindings.uiContext; } @@ -2078,6 +2186,9 @@ export class AgentSession { if (bindings.commandContextActions !== undefined) { this._extensionCommandContextActions = bindings.commandContextActions; } + if (bindings.reloadHandler !== undefined) { + this._runtimeReloadHandler = bindings.reloadHandler; + } if (bindings.abortHandler !== undefined) { this._extensionAbortHandler = bindings.abortHandler; } @@ -2089,8 +2200,10 @@ export class AgentSession { } this._applyExtensionBindings(this._extensionRunner); - await this._extensionRunner.emit(this._sessionStartEvent); - await this.extendResourcesFromExtensions(this._sessionStartEvent.reason === "reload" ? "reload" : "startup"); + await this.withReloadDeferred(async () => { + await this._extensionRunner.emit(this._sessionStartEvent); + await this.extendResourcesFromExtensions(this._sessionStartEvent.reason === "reload" ? "reload" : "startup"); + }); } private async extendResourcesFromExtensions(reason: "startup" | "reload"): Promise { @@ -2148,7 +2261,14 @@ export class AgentSession { private _applyExtensionBindings(runner: ExtensionRunner): void { runner.setUIContext(this._extensionUIContext, this._extensionMode); - runner.bindCommandContext(this._extensionCommandContextActions); + runner.bindCommandContext({ + waitForIdle: this._extensionCommandContextActions?.waitForIdle ?? (() => this.agent.waitForIdle()), + newSession: this._extensionCommandContextActions?.newSession ?? (async () => ({ cancelled: false })), + fork: this._extensionCommandContextActions?.fork ?? (async () => ({ cancelled: false })), + navigateTree: this._extensionCommandContextActions?.navigateTree ?? (async () => ({ cancelled: false })), + switchSession: this._extensionCommandContextActions?.switchSession ?? (async () => ({ cancelled: false })), + reload: this._extensionCommandContextActions?.reload ?? (() => this.requestReload()), + }); this._extensionErrorUnsubscriber?.(); this._extensionErrorUnsubscriber = this._extensionErrorListener @@ -2269,6 +2389,7 @@ export class AgentSession { } })(); }, + reload: () => this.requestReload(), getSystemPrompt: () => this.systemPrompt, getSystemPromptOptions: () => this._baseSystemPromptOptions, }, @@ -2433,6 +2554,10 @@ export class AgentSession { } async reload(): Promise { + await this._runReloadHandler(); + } + + private async _reloadCore(): Promise { const previousFlagValues = this._extensionRunner.getFlagValues(); await emitSessionShutdownEvent(this._extensionRunner, { type: "session_shutdown", reason: "reload" }); await this.settingsManager.reload(); @@ -2446,6 +2571,7 @@ export class AgentSession { }); const hasBindings = + this._extensionsBound || this._extensionUIContext || this._extensionCommandContextActions || this._extensionShutdownHandler || @@ -2569,6 +2695,13 @@ export class AgentSession { // Bash Execution // ========================================================================= + async emitUserBash(event: UserBashEvent): Promise { + if (!this._extensionRunner.hasHandlers("user_bash")) { + return undefined; + } + return this.withReloadDeferred(() => this._extensionRunner.emitUserBash(event)); + } + /** * Execute a bash command. * Adds result to agent context and session. @@ -2883,6 +3016,7 @@ export class AgentSession { return { editorText, cancelled: false, summaryEntry }; } finally { this._branchSummaryAbortController = undefined; + await this._flushRequestedReload(); } } diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 9cb0e6570..c587c7ab8 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -284,6 +284,7 @@ export class ExtensionRunner { private navigateTreeHandler: NavigateTreeHandler = async () => ({ cancelled: false }); private switchSessionHandler: SwitchSessionHandler = async () => ({ cancelled: false }); private reloadHandler: ReloadHandler = async () => {}; + private commandReloadHandler: ReloadHandler | undefined; private shutdownHandler: ShutdownHandler = () => {}; private shortcutDiagnostics: ResourceDiagnostic[] = []; private commandDiagnostics: ResourceDiagnostic[] = []; @@ -338,6 +339,7 @@ export class ExtensionRunner { this.shutdownHandler = contextActions.shutdown; this.getContextUsageFn = contextActions.getContextUsage; this.compactFn = contextActions.compact; + this.reloadHandler = contextActions.reload ?? (async () => {}); this.getSystemPromptFn = contextActions.getSystemPrompt; this.getSystemPromptOptionsFn = contextActions.getSystemPromptOptions ?? (() => ({ cwd: this.cwd })); @@ -385,7 +387,7 @@ export class ExtensionRunner { this.forkHandler = actions.fork; this.navigateTreeHandler = actions.navigateTree; this.switchSessionHandler = actions.switchSession; - this.reloadHandler = actions.reload; + this.commandReloadHandler = actions.reload; return; } @@ -394,7 +396,7 @@ export class ExtensionRunner { this.forkHandler = async () => ({ cancelled: false }); this.navigateTreeHandler = async () => ({ cancelled: false }); this.switchSessionHandler = async () => ({ cancelled: false }); - this.reloadHandler = async () => {}; + this.commandReloadHandler = undefined; } setUIContext(uiContext?: ExtensionUIContext, mode: ExtensionMode = "print"): void { @@ -678,6 +680,10 @@ export class ExtensionRunner { runner.assertActive(); runner.compactFn(options); }, + reload: () => { + runner.assertActive(); + return runner.reloadHandler(); + }, getSystemPrompt: () => { runner.assertActive(); return runner.getSystemPromptFn(); @@ -719,7 +725,7 @@ export class ExtensionRunner { }; context.reload = () => { this.assertActive(); - return this.reloadHandler(); + return (this.commandReloadHandler ?? this.reloadHandler)(); }; return context; } diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index a869a55d1..871e844cc 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -328,6 +328,11 @@ export interface ExtensionContext { getContextUsage(): ContextUsage | undefined; /** Trigger compaction without awaiting completion. */ compact(options?: CompactOptions): void; + /** + * Request a runtime reload. Resolves after the request is accepted; reload may + * run later when the session reaches a safe boundary. + */ + reload(): Promise; /** Get the current effective system prompt. */ getSystemPrompt(): string; } @@ -1537,6 +1542,7 @@ export interface ExtensionContextActions { shutdown: () => void; getContextUsage: () => ContextUsage | undefined; compact: (options?: CompactOptions) => void; + reload?: () => Promise; getSystemPrompt: () => string; getSystemPromptOptions?: () => BuildSystemPromptOptions; } @@ -1564,7 +1570,7 @@ export interface ExtensionCommandContextActions { sessionPath: string, options?: { withSession?: (ctx: ReplacedSessionContext) => Promise }, ) => Promise<{ cancelled: boolean }>; - reload: () => Promise; + reload?: () => Promise; } /** diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index b3424a15f..ab6c0c152 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -1602,10 +1602,8 @@ export class InteractiveMode { switchSession: async (sessionPath, options) => { return this.handleResumeSession(sessionPath, options); }, - reload: async () => { - await this.handleReloadCommand(); - }, }, + reloadHandler: (reloadCore) => this.performReload(reloadCore), shutdownHandler: () => { this.shutdownRequested = true; if (!this.session.isStreaming) { @@ -1718,6 +1716,7 @@ export class InteractiveMode { } })(); }, + reload: () => this.session.requestReload(), getSystemPrompt: () => this.session.systemPrompt, }); @@ -5058,16 +5057,7 @@ export class InteractiveMode { // Command handlers // ========================================================================= - private async handleReloadCommand(): Promise { - if (this.session.isStreaming) { - this.showWarning("Wait for the current response to finish before reloading."); - return; - } - if (this.session.isCompacting) { - this.showWarning("Wait for compaction to finish before reloading."); - return; - } - + private async performReload(reloadCore: () => Promise): Promise { this.resetExtensionUI(); const reloadBox = new Container(); @@ -5095,7 +5085,8 @@ export class InteractiveMode { }; try { - await this.session.reload(); + let themeWarning: string | undefined; + await reloadCore(); configureHttpDispatcher(this.settingsManager.getHttpIdleTimeoutMs()); this.keybindings.reload(); const activeHeader = this.customHeader ?? this.builtInHeader; @@ -5107,7 +5098,7 @@ export class InteractiveMode { const themeName = this.settingsManager.getTheme(); const themeResult = themeName ? setTheme(themeName, true) : { success: true }; if (!themeResult.success) { - this.showError(`Failed to load theme "${themeName}": ${themeResult.error}\nFell back to dark theme.`); + themeWarning = `Failed to load theme "${themeName}": ${themeResult.error}\nFell back to dark theme.`; } const editorPaddingX = this.settingsManager.getEditorPaddingX(); const autocompleteMaxVisible = this.settingsManager.getAutocompleteMaxVisible(); @@ -5124,6 +5115,9 @@ export class InteractiveMode { this.setupExtensionShortcuts(runner); this.rebuildChatFromMessages(); dismissReloadBox(this.editor as Component); + if (themeWarning) { + this.showWarning(themeWarning); + } this.showLoadedResources({ force: false, showDiagnosticsWhenQuiet: true, @@ -5140,6 +5134,23 @@ export class InteractiveMode { ); } catch (error) { dismissReloadBox(previousEditor as Component); + throw error; + } + } + + private async handleReloadCommand(): Promise { + if (this.session.isStreaming) { + this.showWarning("Wait for the current response to finish before reloading."); + return; + } + if (this.session.isCompacting) { + this.showWarning("Wait for compaction to finish before reloading."); + return; + } + + try { + await this.session.reload(); + } catch (error) { this.showError(`Reload failed: ${error instanceof Error ? error.message : String(error)}`); } } @@ -5633,10 +5644,8 @@ export class InteractiveMode { } private async handleBashCommand(command: string, excludeFromContext = false): Promise { - const extensionRunner = this.session.extensionRunner; - // Emit user_bash event to let extensions intercept - const eventResult = await extensionRunner.emitUserBash({ + const eventResult = await this.session.emitUserBash({ type: "user_bash", command, excludeFromContext, diff --git a/packages/coding-agent/test/agent-session-reload-request.test.ts b/packages/coding-agent/test/agent-session-reload-request.test.ts new file mode 100644 index 000000000..b36e8c5cf --- /dev/null +++ b/packages/coding-agent/test/agent-session-reload-request.test.ts @@ -0,0 +1,387 @@ +import { existsSync, mkdirSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { fauxAssistantMessage, fauxToolCall, registerFauxProvider } from "@earendil-works/pi-ai"; +import { Type } from "typebox"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { AuthStorage } from "../src/core/auth-storage.ts"; +import type { ExtensionCommandContextActions } from "../src/core/extensions/index.ts"; +import { DefaultResourceLoader } from "../src/core/resource-loader.ts"; +import { createAgentSession } from "../src/core/sdk.ts"; +import { SessionManager } from "../src/core/session-manager.ts"; +import { SettingsManager } from "../src/core/settings-manager.ts"; + +describe("AgentSession reload requests", () => { + let tempDir: string; + let agentDir: string; + const cleanups: Array<() => void> = []; + + beforeEach(() => { + tempDir = join(tmpdir(), `pi-reload-request-test-${Date.now()}-${Math.random().toString(36).slice(2)}`); + agentDir = join(tempDir, "agent"); + mkdirSync(agentDir, { recursive: true }); + }); + + afterEach(() => { + while (cleanups.length > 0) { + cleanups.pop()?.(); + } + if (existsSync(tempDir)) { + rmSync(tempDir, { recursive: true, force: true }); + } + }); + + it("defers ctx.reload from an extension tool until the agent turn is idle", async () => { + const events: string[] = []; + const faux = registerFauxProvider(); + cleanups.push(() => faux.unregister()); + faux.setResponses([ + fauxAssistantMessage(fauxToolCall("reload_runtime", {}, { id: "reload-1" }), { stopReason: "toolUse" }), + ]); + + const authStorage = AuthStorage.inMemory(); + authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + const settingsManager = SettingsManager.create(tempDir, agentDir); + const sessionManager = SessionManager.inMemory(); + const resourceLoader = new DefaultResourceLoader({ + cwd: tempDir, + agentDir, + settingsManager, + extensionFactories: [ + (pi) => { + pi.registerTool({ + name: "reload_runtime", + label: "Reload Runtime", + description: "Reload extensions, skills, prompts, and themes", + parameters: Type.Object({}), + async execute(_toolCallId, _params, _signal, _onUpdate, ctx) { + events.push("tool_execute"); + await ctx.reload(); + events.push("reload_requested"); + return { + content: [{ type: "text", text: "reload requested" }], + details: {}, + terminate: true, + }; + }, + }); + pi.on("agent_end", () => { + events.push("agent_end"); + }); + }, + ], + }); + await resourceLoader.reload(); + + const { session } = await createAgentSession({ + cwd: tempDir, + agentDir, + authStorage, + model: faux.getModel(), + settingsManager, + sessionManager, + resourceLoader, + }); + cleanups.push(() => session.dispose()); + + const commandContextActions: ExtensionCommandContextActions = { + waitForIdle: () => session.agent.waitForIdle(), + newSession: async () => ({ cancelled: false }), + fork: async () => ({ cancelled: false }), + navigateTree: async () => ({ cancelled: false }), + switchSession: async () => ({ cancelled: false }), + }; + await session.bindExtensions({ + commandContextActions, + reloadHandler: async (reloadCore) => { + events.push("reload"); + await reloadCore(); + }, + }); + + await session.prompt("reload now"); + + expect(events).toEqual(["tool_execute", "reload_requested", "agent_end", "reload"]); + }); + + it("defers ctx.reload from before_agent_start until the prompt finishes", async () => { + const events: string[] = []; + const faux = registerFauxProvider(); + cleanups.push(() => faux.unregister()); + faux.setResponses([fauxAssistantMessage("ok")]); + + const authStorage = AuthStorage.inMemory(); + authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + const settingsManager = SettingsManager.create(tempDir, agentDir); + const sessionManager = SessionManager.inMemory(); + const resourceLoader = new DefaultResourceLoader({ + cwd: tempDir, + agentDir, + settingsManager, + extensionFactories: [ + (pi) => { + pi.on("before_agent_start", async (_event, ctx) => { + events.push("before_agent_start"); + await ctx.reload(); + events.push("reload_requested"); + return { systemPrompt: "modified" }; + }); + pi.on("agent_end", () => { + events.push("agent_end"); + }); + }, + ], + }); + await resourceLoader.reload(); + + const { session } = await createAgentSession({ + cwd: tempDir, + agentDir, + authStorage, + model: faux.getModel(), + settingsManager, + sessionManager, + resourceLoader, + }); + cleanups.push(() => session.dispose()); + + await session.bindExtensions({ + reloadHandler: async (reloadCore) => { + events.push("reload"); + await reloadCore(); + }, + }); + + await session.prompt("reload before start"); + + expect(events).toEqual(["before_agent_start", "reload_requested", "agent_end", "reload"]); + }); + + it("flushes ctx.reload requested from idle user_bash handlers", async () => { + const events: string[] = []; + const faux = registerFauxProvider(); + cleanups.push(() => faux.unregister()); + + const authStorage = AuthStorage.inMemory(); + authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + const settingsManager = SettingsManager.create(tempDir, agentDir); + const sessionManager = SessionManager.inMemory(); + const resourceLoader = new DefaultResourceLoader({ + cwd: tempDir, + agentDir, + settingsManager, + extensionFactories: [ + (pi) => { + pi.on("user_bash", async (_event, ctx) => { + events.push("user_bash"); + await ctx.reload(); + events.push("reload_requested"); + }); + }, + ], + }); + await resourceLoader.reload(); + + const { session } = await createAgentSession({ + cwd: tempDir, + agentDir, + authStorage, + model: faux.getModel(), + settingsManager, + sessionManager, + resourceLoader, + }); + cleanups.push(() => session.dispose()); + + await session.bindExtensions({ + reloadHandler: async (reloadCore) => { + events.push("reload"); + await reloadCore(); + }, + }); + + await session.emitUserBash({ type: "user_bash", command: "echo hi", excludeFromContext: false, cwd: tempDir }); + + expect(events).toEqual(["user_bash", "reload_requested", "reload"]); + }); + + it("flushes ctx.reload requested from idle selection events", async () => { + const events: string[] = []; + const faux = registerFauxProvider({ models: [{ id: "faux-thinker", reasoning: true }] }); + cleanups.push(() => faux.unregister()); + const model = faux.getModel("faux-thinker"); + if (!model) throw new Error("faux-thinker model not found"); + + const authStorage = AuthStorage.inMemory(); + authStorage.setRuntimeApiKey(model.provider, "faux-key"); + const settingsManager = SettingsManager.create(tempDir, agentDir); + const sessionManager = SessionManager.inMemory(); + const resourceLoader = new DefaultResourceLoader({ + cwd: tempDir, + agentDir, + settingsManager, + extensionFactories: [ + (pi) => { + pi.on("thinking_level_select", async (_event, ctx) => { + events.push("thinking_level_select"); + await ctx.reload(); + events.push("reload_requested"); + }); + }, + ], + }); + await resourceLoader.reload(); + + const { session } = await createAgentSession({ + cwd: tempDir, + agentDir, + authStorage, + model, + settingsManager, + sessionManager, + resourceLoader, + }); + cleanups.push(() => session.dispose()); + + await session.bindExtensions({ + reloadHandler: async (reloadCore) => { + events.push("reload"); + await reloadCore(); + }, + }); + + session.setThinkingLevel("low"); + await new Promise((resolve) => setImmediate(resolve)); + + expect(events).toEqual(["thinking_level_select", "reload_requested", "reload"]); + }); + + it("runs a follow-up reload when direct reload is called during reload", async () => { + const events: string[] = []; + const faux = registerFauxProvider(); + cleanups.push(() => faux.unregister()); + const authStorage = AuthStorage.inMemory(); + authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + const settingsManager = SettingsManager.create(tempDir, agentDir); + const sessionManager = SessionManager.inMemory(); + const resourceLoader = new DefaultResourceLoader({ + cwd: tempDir, + agentDir, + settingsManager, + }); + await resourceLoader.reload(); + + const { session } = await createAgentSession({ + cwd: tempDir, + agentDir, + authStorage, + model: faux.getModel(), + settingsManager, + sessionManager, + resourceLoader, + }); + cleanups.push(() => session.dispose()); + + let reloads = 0; + const commandContextActions: ExtensionCommandContextActions = { + waitForIdle: () => session.agent.waitForIdle(), + newSession: async () => ({ cancelled: false }), + fork: async () => ({ cancelled: false }), + navigateTree: async () => ({ cancelled: false }), + switchSession: async () => ({ cancelled: false }), + }; + await session.bindExtensions({ + commandContextActions, + reloadHandler: async (reloadCore) => { + reloads++; + events.push(`reload_${reloads}`); + if (reloads === 1) { + await session.reload(); + } + await reloadCore(); + }, + }); + + await session.requestReload(); + + expect(events).toEqual(["reload_1", "reload_2"]); + }); + + it("rejects the prompt when a deferred reload fails", async () => { + const events: string[] = []; + const extensionErrors: string[] = []; + const faux = registerFauxProvider(); + cleanups.push(() => faux.unregister()); + faux.setResponses([ + fauxAssistantMessage(fauxToolCall("reload_runtime", {}, { id: "reload-fail-1" }), { stopReason: "toolUse" }), + ]); + + const authStorage = AuthStorage.inMemory(); + authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + const settingsManager = SettingsManager.create(tempDir, agentDir); + const sessionManager = SessionManager.inMemory(); + const resourceLoader = new DefaultResourceLoader({ + cwd: tempDir, + agentDir, + settingsManager, + extensionFactories: [ + (pi) => { + pi.registerTool({ + name: "reload_runtime", + label: "Reload Runtime", + description: "Reload extensions, skills, prompts, and themes", + parameters: Type.Object({}), + async execute(_toolCallId, _params, _signal, _onUpdate, ctx) { + events.push("tool_execute"); + await ctx.reload(); + events.push("reload_requested"); + return { + content: [{ type: "text", text: "reload requested" }], + details: {}, + terminate: true, + }; + }, + }); + pi.on("agent_end", () => { + events.push("agent_end"); + }); + }, + ], + }); + await resourceLoader.reload(); + + const { session } = await createAgentSession({ + cwd: tempDir, + agentDir, + authStorage, + model: faux.getModel(), + settingsManager, + sessionManager, + resourceLoader, + }); + cleanups.push(() => session.dispose()); + + const commandContextActions: ExtensionCommandContextActions = { + waitForIdle: () => session.agent.waitForIdle(), + newSession: async () => ({ cancelled: false }), + fork: async () => ({ cancelled: false }), + navigateTree: async () => ({ cancelled: false }), + switchSession: async () => ({ cancelled: false }), + }; + await session.bindExtensions({ + commandContextActions, + reloadHandler: async () => { + events.push("reload"); + throw new Error("reload failed"); + }, + onError: (error) => { + extensionErrors.push(error.error); + }, + }); + + await expect(session.prompt("reload now")).rejects.toThrow("reload failed"); + + expect(events).toEqual(["tool_execute", "reload_requested", "agent_end", "reload"]); + expect(extensionErrors).toContain("reload failed"); + }); +}); diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index cd6119624..09656ff4b 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -873,6 +873,18 @@ describe("ExtensionRunner", () => { await commandContext.fork("entry-2", { position: "at" }); expect(fork).toHaveBeenLastCalledWith("entry-2", { position: "at" }); }); + + it("keeps reload awaitable on command contexts", async () => { + const runtime = createExtensionRuntime(); + const runner = new ExtensionRunner([], runtime, tempDir, sessionManager, modelRegistry); + const reload = vi.fn(async () => {}); + + runner.bindCore(extensionActions, { ...extensionContextActions, reload }); + + const commandContext = runner.createCommandContext(); + await commandContext.reload(); + expect(reload).toHaveBeenCalledOnce(); + }); }); describe("hasHandlers", () => { diff --git a/packages/coding-agent/test/trigger-compact-extension.test.ts b/packages/coding-agent/test/trigger-compact-extension.test.ts index 80f3d46c8..bc3d805f0 100644 --- a/packages/coding-agent/test/trigger-compact-extension.test.ts +++ b/packages/coding-agent/test/trigger-compact-extension.test.ts @@ -19,6 +19,7 @@ function createContext(tokens: number | null, compact = vi.fn()): ExtensionConte shutdown: vi.fn(), getContextUsage: () => ({ tokens, contextWindow: 200_000, percent: tokens === null ? null : tokens / 2000 }), compact, + reload: vi.fn(), getSystemPrompt: () => "", }; }