From 7e384fbc4fc50b9a9dc0b4465a1a11dcb00699aa Mon Sep 17 00:00:00 2001 From: Seljuke Date: Wed, 12 Aug 2026 00:24:09 +0200 Subject: [PATCH 1/2] fix(extensions): pause tool-call timeout during human dialogs --- .../src/config/settings-schema.ts | 12 ++ .../src/extensibility/extensions/runner.ts | 123 ++++++++++--- .../src/extensibility/extensions/wrapper.ts | 21 ++- .../coding-agent/src/session/agent-session.ts | 19 +- .../test/extensions-runner.test.ts | 170 ++++++++++++++++-- 5 files changed, 295 insertions(+), 50 deletions(-) diff --git a/packages/coding-agent/src/config/settings-schema.ts b/packages/coding-agent/src/config/settings-schema.ts index 95325f2c2..238543076 100644 --- a/packages/coding-agent/src/config/settings-schema.ts +++ b/packages/coding-agent/src/config/settings-schema.ts @@ -5494,6 +5494,18 @@ export const SETTINGS_SCHEMA = { "commit.changelogMaxDiffChars": { type: "number", default: 120000 }, + "extensionHandlers.toolCallTimeoutMs": { + type: "number", + default: 30_000, + ui: { + tab: "tools", + group: "Extensions", + label: "Tool Call Handler Timeout (ms)", + description: + "Active-work timeout for extension tool_call handlers; time awaiting OMP-owned dialogs does not count", + }, + }, + "dev.autoqa": { type: "boolean", default: true, diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index ad7e5417c..842a112fc 100644 --- a/packages/coding-agent/src/extensibility/extensions/runner.ts +++ b/packages/coding-agent/src/extensibility/extensions/runner.ts @@ -116,6 +116,11 @@ function handlerTimeoutForEvent(eventType: string): number { const EXTENSION_HANDLER_TIMEOUT = Symbol("extensionHandlerTimeout"); const EXTENSION_HANDLER_ABORTED = Symbol("extensionHandlerAborted"); +interface HandlerTimeoutBudget { + pause(): void; + resume(): void; +} + function attachHandlerSignal( dialogOptions: ExtensionUIDialogOptions | undefined, handlerSignal: AbortSignal, @@ -126,22 +131,35 @@ function attachHandlerSignal( return { ...dialogOptions, signal: AbortSignal.any([dialogOptions.signal, handlerSignal]) }; } -function createHandlerUIContext(ui: ExtensionUIContext, handlerSignal: AbortSignal): ExtensionUIContext { +function createHandlerUIContext( + ui: ExtensionUIContext, + handlerSignal: AbortSignal, + timeoutBudget?: HandlerTimeoutBudget, +): ExtensionUIContext { const askDialog = ui.askDialog; + const runDialog = async (dialog: () => Promise): Promise => { + timeoutBudget?.pause(); + try { + return await dialog(); + } finally { + timeoutBudget?.resume(); + } + }; const dialogMethods = { select: (title, options, dialogOptions) => - ui.select(title, options, attachHandlerSignal(dialogOptions, handlerSignal)), + runDialog(() => ui.select(title, options, attachHandlerSignal(dialogOptions, handlerSignal))), confirm: (title, message, dialogOptions) => - ui.confirm(title, message, attachHandlerSignal(dialogOptions, handlerSignal)), + runDialog(() => ui.confirm(title, message, attachHandlerSignal(dialogOptions, handlerSignal))), input: (title, placeholder, dialogOptions) => - ui.input(title, placeholder, attachHandlerSignal(dialogOptions, handlerSignal)), + runDialog(() => ui.input(title, placeholder, attachHandlerSignal(dialogOptions, handlerSignal))), askDialog: askDialog ? (questions, dialogOptions) => - askDialog.call(ui, questions, attachHandlerSignal(dialogOptions, handlerSignal)) + runDialog(() => askDialog.call(ui, questions, attachHandlerSignal(dialogOptions, handlerSignal))) : undefined, + custom: (factory, options) => runDialog(() => ui.custom(factory, options)), editor: (title, prefill, dialogOptions, editorOptions) => - ui.editor(title, prefill, attachHandlerSignal(dialogOptions, handlerSignal), editorOptions), - } satisfies Pick; + runDialog(() => ui.editor(title, prefill, attachHandlerSignal(dialogOptions, handlerSignal), editorOptions)), + } satisfies Pick; const delegatedMethods = new Map(); return new Proxy(ui, { @@ -166,10 +184,14 @@ function createHandlerUIContext(ui: ExtensionUIContext, handlerSignal: AbortSign * `pi.setModel()` and then reading `ctx.model` would see a stale model. * Prototype delegation keeps every getter live while overriding `ui`. */ -function createHandlerContext(ctx: ExtensionContext, handlerSignal: AbortSignal): ExtensionContext { +function createHandlerContext( + ctx: ExtensionContext, + handlerSignal: AbortSignal, + timeoutBudget?: HandlerTimeoutBudget, +): ExtensionContext { const scoped: ExtensionContext = Object.create(ctx); Object.defineProperty(scoped, "ui", { - value: createHandlerUIContext(ctx.ui, handlerSignal), + value: createHandlerUIContext(ctx.ui, handlerSignal, timeoutBudget), enumerable: true, configurable: true, }); @@ -189,7 +211,7 @@ function createHandlerContext(ctx: ExtensionContext, handlerSignal: AbortSignal) * can `clearTimeout` on the winning branch. */ async function raceHandlerWithTimeout( - work: (handlerSignal: AbortSignal) => Promise | T, + work: (handlerSignal: AbortSignal, timeoutBudget: HandlerTimeoutBudget) => Promise | T, timeoutMs: number, signal?: AbortSignal, ): Promise { @@ -202,13 +224,52 @@ async function raceHandlerWithTimeout( >(); const onAbort = () => resolveInterrupt(EXTENSION_HANDLER_ABORTED); signal?.addEventListener("abort", onAbort, { once: true }); - const timer = setTimeout(() => { + let timer: Timer | undefined; + let remainingMs = timeoutMs; + let activeSince = performance.now(); + let pauseDepth = 0; + let settled = false; + const clearTimer = () => { + if (timer === undefined) return; + clearTimeout(timer); + timer = undefined; + }; + const expire = () => { + if (settled) return; + settled = true; + clearTimer(); timeoutController.abort(new DOMException(`Handler timed out after ${timeoutMs}ms`, "TimeoutError")); resolveInterrupt(EXTENSION_HANDLER_TIMEOUT); - }, timeoutMs); + }; + const armTimer = () => { + if (settled || pauseDepth > 0) return; + activeSince = performance.now(); + timer = setTimeout(expire, Math.max(0, remainingMs)); + }; + const settle = () => { + if (settled) return; + settled = true; + clearTimer(); + }; + const timeoutBudget: HandlerTimeoutBudget = { + pause: () => { + if (settled) return; + pauseDepth++; + if (pauseDepth !== 1) return; + remainingMs = Math.max(0, remainingMs - (performance.now() - activeSince)); + clearTimer(); + if (remainingMs <= 0) expire(); + }, + resume: () => { + if (settled || pauseDepth === 0) return; + pauseDepth--; + if (pauseDepth === 0) armTimer(); + }, + }; + armTimer(); try { if (signal?.aborted) return EXTENSION_HANDLER_ABORTED; - const workPromise = Promise.resolve(work(handlerSignal)); + const workPromise = Promise.resolve(work(handlerSignal, timeoutBudget)); const result = await Promise.race([workPromise, interruptPromise]); if (result === EXTENSION_HANDLER_TIMEOUT) { await Promise.race([ @@ -221,7 +282,7 @@ async function raceHandlerWithTimeout( } return result; } finally { - clearTimeout(timer); + settle(); signal?.removeEventListener("abort", onAbort); } } @@ -1038,23 +1099,33 @@ export class ExtensionRunner { ext: Extension, timeoutMs: number, onFailure?: (kind: "timeout" | "error", message: string) => TResult, + outerSignal?: AbortSignal, ): Promise { - const signal = + // `session_stop` carries its own signal on the event; `tool_call` receives + // the outer dispatch signal (loop request or wrapper execute) so an abort + // while a handler awaits a human dialog cancels the dialog and settles the + // gate without executing the underlying tool. Compose whichever apply. + const sessionStopSignal = event.type === "session_stop" && "signal" in event && event.signal instanceof AbortSignal ? event.signal : undefined; + const signals = [outerSignal, sessionStopSignal].filter((s): s is AbortSignal => s !== undefined); + const signal = signals.length === 0 ? undefined : signals.length === 1 ? signals[0] : AbortSignal.any(signals); if (signal?.aborted) return undefined; const registrationScope: ToolRegistrationScope = { pending: new Set(), closed: false }; let handlerResult: TResult | typeof EXTENSION_HANDLER_TIMEOUT | typeof EXTENSION_HANDLER_ABORTED | undefined; let handlerFailure: { error: unknown } | undefined; try { handlerResult = await raceHandlerWithTimeout( - async handlerSignal => { + async (handlerSignal, budget) => { registrationScope.signal = handlerSignal; let result: TResult | undefined; try { result = await this.#toolRegistrationScope.run(registrationScope, () => - handler(event, createHandlerContext(ctx, handlerSignal)), + handler( + event, + createHandlerContext(ctx, handlerSignal, event.type === "tool_call" ? budget : undefined), + ), ); } catch (error) { handlerFailure = { error }; @@ -1215,8 +1286,8 @@ export class ExtensionRunner { /** * Emit a `tool_call` event to every subscribed extension before the tool executes. * - * Each handler is bounded by `extensionHandlerTimeoutMs` (default 30s). This - * matches the timeout policy already applied to `emitToolResult` and every + * Each handler is bounded by `extensionHandlers.toolCallTimeoutMs` (default + * 30s). This matches the timeout policy already applied to `emitToolResult` and every * other handler routed through `#runHandlerWithTimeout`; without it a single * hung extension (unresolved `await`, network call with no timeout) would * park `ExtensionToolWrapper.execute` indefinitely and freeze tool @@ -1227,9 +1298,9 @@ export class ExtensionRunner { * pre-execution gate — an unresponsive extension MUST NOT be treated as * silent consent to run the tool. */ - async emitToolCall(event: ToolCallEvent): Promise { + async emitToolCall(event: ToolCallEvent, signal?: AbortSignal): Promise { const ctx = this.createContext(); - const timeoutMs = extensionHandlerTimeoutMs; + const timeoutMs = this.settings?.get("extensionHandlers.toolCallTimeoutMs") ?? extensionHandlerTimeoutMs; let result: ToolCallEventResult | undefined; for (const ext of this.extensions) { @@ -1250,6 +1321,7 @@ export class ExtensionRunner { ? `Extension ${ext.path} timed out after ${timeoutMs}ms` : `Extension ${ext.path} failed: ${message}`, }), + signal, ); if (handlerResult) { @@ -1258,9 +1330,18 @@ export class ExtensionRunner { return result; } } + // Fail closed when the outer dispatch aborted while a handler was + // pending: an aborted gate MUST NOT become silent consent to run the + // underlying tool. Symmetric with the timeout policy above. + if (signal?.aborted) { + return { block: true, reason: `Tool execution was cancelled while an extension handler was pending` }; + } } } + if (signal?.aborted) { + return { block: true, reason: `Tool execution was cancelled while an extension handler was pending` }; + } return result; } diff --git a/packages/coding-agent/src/extensibility/extensions/wrapper.ts b/packages/coding-agent/src/extensibility/extensions/wrapper.ts index 335a7abcd..3be1ab950 100644 --- a/packages/coding-agent/src/extensibility/extensions/wrapper.ts +++ b/packages/coding-agent/src/extensibility/extensions/wrapper.ts @@ -205,15 +205,18 @@ export class ExtensionToolWrapper this.#beforeToolCall(ctx); + this.agent.beforeToolCall = (ctx, signal) => this.#beforeToolCall(ctx, signal); this.agent.providerSessionState = this.#providerSessionState; this.#syncAgentSessionId(); this.#todo.syncFromBranch(); @@ -3276,7 +3276,7 @@ export class AgentSession { * emit a second event (nested xd:// device dispatches and direct non-loop * execution still emit there). */ - async #beforeToolCall(ctx: BeforeToolCallContext): Promise { + async #beforeToolCall(ctx: BeforeToolCallContext, signal?: AbortSignal): Promise { const runner = this.#extensionRunner; if (!runner?.hasHandlers("tool_call")) return undefined; const metadata = ctx.toolCall.providerMetadata; @@ -3294,12 +3294,15 @@ export class AgentSession { ? { actions: computer.actions, pendingSafetyChecks: computer.pendingSafetyChecks } : ctx.args; runner.markToolCallEmitted(ctx.toolCall.id, ctx.tool.name); - const callResult = await runner.emitToolCall({ - type: "tool_call", - toolName: ctx.tool.name, - toolCallId: ctx.toolCall.id, - input: normalizeToolEventInput(ctx.tool.name, resolveToolEventInput(ctx.tool, eventArgs)), - }); + const callResult = await runner.emitToolCall( + { + type: "tool_call", + toolName: ctx.tool.name, + toolCallId: ctx.toolCall.id, + input: normalizeToolEventInput(ctx.tool.name, resolveToolEventInput(ctx.tool, eventArgs)), + }, + signal, + ); if (callResult?.block) { return { block: true, reason: callResult.reason || "Tool execution was blocked by an extension" }; } diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 059cf39c3..0fda1887e 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -10,6 +10,7 @@ import type { AgentMessage, AgentTool } from "@oh-my-pi/pi-agent-core"; import type { ImageContent, TextContent } from "@oh-my-pi/pi-ai"; import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { discoverAndLoadExtensions, ExtensionRuntime } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/loader"; import { EXTENSION_HANDLER_TIMEOUT_MS, @@ -1301,7 +1302,7 @@ describe("ExtensionRunner", () => { }); }); - it("times out tool_call handlers with fail-closed policy so a hung extension cannot indefinitely block tool execution (#3948)", async () => { + it("uses the configured tool_call timeout and fails closed so a hung extension cannot block execution (#3948)", async () => { const hangExtensionPath = path.join(tempDir.path(), "hang-tool-call.ts"); fs.writeFileSync( hangExtensionPath, @@ -1321,14 +1322,14 @@ describe("ExtensionRunner", () => { tempDir.path(), sessionManager, modelRegistry, + undefined, + Settings.isolated({ "extensionHandlers.toolCallTimeoutMs": 10 }), ); const warnSpy = vi.spyOn(logger, "warn").mockImplementation(() => {}); const errors: Array<{ extensionPath: string; event: string; error: string }> = []; runner.onError(err => { errors.push(err); }); - testSetExtensionHandlerTimeoutMs(10); - const executeCalls: unknown[] = []; const tool: AgentTool = { name: "sleepy", @@ -1481,7 +1482,7 @@ describe("ExtensionRunner", () => { expect(errors).toEqual([]); }); - it("aborts a tool_call handler's confirmation before returning its timeout block", async () => { + it("pauses a tool_call handler timeout during standard and custom dialogs, then resumes its budget", async () => { const extensionPath = path.join(tempDir.path(), "confirm-tool-call.ts"); const markerPath = path.join(tempDir.path(), "confirm-settled.txt"); fs.writeFileSync( @@ -1492,8 +1493,11 @@ describe("ExtensionRunner", () => { export default function(pi) { pi.on("tool_call", async (_event, ctx) => { ctx.ui.notify("Waiting for confirmation"); + await new Promise(resolve => setTimeout(resolve, 8)); await ctx.ui.confirm("High-risk command", "Allow this command?"); + await ctx.ui.custom(() => ({})); fs.writeFileSync(${JSON.stringify(markerPath)}, "settled"); + await Promise.withResolvers().promise; }); } `, @@ -1515,8 +1519,16 @@ describe("ExtensionRunner", () => { dialogSignal?.addEventListener("abort", () => dialog.resolve(false), { once: true }); return await dialog.promise; }; + const customDialog = Promise.withResolvers(); + let customPending = false; + const custom: ExtensionUIContext["custom"] = async () => { + customPending = true; + await customDialog.promise; + return undefined as T; + }; const uiPrototype = Object.create(runner.getUIContext(), { confirm: { value: confirm }, + custom: { value: custom }, notify: { value: notify }, }); const uiContext: ExtensionUIContext = Object.create(uiPrototype); @@ -1549,25 +1561,159 @@ describe("ExtensionRunner", () => { undefined, uiContext, ); - testSetExtensionHandlerTimeoutMs(10); + vi.useFakeTimers(); + let now = 0; + const performanceNow = vi.spyOn(performance, "now").mockImplementation(() => now); + try { + testSetExtensionHandlerTimeoutMs(25); + + const tool: AgentTool = { + name: "guarded", + label: "Guarded", + description: "must not execute after the extension gate times out", + parameters: Type.Object({}), + strict: true, + execute: async () => ({ content: [{ type: "text", text: "ran" }] }), + }; + const wrapped = new ExtensionToolWrapper(tool, runner); + const flush = async () => { + for (let attempts = 0; attempts < 10; attempts++) await Promise.resolve(); + }; + + const execution = wrapped.execute("tool-call-id", {}); + await flush(); + expect(notify).toHaveBeenCalledWith("Waiting for confirmation"); + expect(dialogSignal).toBeUndefined(); + + now = 8; + vi.advanceTimersByTime(8); + await flush(); + expect(dialogSignal).toBeDefined(); + + now = 108; + vi.advanceTimersByTime(100); + await flush(); + expect(dialogSignal?.aborted).toBe(false); + + dialog.resolve(true); + await flush(); + expect(customPending).toBe(true); + expect(fs.existsSync(markerPath)).toBe(false); + + now = 208; + vi.advanceTimersByTime(100); + await flush(); + expect(dialogSignal?.aborted).toBe(false); + expect(fs.existsSync(markerPath)).toBe(false); + + customDialog.resolve(); + await flush(); + expect(fs.readFileSync(markerPath, "utf8")).toBe("settled"); + + now = 225; + vi.advanceTimersByTime(17); + await flush(); + vi.advanceTimersByTime(0); + await flush(); + await expect(execution).rejects.toThrow(`Extension ${extensionPath} timed out after 25ms`); + } finally { + performanceNow.mockRestore(); + vi.useRealTimers(); + } + }); + + it("cancels a pending confirmation and blocks tool execution when the outer dispatch aborts (#4223)", async () => { + const extensionPath = path.join(tempDir.path(), "confirm-abort-tool-call.ts"); + const recordPath = path.join(tempDir.path(), "confirm-abort-executed.jsonl"); + fs.writeFileSync( + extensionPath, + ` + export default function(pi) { + pi.on("tool_call", async (_event, ctx) => { + await ctx.ui.confirm("High-risk command", "Allow this command?"); + }); + } + `, + ); + + const result = await loadTestExtensions([extensionPath]); + const runner = new ExtensionRunner( + result.extensions, + result.runtime, + tempDir.path(), + sessionManager, + modelRegistry, + ); + let dialogSignal: AbortSignal | undefined; + const dialog = Promise.withResolvers(); + const confirm: ExtensionUIContext["confirm"] = async (_title, _message, dialogOptions) => { + dialogSignal = dialogOptions?.signal; + dialogSignal?.addEventListener("abort", () => dialog.resolve(false), { once: true }); + return await dialog.promise; + }; + const uiPrototype = Object.create(runner.getUIContext(), { + confirm: { value: confirm }, + }); + const uiContext: ExtensionUIContext = Object.create(uiPrototype); + runner.initialize( + { + sendMessage: () => {}, + sendUserMessage: () => {}, + appendEntry: () => {}, + setLabel: () => {}, + getActiveTools: () => [], + getAllTools: () => [], + setActiveTools: async () => {}, + getCommands: () => [], + setModel: async () => false, + getThinkingLevel: () => undefined, + setThinkingLevel: () => {}, + getSessionName: () => undefined, + setSessionName: async () => {}, + }, + { + getModel: () => undefined, + isIdle: () => true, + abort: () => {}, + hasPendingMessages: () => false, + shutdown: () => {}, + getContextUsage: () => undefined, + compact: async () => {}, + getSystemPrompt: () => [], + }, + undefined, + uiContext, + ); const tool: AgentTool = { name: "guarded", label: "Guarded", - description: "must not execute after the extension gate times out", + description: "must not execute after the dispatch aborts", parameters: Type.Object({}), strict: true, - execute: async () => ({ content: [{ type: "text", text: "ran" }] }), + execute: async () => { + fs.appendFileSync(recordPath, "ran\n"); + return { content: [{ type: "text", text: "ran" }] }; + }, }; const wrapped = new ExtensionToolWrapper(tool, runner); + const flush = async () => { + for (let attempts = 0; attempts < 10; attempts++) await Promise.resolve(); + }; - await expect(wrapped.execute("tool-call-id", {})).rejects.toThrow( - `Extension ${extensionPath} timed out after 10ms`, - ); - expect(notify).toHaveBeenCalledWith("Waiting for confirmation"); + const controller = new AbortController(); + const execution = wrapped.execute("tool-call-id", {} as never, controller.signal); + await flush(); + + expect(dialogSignal).toBeDefined(); + expect(dialogSignal?.aborted).toBe(false); + + controller.abort(); + await flush(); expect(dialogSignal?.aborted).toBe(true); - expect(fs.readFileSync(markerPath, "utf8")).toBe("settled"); + await expect(execution).rejects.toThrow(); + expect(fs.existsSync(recordPath)).toBe(false); }); }); From fd28acf5a1fca96c9866618363f49cfc39565db6 Mon Sep 17 00:00:00 2001 From: Seljuke Date: Thu, 13 Aug 2026 20:39:12 +0200 Subject: [PATCH 2/2] fix(extensions): harden dialog timeouts --- packages/coding-agent/CHANGELOG.md | 4 + .../src/config/settings-schema.ts | 2 +- .../src/extensibility/extensions/runner.ts | 22 +- .../src/extensibility/extensions/types.ts | 2 + .../controllers/extension-ui-controller.ts | 80 ++++--- .../test/extensions-runner.test.ts | 212 ++++++++++-------- .../extension-ui-controller.test.ts | 32 ++- 7 files changed, 221 insertions(+), 133 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 14334fddd..de9c516e1 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Made extension tool-call timeouts configurable and paused them during user dialogs. + ## [17.3.0] - 2026-08-13 ### Breaking Changes diff --git a/packages/coding-agent/src/config/settings-schema.ts b/packages/coding-agent/src/config/settings-schema.ts index 238543076..5fc05954b 100644 --- a/packages/coding-agent/src/config/settings-schema.ts +++ b/packages/coding-agent/src/config/settings-schema.ts @@ -5502,7 +5502,7 @@ export const SETTINGS_SCHEMA = { group: "Extensions", label: "Tool Call Handler Timeout (ms)", description: - "Active-work timeout for extension tool_call handlers; time awaiting OMP-owned dialogs does not count", + "Positive finite active-work timeout for extension tool_call handlers; invalid values use 30000ms, and time awaiting OMP-owned dialogs does not count", }, }, diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index 842a112fc..dbc6d40f9 100644 --- a/packages/coding-agent/src/extensibility/extensions/runner.ts +++ b/packages/coding-agent/src/extensibility/extensions/runner.ts @@ -91,6 +91,10 @@ export function testSetExtensionHandlerTimeoutMs(timeoutMs: number): void { extensionHandlerTimeoutMs = timeoutMs; } +function normalizeHandlerTimeout(timeoutMs: number): number { + return Number.isFinite(timeoutMs) && timeoutMs > 0 ? timeoutMs : EXTENSION_HANDLER_TIMEOUT_MS; +} + /** * Dedicated cap for `session_shutdown` handlers. The generic 30s budget is * appropriate for events extensions can observe (e.g. `session_start`, @@ -156,7 +160,13 @@ function createHandlerUIContext( ? (questions, dialogOptions) => runDialog(() => askDialog.call(ui, questions, attachHandlerSignal(dialogOptions, handlerSignal))) : undefined, - custom: (factory, options) => runDialog(() => ui.custom(factory, options)), + custom: (factory, options) => + runDialog(() => + ui.custom(factory, { + ...options, + signal: options?.signal ? AbortSignal.any([options.signal, handlerSignal]) : handlerSignal, + }), + ), editor: (title, prefill, dialogOptions, editorOptions) => runDialog(() => ui.editor(title, prefill, attachHandlerSignal(dialogOptions, handlerSignal), editorOptions)), } satisfies Pick; @@ -1300,7 +1310,9 @@ export class ExtensionRunner { */ async emitToolCall(event: ToolCallEvent, signal?: AbortSignal): Promise { const ctx = this.createContext(); - const timeoutMs = this.settings?.get("extensionHandlers.toolCallTimeoutMs") ?? extensionHandlerTimeoutMs; + const timeoutMs = normalizeHandlerTimeout( + this.settings?.get("extensionHandlers.toolCallTimeoutMs") ?? extensionHandlerTimeoutMs, + ); let result: ToolCallEventResult | undefined; for (const ext of this.extensions) { @@ -1330,12 +1342,6 @@ export class ExtensionRunner { return result; } } - // Fail closed when the outer dispatch aborted while a handler was - // pending: an aborted gate MUST NOT become silent consent to run the - // underlying tool. Symmetric with the timeout policy above. - if (signal?.aborted) { - return { block: true, reason: `Tool execution was cancelled while an extension handler was pending` }; - } } } diff --git a/packages/coding-agent/src/extensibility/extensions/types.ts b/packages/coding-agent/src/extensibility/extensions/types.ts index 766519f7a..d2aaa6bca 100644 --- a/packages/coding-agent/src/extensibility/extensions/types.ts +++ b/packages/coding-agent/src/extensibility/extensions/types.ts @@ -234,6 +234,8 @@ export interface ExtensionCustomOptions { overlayOptions?: OverlayOptions | (() => OverlayOptions); /** Invoked with the overlay handle once the overlay is created (overlay mode only). */ onHandle?: (handle: OverlayHandle) => void; + /** Abort the custom UI and reject its promise. */ + signal?: AbortSignal; } /** Wrap the current autocomplete provider with additional behavior (pi-compatible). */ diff --git a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts index 1123ba668..af8872451 100644 --- a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts +++ b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts @@ -1053,14 +1053,12 @@ export class ExtensionUiController { const savedText = this.ctx.editor.getText(); const keybindings = KeybindingsManager.inMemory(); - const { promise, resolve } = Promise.withResolvers(); + const { promise, resolve, reject } = Promise.withResolvers(); let component: (Component & { dispose?(): void }) | undefined; let overlayHandle: OverlayHandle | undefined; let closed = false; - const close = (result: T) => { - if (closed) return; - closed = true; + const cleanup = () => { component?.dispose?.(); overlayHandle?.hide(); overlayHandle = undefined; @@ -1071,35 +1069,55 @@ export class ExtensionUiController { } this.ctx.ui.setFocus(this.ctx.editor); this.ctx.ui.requestRender(); - resolve(result); }; + const finish = (settle: () => void) => { + if (closed) return; + closed = true; + options?.signal?.removeEventListener("abort", onAbort); + try { + cleanup(); + } finally { + settle(); + } + }; + const fail = (error: unknown) => finish(() => reject(error)); + const onAbort = () => fail(options?.signal?.reason ?? new DOMException("Dialog aborted", "AbortError")); + const close = (result: T) => finish(() => resolve(result)); - Promise.try(() => factory(this.ctx.ui, theme, keybindings, close)).then(c => { - if (closed) { - c.dispose?.(); - return; - } - component = c; - if (options?.overlay) { - const overlayOptions = - typeof options.overlayOptions === "function" ? options.overlayOptions() : options.overlayOptions; - overlayHandle = this.ctx.ui.showOverlay( - component, - overlayOptions ?? { - anchor: "bottom-center", - width: "100%", - maxHeight: "100%", - margin: 0, - }, - ); - options.onHandle?.(overlayHandle); - return; - } - this.ctx.editorContainer.clear(); - this.ctx.editorContainer.addChild(component); - this.ctx.ui.setFocus(component); - this.ctx.ui.requestRender(); - }); + if (options?.signal?.aborted) { + fail(options.signal.reason ?? new DOMException("Dialog aborted", "AbortError")); + return promise; + } + options?.signal?.addEventListener("abort", onAbort, { once: true }); + + Promise.try(() => factory(this.ctx.ui, theme, keybindings, close)) + .then(c => { + if (closed) { + c.dispose?.(); + return; + } + component = c; + if (options?.overlay) { + const overlayOptions = + typeof options.overlayOptions === "function" ? options.overlayOptions() : options.overlayOptions; + overlayHandle = this.ctx.ui.showOverlay( + component, + overlayOptions ?? { + anchor: "bottom-center", + width: "100%", + maxHeight: "100%", + margin: 0, + }, + ); + options.onHandle?.(overlayHandle); + return; + } + this.ctx.editorContainer.clear(); + this.ctx.editorContainer.addChild(component); + this.ctx.ui.setFocus(component); + this.ctx.ui.requestRender(); + }) + .catch(fail); return promise; } diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 0fda1887e..a4f0cb91e 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -1189,6 +1189,38 @@ describe("ExtensionRunner", () => { }); describe("handler timeouts", () => { + const initializeRunner = (runner: ExtensionRunner, uiContext: ExtensionUIContext): void => { + runner.initialize( + { + sendMessage: () => {}, + sendUserMessage: () => {}, + appendEntry: () => {}, + setLabel: () => {}, + getActiveTools: () => [], + getAllTools: () => [], + setActiveTools: async () => {}, + getCommands: () => [], + setModel: async () => false, + getThinkingLevel: () => undefined, + setThinkingLevel: () => {}, + getSessionName: () => undefined, + setSessionName: async () => {}, + }, + { + getModel: () => undefined, + isIdle: () => true, + abort: () => {}, + hasPendingMessages: () => false, + shutdown: () => {}, + getContextUsage: () => undefined, + compact: async () => {}, + getSystemPrompt: () => [], + }, + undefined, + uiContext, + ); + }; + it("times out session_start handlers, emits an error, and continues to sibling extensions", async () => { const hangExtensionPath = path.join(tempDir.path(), "hang-session-start.ts"); const fastExtensionPath = path.join(tempDir.path(), "fast-session-start.ts"); @@ -1370,6 +1402,62 @@ describe("ExtensionRunner", () => { warnSpy.mockRestore(); }); + it("falls back to the default tool_call timeout for invalid configured values", async () => { + const extensionPath = path.join(tempDir.path(), "invalid-timeout-tool-call.ts"); + fs.writeFileSync( + extensionPath, + ` + export default function(pi) { + pi.on("tool_call", async () => { + await Promise.withResolvers().promise; + }); + } + `, + ); + const loaded = await loadTestExtensions([extensionPath]); + + vi.useFakeTimers(); + try { + for (const configuredTimeout of [0, -1, Number.NaN, Number.POSITIVE_INFINITY]) { + const runner = new ExtensionRunner( + loaded.extensions, + loaded.runtime, + tempDir.path(), + sessionManager, + modelRegistry, + undefined, + Settings.isolated({ "extensionHandlers.toolCallTimeoutMs": configuredTimeout }), + ); + let settled = false; + const decision = runner + .emitToolCall({ + type: "tool_call", + toolName: "guarded", + toolCallId: "invalid-timeout-call", + input: {}, + }) + .then(result => { + settled = true; + return result; + }); + + vi.advanceTimersByTime(EXTENSION_HANDLER_TIMEOUT_MS - 1); + expect(settled).toBe(false); + + vi.advanceTimersByTime(1); + await Promise.resolve(); + await Promise.resolve(); + vi.advanceTimersByTime(0); + expect(await decision).toEqual({ + block: true, + reason: `Extension ${extensionPath} timed out after ${EXTENSION_HANDLER_TIMEOUT_MS}ms`, + }); + } + } finally { + vi.useRealTimers(); + } + }); + it("fails closed when a tool_call handler registration cannot activate", async () => { const extensionPath = path.join(tempDir.path(), "tool-call-registration.ts"); fs.writeFileSync( @@ -1484,19 +1572,16 @@ describe("ExtensionRunner", () => { it("pauses a tool_call handler timeout during standard and custom dialogs, then resumes its budget", async () => { const extensionPath = path.join(tempDir.path(), "confirm-tool-call.ts"); - const markerPath = path.join(tempDir.path(), "confirm-settled.txt"); fs.writeFileSync( extensionPath, ` - import * as fs from "node:fs"; - export default function(pi) { pi.on("tool_call", async (_event, ctx) => { ctx.ui.notify("Waiting for confirmation"); await new Promise(resolve => setTimeout(resolve, 8)); await ctx.ui.confirm("High-risk command", "Allow this command?"); await ctx.ui.custom(() => ({})); - fs.writeFileSync(${JSON.stringify(markerPath)}, "settled"); + ctx.ui.notify("Custom settled"); await Promise.withResolvers().promise; }); } @@ -1512,17 +1597,26 @@ describe("ExtensionRunner", () => { modelRegistry, ); const dialog = Promise.withResolvers(); + const handlerStarted = Promise.withResolvers(); + const confirmationStarted = Promise.withResolvers(); + const customStarted = Promise.withResolvers(); + const customCompleted = Promise.withResolvers(); let dialogSignal: AbortSignal | undefined; - const notify = vi.fn(); + const notify: ExtensionUIContext["notify"] = message => { + if (message === "Waiting for confirmation") handlerStarted.resolve(); + if (message === "Custom settled") customCompleted.resolve(); + }; const confirm: ExtensionUIContext["confirm"] = async (_title, _message, dialogOptions) => { dialogSignal = dialogOptions?.signal; + confirmationStarted.resolve(); dialogSignal?.addEventListener("abort", () => dialog.resolve(false), { once: true }); return await dialog.promise; }; const customDialog = Promise.withResolvers(); - let customPending = false; - const custom: ExtensionUIContext["custom"] = async () => { - customPending = true; + let customSignal: AbortSignal | undefined; + const custom: ExtensionUIContext["custom"] = async (...args: Parameters) => { + customSignal = args[1]?.signal; + customStarted.resolve(); await customDialog.promise; return undefined as T; }; @@ -1532,35 +1626,7 @@ describe("ExtensionRunner", () => { notify: { value: notify }, }); const uiContext: ExtensionUIContext = Object.create(uiPrototype); - runner.initialize( - { - sendMessage: () => {}, - sendUserMessage: () => {}, - appendEntry: () => {}, - setLabel: () => {}, - getActiveTools: () => [], - getAllTools: () => [], - setActiveTools: async () => {}, - getCommands: () => [], - setModel: async () => false, - getThinkingLevel: () => undefined, - setThinkingLevel: () => {}, - getSessionName: () => undefined, - setSessionName: async () => {}, - }, - { - getModel: () => undefined, - isIdle: () => true, - abort: () => {}, - hasPendingMessages: () => false, - shutdown: () => {}, - getContextUsage: () => undefined, - compact: async () => {}, - getSystemPrompt: () => [], - }, - undefined, - uiContext, - ); + initializeRunner(runner, uiContext); vi.useFakeTimers(); let now = 0; const performanceNow = vi.spyOn(performance, "now").mockImplementation(() => now); @@ -1576,45 +1642,37 @@ describe("ExtensionRunner", () => { execute: async () => ({ content: [{ type: "text", text: "ran" }] }), }; const wrapped = new ExtensionToolWrapper(tool, runner); - const flush = async () => { - for (let attempts = 0; attempts < 10; attempts++) await Promise.resolve(); - }; const execution = wrapped.execute("tool-call-id", {}); - await flush(); - expect(notify).toHaveBeenCalledWith("Waiting for confirmation"); + await handlerStarted.promise; expect(dialogSignal).toBeUndefined(); now = 8; vi.advanceTimersByTime(8); - await flush(); + await confirmationStarted.promise; expect(dialogSignal).toBeDefined(); now = 108; vi.advanceTimersByTime(100); - await flush(); expect(dialogSignal?.aborted).toBe(false); dialog.resolve(true); - await flush(); - expect(customPending).toBe(true); - expect(fs.existsSync(markerPath)).toBe(false); + await customStarted.promise; + expect(customSignal).toBeDefined(); + expect(customSignal?.aborted).toBe(false); now = 208; vi.advanceTimersByTime(100); - await flush(); - expect(dialogSignal?.aborted).toBe(false); - expect(fs.existsSync(markerPath)).toBe(false); + expect(customSignal?.aborted).toBe(false); customDialog.resolve(); - await flush(); - expect(fs.readFileSync(markerPath, "utf8")).toBe("settled"); + await customCompleted.promise; now = 225; vi.advanceTimersByTime(17); - await flush(); + await Promise.resolve(); + await Promise.resolve(); vi.advanceTimersByTime(0); - await flush(); await expect(execution).rejects.toThrow(`Extension ${extensionPath} timed out after 25ms`); } finally { performanceNow.mockRestore(); @@ -1624,7 +1682,6 @@ describe("ExtensionRunner", () => { it("cancels a pending confirmation and blocks tool execution when the outer dispatch aborts (#4223)", async () => { const extensionPath = path.join(tempDir.path(), "confirm-abort-tool-call.ts"); - const recordPath = path.join(tempDir.path(), "confirm-abort-executed.jsonl"); fs.writeFileSync( extensionPath, ` @@ -1646,8 +1703,10 @@ describe("ExtensionRunner", () => { ); let dialogSignal: AbortSignal | undefined; const dialog = Promise.withResolvers(); + const confirmationStarted = Promise.withResolvers(); const confirm: ExtensionUIContext["confirm"] = async (_title, _message, dialogOptions) => { dialogSignal = dialogOptions?.signal; + confirmationStarted.resolve(); dialogSignal?.addEventListener("abort", () => dialog.resolve(false), { once: true }); return await dialog.promise; }; @@ -1655,35 +1714,8 @@ describe("ExtensionRunner", () => { confirm: { value: confirm }, }); const uiContext: ExtensionUIContext = Object.create(uiPrototype); - runner.initialize( - { - sendMessage: () => {}, - sendUserMessage: () => {}, - appendEntry: () => {}, - setLabel: () => {}, - getActiveTools: () => [], - getAllTools: () => [], - setActiveTools: async () => {}, - getCommands: () => [], - setModel: async () => false, - getThinkingLevel: () => undefined, - setThinkingLevel: () => {}, - getSessionName: () => undefined, - setSessionName: async () => {}, - }, - { - getModel: () => undefined, - isIdle: () => true, - abort: () => {}, - hasPendingMessages: () => false, - shutdown: () => {}, - getContextUsage: () => undefined, - compact: async () => {}, - getSystemPrompt: () => [], - }, - undefined, - uiContext, - ); + initializeRunner(runner, uiContext); + let executed = false; const tool: AgentTool = { name: "guarded", @@ -1692,28 +1724,24 @@ describe("ExtensionRunner", () => { parameters: Type.Object({}), strict: true, execute: async () => { - fs.appendFileSync(recordPath, "ran\n"); + executed = true; return { content: [{ type: "text", text: "ran" }] }; }, }; const wrapped = new ExtensionToolWrapper(tool, runner); - const flush = async () => { - for (let attempts = 0; attempts < 10; attempts++) await Promise.resolve(); - }; const controller = new AbortController(); const execution = wrapped.execute("tool-call-id", {} as never, controller.signal); - await flush(); + await confirmationStarted.promise; expect(dialogSignal).toBeDefined(); expect(dialogSignal?.aborted).toBe(false); controller.abort(); - await flush(); + await expect(execution).rejects.toThrow(); expect(dialogSignal?.aborted).toBe(true); - await expect(execution).rejects.toThrow(); - expect(fs.existsSync(recordPath)).toBe(false); + expect(executed).toBe(false); }); }); diff --git a/packages/coding-agent/test/modes/controllers/extension-ui-controller.test.ts b/packages/coding-agent/test/modes/controllers/extension-ui-controller.test.ts index fafbaba79..116d956cc 100644 --- a/packages/coding-agent/test/modes/controllers/extension-ui-controller.test.ts +++ b/packages/coding-agent/test/modes/controllers/extension-ui-controller.test.ts @@ -1,4 +1,4 @@ -import { afterEach, beforeAll, describe, expect, it, vi } from "bun:test"; +import { afterEach, beforeAll, describe, expect, it, type Mock, vi } from "bun:test"; import { Container, type OverlayOptions, setKeybindings } from "@oh-my-pi/pi-tui"; import { KeybindingsManager } from "../../../src/config/keybindings"; import type { ExtensionAskDialogQuestion, ExtensionUIContext } from "../../../src/extensibility/extensions"; @@ -316,4 +316,34 @@ describe("ExtensionUiController custom overlay", () => { margin: 0, }); }); + + it("rejects and restores the editor when a custom factory fails", async () => { + const harness = makeHarness(); + const ui = await harness.init(); + const failure = new Error("custom factory failed"); + + await expect(ui.custom(() => Promise.reject(failure))).rejects.toBe(failure); + + expect(harness.editorContainer.children).toEqual([harness.editor]); + expect(harness.setFocus).toHaveBeenLastCalledWith(harness.editor); + }); + + it("aborts a pending custom factory and disposes its late component", async () => { + const harness = makeHarness(); + const ui = await harness.init(); + const controller = new AbortController(); + const factory = Promise.withResolvers(); + const component = new Container() as Container & { dispose: Mock<() => void> }; + component.dispose = vi.fn(); + + const pending = ui.custom(() => factory.promise, { signal: controller.signal }); + controller.abort(); + + await expect(pending).rejects.toBe(controller.signal.reason); + factory.resolve(component); + await flushMicrotasks(); + + expect(component.dispose).toHaveBeenCalledTimes(1); + expect(harness.editorContainer.children).toEqual([harness.editor]); + }); });