diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index e2b55fc74..99f1ea9b7 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -12,6 +12,7 @@ - Fixed MCP resources with custom URI schemes being treated as missing filesystem paths. `read` and `omp read` now resolve server-advertised native resource URIs such as `ags://capabilities/current-host`, while preserving the existing `mcp://` form. - Fixed custom LSP servers sending `languageId: "plaintext"` for extensions outside the built-in language map by honoring an optional per-server `languageId` in `lsp.json` for disk and in-memory document opens ([#6800](https://github.com/can1357/oh-my-pi/issues/6800)). +- 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)). ### Fixed diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index 477ed18be..d7df257ae 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,50 @@ 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; + const dialogMethods = { + 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), + } satisfies Pick; + const delegatedMethods = new Map(); + + return new Proxy(ui, { + get(target, property) { + if (Object.hasOwn(dialogMethods, property)) { + return Reflect.get(dialogMethods, property, dialogMethods); + } + const cached = delegatedMethods.get(property); + if (cached) return cached; + const value: unknown = Reflect.get(target, property, target); + if (typeof value !== "function") return value; + const delegated: unknown = value.bind(target); + delegatedMethods.set(property, delegated); + return delegated; + }, + }); +} + /** * Race `work` against a `timeoutMs` budget and optional cancellation signal, * clearing the timer and abort listener as soon as one branch settles. @@ -118,19 +163,37 @@ 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 { promise: interruptPromise, resolve: resolveInterrupt } = Promise.withResolvers< typeof EXTENSION_HANDLER_TIMEOUT | typeof EXTENSION_HANDLER_ABORTED >(); - const timer = setTimeout(() => resolveInterrupt(EXTENSION_HANDLER_TIMEOUT), timeoutMs); const onAbort = () => resolveInterrupt(EXTENSION_HANDLER_ABORTED); signal?.addEventListener("abort", onAbort, { once: true }); + const timer = setTimeout(() => { + timeoutController.abort(new DOMException(`Handler timed out after ${timeoutMs}ms`, "TimeoutError")); + resolveInterrupt(EXTENSION_HANDLER_TIMEOUT); + }, timeoutMs); try { - return await Promise.race([work, interruptPromise]); + if (signal?.aborted) return EXTENSION_HANDLER_ABORTED; + const workPromise = Promise.resolve(work(handlerSignal)); + const result = await Promise.race([workPromise, interruptPromise]); + if (result === EXTENSION_HANDLER_TIMEOUT) { + await Promise.race([ + workPromise.then( + () => undefined, + () => undefined, + ), + Bun.sleep(0), + ]); + } + return result; } finally { clearTimeout(timer); signal?.removeEventListener("abort", onAbort); @@ -637,7 +700,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 +870,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/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index 5fe031bb1..19b74c64c 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -603,6 +603,57 @@ export function requestRpcEditor( } as RpcExtensionUIRequest); return promise; } + +/** Sends an RPC extension dialog and cancels the remote presentation when its signal aborts. */ +export function requestRpcDialog( + pendingRequests: Map, + output: RpcOutput, + opts: ExtensionUIDialogOptions | undefined, + defaultValue: T, + request: Record, + parseResponse: (response: RpcExtensionUIResponse) => T, +): Promise { + if (opts?.signal?.aborted) return Promise.resolve(defaultValue); + + const id = Snowflake.next() as string; + const { promise, resolve, reject } = Promise.withResolvers(); + let timeoutId: NodeJS.Timeout | undefined; + + const cleanup = () => { + clearTimeout(timeoutId); + opts?.signal?.removeEventListener("abort", onAbort); + pendingRequests.delete(id); + }; + const onAbort = () => { + output({ + type: "extension_ui_request", + id: Snowflake.next() as string, + method: "cancel", + targetId: id, + } as RpcExtensionUIRequest); + cleanup(); + resolve(defaultValue); + }; + opts?.signal?.addEventListener("abort", onAbort, { once: true }); + + if (opts?.timeout !== undefined) { + timeoutId = setTimeout(() => { + opts.onTimeout?.(); + cleanup(); + resolve(defaultValue); + }, opts.timeout); + } + + pendingRequests.set(id, { + resolve: response => { + cleanup(); + resolve(parseResponse(response)); + }, + reject, + }); + output({ type: "extension_ui_request", id, ...request } as RpcExtensionUIRequest); + return promise; +} /** * Run in RPC mode. * Listens for JSON commands on stdin, outputs events and responses on stdout. @@ -685,56 +736,14 @@ export async function runRpcMode( private output: (obj: RpcResponse | RpcExtensionUIRequest | object) => void, ) {} - /** Helper for dialog methods with signal/timeout support */ - #createDialogPromise( - opts: ExtensionUIDialogOptions | undefined, - defaultValue: T, - request: Record, - parseResponse: (response: RpcExtensionUIResponse) => T, - ): Promise { - if (opts?.signal?.aborted) return Promise.resolve(defaultValue); - - const id = Snowflake.next() as string; - const { promise, resolve, reject } = Promise.withResolvers(); - let timeoutId: NodeJS.Timeout | undefined; - - const cleanup = () => { - if (timeoutId) clearTimeout(timeoutId); - opts?.signal?.removeEventListener("abort", onAbort); - this.pendingRequests.delete(id); - }; - - const onAbort = () => { - cleanup(); - resolve(defaultValue); - }; - opts?.signal?.addEventListener("abort", onAbort, { once: true }); - - if (opts?.timeout !== undefined) { - timeoutId = setTimeout(() => { - opts.onTimeout?.(); - cleanup(); - resolve(defaultValue); - }, opts.timeout); - } - - this.pendingRequests.set(id, { - resolve: (response: RpcExtensionUIResponse) => { - cleanup(); - resolve(parseResponse(response)); - }, - reject, - }); - this.output({ type: "extension_ui_request", id, ...request } as RpcExtensionUIRequest); - return promise; - } - select( title: string, options: ExtensionUISelectItem[], dialogOptions?: ExtensionUIDialogOptions, ): Promise { - return this.#createDialogPromise( + return requestRpcDialog( + this.pendingRequests, + this.output, dialogOptions, undefined, { @@ -748,7 +757,9 @@ export async function runRpcMode( } confirm(title: string, message: string, dialogOptions?: ExtensionUIDialogOptions): Promise { - return this.#createDialogPromise( + return requestRpcDialog( + this.pendingRequests, + this.output, dialogOptions, false, { method: "confirm", title, message, timeout: dialogOptions?.timeout }, @@ -768,7 +779,9 @@ export async function runRpcMode( placeholder?: string, dialogOptions?: ExtensionUIDialogOptions, ): Promise { - return this.#createDialogPromise( + return requestRpcDialog( + this.pendingRequests, + this.output, dialogOptions, undefined, { method: "input", title, placeholder, timeout: dialogOptions?.timeout }, diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index f888eaae4..929ef9af9 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"; @@ -855,6 +859,81 @@ describe("ExtensionRunner", () => { ]); }); + it("observes a session_stop signal aborted synchronously by the handler", async () => { + const extensionPath = path.join(tempDir.path(), "self-cancel-session-stop.ts"); + await Bun.write( + extensionPath, + ` + export default function(pi) { + pi.on("session_stop", async (_event, ctx) => { + ctx.abort(); + await Promise.withResolvers().promise; + }); + } + `, + ); + + const result = await loadTestExtensions([extensionPath]); + const runner = new ExtensionRunner( + result.extensions, + result.runtime, + tempDir.path(), + sessionManager, + modelRegistry, + ); + const controller = new AbortController(); + 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: () => controller.abort(), + hasPendingMessages: () => false, + shutdown: () => {}, + getContextUsage: () => undefined, + compact: async () => {}, + getSystemPrompt: () => [], + }, + ); + vi.useFakeTimers(); + try { + testSetExtensionHandlerTimeoutMs(100); + const emission = runner.emitSessionStop({ + messages: [], + turn_id: 0, + session_id: "session-123", + stop_hook_active: false, + signal: controller.signal, + }); + let settled = false; + void emission.then(() => { + settled = true; + }); + for (let attempts = 0; attempts < 10 && !settled; attempts++) { + await Promise.resolve(); + } + + expect(controller.signal.aborted).toBe(true); + expect(settled).toBe(true); + await emission; + } finally { + vi.useRealTimers(); + } + }); it("continues to later handlers after empty continuation feedback", async () => { await Bun.write( path.join(extensionsDir, "session-stop-empty.ts"), @@ -1236,6 +1315,95 @@ 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) => { + ctx.ui.notify("Waiting for confirmation"); + 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 notify = vi.fn(); + 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 }, + 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, + ); + 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(notify).toHaveBeenCalledWith("Waiting for confirmation"); + + 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]); + }); }); diff --git a/packages/coding-agent/test/rpc-extension-ui.test.ts b/packages/coding-agent/test/rpc-extension-ui.test.ts new file mode 100644 index 000000000..fede2aba1 --- /dev/null +++ b/packages/coding-agent/test/rpc-extension-ui.test.ts @@ -0,0 +1,40 @@ +import { describe, expect, it, vi } from "bun:test"; +import { type PendingExtensionRequest, requestRpcDialog } from "@oh-my-pi/pi-coding-agent/modes/rpc/rpc-mode"; + +describe("RPC extension UI", () => { + it("cancels the remote dialog when its signal aborts", async () => { + const pendingRequests = new Map(); + const output = vi.fn<(frame: object) => void>(); + const controller = new AbortController(); + const result = requestRpcDialog( + pendingRequests, + output, + { signal: controller.signal }, + false, + { method: "confirm", title: "High-risk command", message: "Allow this command?" }, + response => ("confirmed" in response ? response.confirmed : false), + ); + const request = output.mock.calls[0]?.[0]; + if (!request || !("id" in request) || typeof request.id !== "string") { + throw new Error("Expected the RPC dialog request to carry an id"); + } + + controller.abort(); + + expect(await result).toBe(false); + expect(output).toHaveBeenNthCalledWith(1, { + type: "extension_ui_request", + id: request.id, + method: "confirm", + title: "High-risk command", + message: "Allow this command?", + }); + expect(output).toHaveBeenNthCalledWith(2, { + type: "extension_ui_request", + id: expect.any(String), + method: "cancel", + targetId: request.id, + }); + expect(pendingRequests.size).toBe(0); + }); +});