diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index a59a4e7c1..9b985bbb6 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -9,6 +9,7 @@ ### Fixed +- Fixed interactive extension confirmations ignoring `dialogOptions`, and cancelled handler-owned dialogs when the extension watchdog times out so stale approval UI cannot outlive a blocked tool call ([#6805](https://github.com/can1357/oh-my-pi/issues/6805)). - Fixed Python cell errors (`$` commands and the eval tool) leaking runner-internal traceback frames. Cell syntax errors now render as the bare caret display with a `` filename instead of a `_handle_request_async`/`ast.parse` stack dump, and runtime tracebacks start at user code, matching the Ruby runner's user-frame filtering. - Dropped unavailable forced tool choices through the queue rejection lifecycle and discarded their remaining sequence yields so a skipped force cannot disable tools on the next request ([#6543](https://github.com/can1357/oh-my-pi/pull/6543) by [@paralin](https://github.com/paralin)). diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index 477ed18be..d28ca07c4 100644 --- a/packages/coding-agent/src/extensibility/extensions/runner.ts +++ b/packages/coding-agent/src/extensibility/extensions/runner.ts @@ -37,6 +37,7 @@ import type { ExtensionRuntime, ExtensionShortcut, ExtensionUIContext, + ExtensionUIDialogOptions, InputEvent, InputEventResult, MessageRenderer, @@ -105,6 +106,35 @@ function handlerTimeoutForEvent(eventType: string): number { const EXTENSION_HANDLER_TIMEOUT = Symbol("extensionHandlerTimeout"); const EXTENSION_HANDLER_ABORTED = Symbol("extensionHandlerAborted"); +function attachHandlerSignal( + dialogOptions: ExtensionUIDialogOptions | undefined, + handlerSignal: AbortSignal, +): ExtensionUIDialogOptions { + if (!dialogOptions) return { signal: handlerSignal }; + if (!dialogOptions.signal) return { ...dialogOptions, signal: handlerSignal }; + if (dialogOptions.signal === handlerSignal) return dialogOptions; + return { ...dialogOptions, signal: AbortSignal.any([dialogOptions.signal, handlerSignal]) }; +} + +function createHandlerUIContext(ui: ExtensionUIContext, handlerSignal: AbortSignal): ExtensionUIContext { + const askDialog = ui.askDialog; + return { + ...ui, + select: (title, options, dialogOptions) => + ui.select(title, options, attachHandlerSignal(dialogOptions, handlerSignal)), + confirm: (title, message, dialogOptions) => + ui.confirm(title, message, attachHandlerSignal(dialogOptions, handlerSignal)), + input: (title, placeholder, dialogOptions) => + ui.input(title, placeholder, attachHandlerSignal(dialogOptions, handlerSignal)), + askDialog: askDialog + ? (questions, dialogOptions) => + askDialog.call(ui, questions, attachHandlerSignal(dialogOptions, handlerSignal)) + : undefined, + editor: (title, prefill, dialogOptions, editorOptions) => + ui.editor(title, prefill, attachHandlerSignal(dialogOptions, handlerSignal), editorOptions), + }; +} + /** * Race `work` against a `timeoutMs` budget and optional cancellation signal, * clearing the timer and abort listener as soon as one branch settles. @@ -118,19 +148,36 @@ const EXTENSION_HANDLER_ABORTED = Symbol("extensionHandlerAborted"); * can `clearTimeout` on the winning branch. */ async function raceHandlerWithTimeout( - work: Promise, + work: (handlerSignal: AbortSignal) => Promise | T, timeoutMs: number, signal?: AbortSignal, ): Promise { if (signal?.aborted) return EXTENSION_HANDLER_ABORTED; + + const timeoutController = new AbortController(); + const handlerSignal = signal ? AbortSignal.any([signal, timeoutController.signal]) : timeoutController.signal; + const workPromise = Promise.resolve(work(handlerSignal)); const { promise: interruptPromise, resolve: resolveInterrupt } = Promise.withResolvers< typeof EXTENSION_HANDLER_TIMEOUT | typeof EXTENSION_HANDLER_ABORTED >(); - const timer = setTimeout(() => resolveInterrupt(EXTENSION_HANDLER_TIMEOUT), timeoutMs); + const timer = setTimeout(() => { + timeoutController.abort(new DOMException(`Handler timed out after ${timeoutMs}ms`, "TimeoutError")); + resolveInterrupt(EXTENSION_HANDLER_TIMEOUT); + }, timeoutMs); const onAbort = () => resolveInterrupt(EXTENSION_HANDLER_ABORTED); signal?.addEventListener("abort", onAbort, { once: true }); try { - return await Promise.race([work, interruptPromise]); + const result = await Promise.race([workPromise, interruptPromise]); + if (result === EXTENSION_HANDLER_TIMEOUT || result === EXTENSION_HANDLER_ABORTED) { + await Promise.race([ + workPromise.then( + () => undefined, + () => undefined, + ), + Bun.sleep(0), + ]); + } + return result; } finally { clearTimeout(timer); signal?.removeEventListener("abort", onAbort); @@ -637,7 +684,15 @@ export class ExtensionRunner { : undefined; if (signal?.aborted) return undefined; try { - const handlerResult = await raceHandlerWithTimeout(Promise.resolve(handler(event, ctx)), timeoutMs, signal); + const handlerResult = await raceHandlerWithTimeout( + handlerSignal => + handler(event, { + ...ctx, + ui: createHandlerUIContext(ctx.ui, handlerSignal), + }), + timeoutMs, + signal, + ); if (handlerResult === EXTENSION_HANDLER_ABORTED) return undefined; if (handlerResult === EXTENSION_HANDLER_TIMEOUT) { const error = `handler timed out after ${timeoutMs}ms`; @@ -799,7 +854,14 @@ export class ExtensionRunner { for (const handler of handlers) { try { - const handlerResult = await raceHandlerWithTimeout(Promise.resolve(handler(event, ctx)), timeoutMs); + const handlerResult = await raceHandlerWithTimeout( + handlerSignal => + handler(event, { + ...ctx, + ui: createHandlerUIContext(ctx.ui, handlerSignal), + }), + timeoutMs, + ); if (handlerResult === EXTENSION_HANDLER_TIMEOUT) { const error = `handler timed out after ${timeoutMs}ms`; 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 097b42aea..21014ced4 100644 --- a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts +++ b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts @@ -86,7 +86,7 @@ export class ExtensionUiController { const uiContext: ExtensionUIContext = { timeoutStartsOnPresentation: true, select: (title, options, dialogOptions) => this.showCollabAwareSelector(title, options, dialogOptions), - confirm: (title, message, _dialogOptions) => this.showHookConfirm(title, message), + confirm: (title, message, dialogOptions) => this.showHookConfirm(title, message, dialogOptions), input: (title, placeholder, dialogOptions) => this.showHookInput(title, placeholder, dialogOptions), askDialog: (questions, dialogOptions) => this.showAskDialog(questions, dialogOptions), notify: (message, type) => this.showHookNotify(message, type), @@ -941,8 +941,8 @@ export class ExtensionUiController { /** * Show a confirmation dialog for hooks. */ - async showHookConfirm(title: string, message: string): Promise { - const result = await this.showHookSelector(`${title}\n${message}`, ["Yes", "No"]); + async showHookConfirm(title: string, message: string, dialogOptions?: ExtensionUIDialogOptions): Promise { + const result = await this.showHookSelector(`${title}\n${message}`, ["Yes", "No"], dialogOptions); return result === "Yes"; } diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index f888eaae4..596cf10f7 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -15,7 +15,11 @@ import { ExtensionRunner, testSetExtensionHandlerTimeoutMs, } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/runner"; -import type { ExtensionError, ExtensionServiceTier } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/types"; +import type { + ExtensionError, + ExtensionServiceTier, + ExtensionUIContext, +} from "@oh-my-pi/pi-coding-agent/extensibility/extensions/types"; import { ExtensionToolWrapper } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/wrapper"; import { Type } from "@oh-my-pi/pi-coding-agent/extensibility/typebox"; import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; @@ -1236,6 +1240,90 @@ describe("ExtensionRunner", () => { warnSpy.mockRestore(); }); + + it("aborts a tool_call handler's confirmation before returning its timeout block", 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) => { + await ctx.ui.confirm("High-risk command", "Allow this command?"); + fs.writeFileSync(${JSON.stringify(markerPath)}, "settled"); + }); + } + `, + ); + + const result = await loadTestExtensions([extensionPath]); + const runner = new ExtensionRunner( + result.extensions, + result.runtime, + tempDir.path(), + sessionManager, + modelRegistry, + ); + const dialog = Promise.withResolvers(); + let dialogSignal: AbortSignal | undefined; + const uiContext: ExtensionUIContext = { + ...runner.getUIContext(), + confirm: async (_title, _message, dialogOptions) => { + dialogSignal = dialogOptions?.signal; + dialogSignal?.addEventListener("abort", () => dialog.resolve(false), { once: true }); + return await dialog.promise; + }, + }; + 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, + ); + testSetExtensionHandlerTimeoutMs(10); + + 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); + + await expect(wrapped.execute("tool-call-id", {})).rejects.toThrow( + `Extension ${extensionPath} timed out after 10ms`, + ); + + expect(dialogSignal?.aborted).toBe(true); + expect(fs.readFileSync(markerPath, "utf8")).toBe("settled"); + }); }); describe("memory context", () => { diff --git a/packages/coding-agent/test/hook-editor.test.ts b/packages/coding-agent/test/hook-editor.test.ts index 06ff5d95b..f55c33531 100644 --- a/packages/coding-agent/test/hook-editor.test.ts +++ b/packages/coding-agent/test/hook-editor.test.ts @@ -587,4 +587,20 @@ describe("ExtensionUiController dialog serialization", () => { expect(ctx.hookSelector).toBeUndefined(); expect(editorContainer.children).toEqual([editor]); }); + it("dismisses a confirmation and restores the editor when its signal aborts", async () => { + const { ctx, editor, editorContainer } = createControllerContext(); + const controller = new ExtensionUiController(ctx); + const abortController = new AbortController(); + + const result = controller.showHookConfirm("High-risk command", "Allow this command?", { + signal: abortController.signal, + }); + expect(ctx.hookSelector).toBeDefined(); + + abortController.abort(); + + expect(await result).toBe(false); + expect(ctx.hookSelector).toBeUndefined(); + expect(editorContainer.children).toEqual([editor]); + }); });