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); }); });