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

@@ -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 `<cell>` 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)).
 

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<T>(
-	work: Promise<T>,
+	work: (handlerSignal: AbortSignal) => Promise<T> | T,
 	timeoutMs: number,
 	signal?: AbortSignal,
 ): Promise<T | typeof EXTENSION_HANDLER_TIMEOUT | typeof EXTENSION_HANDLER_ABORTED> {
 	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`;

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<boolean> {
-		const result = await this.showHookSelector(`${title}\n${message}`, ["Yes", "No"]);
+	async showHookConfirm(title: string, message: string, dialogOptions?: ExtensionUIDialogOptions): Promise<boolean> {
+		const result = await this.showHookSelector(`${title}\n${message}`, ["Yes", "No"], dialogOptions);
 		return result === "Yes";
 	}
 

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<boolean>();
+			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", () => {

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