From 5e6f12b278546c5477df867a421e7f1d909c64cf Mon Sep 17 00:00:00 2001 From: roboomp Date: Mon, 27 Jul 2026 13:28:24 +0000 Subject: [PATCH 1/4] fix(extensions): cancelled timed-out handler dialogs Forwarded confirmation dialog options in the interactive TUI and scoped extension UI dialogs to each handler watchdog signal. Added regressions for direct confirmation cancellation and fail-closed tool-call timeout cleanup. Fixes #6805 --- packages/coding-agent/CHANGELOG.md | 1 + .../src/extensibility/extensions/runner.ts | 72 +++++++++++++-- .../controllers/extension-ui-controller.ts | 6 +- .../test/extensions-runner.test.ts | 90 ++++++++++++++++++- .../coding-agent/test/hook-editor.test.ts | 16 ++++ 5 files changed, 176 insertions(+), 9 deletions(-) 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]); + }); }); From 4176a6c7993a25c97ed16f83191d33ba416c139b Mon Sep 17 00:00:00 2001 From: roboomp Date: Mon, 27 Jul 2026 13:37:05 +0000 Subject: [PATCH 2/4] fix(extensions): registered abort before handler start Attached the session-stop abort listener and rechecked cancellation before invoking extension work, preventing synchronous ctx.abort() calls from being missed. Added deterministic coverage for a handler that aborts and then waits on non-UI work. --- .../src/extensibility/extensions/runner.ts | 9 ++- .../test/extensions-runner.test.ts | 75 +++++++++++++++++++ 2 files changed, 80 insertions(+), 4 deletions(-) diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index d28ca07c4..16a028b3a 100644 --- a/packages/coding-agent/src/extensibility/extensions/runner.ts +++ b/packages/coding-agent/src/extensibility/extensions/runner.ts @@ -156,19 +156,20 @@ async function raceHandlerWithTimeout( 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 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); - const onAbort = () => resolveInterrupt(EXTENSION_HANDLER_ABORTED); - signal?.addEventListener("abort", onAbort, { once: true }); try { + 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 || result === EXTENSION_HANDLER_ABORTED) { + if (result === EXTENSION_HANDLER_TIMEOUT) { await Promise.race([ workPromise.then( () => undefined, diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 596cf10f7..73c16f666 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -859,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"), From eb40353a9523d4184b9c1c166e176c716f14dc54 Mon Sep 17 00:00:00 2001 From: roboomp Date: Mon, 27 Jul 2026 13:44:33 +0000 Subject: [PATCH 3/4] fix(extensions): preserved prototype-backed UI methods Replaced the scoped UI object spread with a delegating proxy that binds inherited methods to the original context while overriding only abort-capable dialogs. Extended watchdog coverage with a prototype-backed notification method matching RPC UI contexts. --- .../src/extensibility/extensions/runner.ts | 21 ++++++++++++++++--- .../test/extensions-runner.test.ts | 19 ++++++++++------- 2 files changed, 30 insertions(+), 10 deletions(-) diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index 16a028b3a..d7df257ae 100644 --- a/packages/coding-agent/src/extensibility/extensions/runner.ts +++ b/packages/coding-agent/src/extensibility/extensions/runner.ts @@ -118,8 +118,7 @@ function attachHandlerSignal( function createHandlerUIContext(ui: ExtensionUIContext, handlerSignal: AbortSignal): ExtensionUIContext { const askDialog = ui.askDialog; - return { - ...ui, + const dialogMethods = { select: (title, options, dialogOptions) => ui.select(title, options, attachHandlerSignal(dialogOptions, handlerSignal)), confirm: (title, message, dialogOptions) => @@ -132,7 +131,23 @@ function createHandlerUIContext(ui: ExtensionUIContext, handlerSignal: AbortSign : 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; + }, + }); } /** diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 73c16f666..929ef9af9 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -1326,6 +1326,7 @@ describe("ExtensionRunner", () => { 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"); }); @@ -1343,14 +1344,17 @@ describe("ExtensionRunner", () => { ); 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; - }, + 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: () => {}, @@ -1395,6 +1399,7 @@ describe("ExtensionRunner", () => { 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"); From 452932b291b54612c975fab3554b3eb3755883b6 Mon Sep 17 00:00:00 2001 From: roboomp Date: Mon, 27 Jul 2026 13:53:30 +0000 Subject: [PATCH 4/4] fix(rpc): cancelled aborted extension dialogs Moved RPC dialog request lifecycle into a reusable helper that emits a cancel frame targeting the original request before settling an aborted local promise. Added coverage for remote confirmation cancellation and pending-request cleanup. --- .../coding-agent/src/modes/rpc/rpc-mode.ts | 107 ++++++++++-------- .../test/rpc-extension-ui.test.ts | 40 +++++++ 2 files changed, 100 insertions(+), 47 deletions(-) create mode 100644 packages/coding-agent/test/rpc-extension-ui.test.ts 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/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); + }); +});