diff --git a/packages/coding-agent/src/extensibility/extensions/types.ts b/packages/coding-agent/src/extensibility/extensions/types.ts index cfff0ed01..59b67e3b9 100644 --- a/packages/coding-agent/src/extensibility/extensions/types.ts +++ b/packages/coding-agent/src/extensibility/extensions/types.ts @@ -125,6 +125,8 @@ export interface ExtensionUIDialogOptions { timeout?: number; /** Invoked when the UI times out while waiting for a selection/input */ onTimeout?: () => void; + /** Invoked when user input resets a UI-managed timeout countdown */ + onTimeoutReset?: () => void; /** Initial cursor position for select dialogs (0-indexed) */ initialIndex?: number; /** Render an outlined list for select dialogs */ diff --git a/packages/coding-agent/src/modes/components/hook-selector.ts b/packages/coding-agent/src/modes/components/hook-selector.ts index 99ea1ea1e..f5e5570b3 100644 --- a/packages/coding-agent/src/modes/components/hook-selector.ts +++ b/packages/coding-agent/src/modes/components/hook-selector.ts @@ -61,6 +61,7 @@ export interface HookSelectorOptions { tui?: TUI; timeout?: number; onTimeout?: () => void; + onTimeoutReset?: () => void; initialIndex?: number; outline?: boolean; maxVisible?: number; @@ -178,6 +179,7 @@ export class HookSelectorComponent extends Container { #onLeftCallback: (() => void) | undefined; #onRightCallback: (() => void) | undefined; #onExternalEditorCallback: (() => void) | undefined; + #onTimeoutResetCallback: (() => void) | undefined; #slider: HookSelectorSlider | undefined; #sliderIndex: number = 0; #sliderComponent: Text | undefined; @@ -213,6 +215,7 @@ export class HookSelectorComponent extends Container { this.#onLeftCallback = opts?.onLeft; this.#onRightCallback = opts?.onRight; this.#onExternalEditorCallback = opts?.onExternalEditor; + this.#onTimeoutResetCallback = opts?.onTimeoutReset; if (opts?.slider && opts.slider.segments.length > 0) { this.#slider = opts.slider; this.#sliderIndex = Math.max(0, Math.min(opts.slider.index, opts.slider.segments.length - 1)); @@ -633,8 +636,10 @@ export class HookSelectorComponent extends Container { } handleInput(keyData: string): void { - // Reset countdown on any interaction - this.#countdown?.reset(); + if (this.#countdown) { + this.#countdown.reset(); + this.#onTimeoutResetCallback?.(); + } if (matchesSelectCancel(keyData)) { this.#onCancelCallback(); 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 684642b12..ef27b3644 100644 --- a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts +++ b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts @@ -624,6 +624,7 @@ export class ExtensionUiController { initialIndex: dialogOptions?.initialIndex, timeout: dialogOptions?.timeout, onTimeout: dialogOptions?.onTimeout, + onTimeoutReset: dialogOptions?.onTimeoutReset, tui: this.ctx.ui, outline: dialogOptions?.outline, disabledIndices: dialogOptions?.disabledIndices, diff --git a/packages/coding-agent/src/tools/ask.ts b/packages/coding-agent/src/tools/ask.ts index c58ef975a..15fdc8994 100644 --- a/packages/coding-agent/src/tools/ask.ts +++ b/packages/coding-agent/src/tools/ask.ts @@ -389,6 +389,7 @@ interface UIContext { signal?: AbortSignal; outline?: boolean; onTimeout?: () => void; + onTimeoutReset?: () => void; onLeft?: () => void; onRight?: () => void; helpText?: string; @@ -432,18 +433,29 @@ async function askSingleQuestion( const helpText = navigation ? "up/down navigate enter select ←/→ question esc cancel" : "up/down navigate enter select esc cancel"; - const timeoutController = typeof timeout === "number" && timeout > 0 ? new AbortController() : undefined; + const timeoutMs = typeof timeout === "number" && timeout > 0 ? timeout : undefined; + const timeoutController = timeoutMs === undefined ? undefined : new AbortController(); const dialogSignal = signal && timeoutController ? AbortSignal.any([signal, timeoutController.signal]) : (timeoutController?.signal ?? signal); let timeoutId: NodeJS.Timeout | undefined; + let timeoutStartedMs = Date.now(); + const armFallbackTimeout = (durationMs: number) => { + clearTimeout(timeoutId); + timeoutStartedMs = Date.now(); + timeoutId = setTimeout(() => { + timeoutTriggered = true; + timeoutController?.abort(); + }, durationMs); + }; const dialogOptions = { initialIndex, timeout, signal: dialogSignal, outline: true, onTimeout, + onTimeoutReset: timeoutMs === undefined ? undefined : () => armFallbackTimeout(timeoutMs), helpText, selectionMarker: marker?.selectionMarker, checkedIndices: marker?.checkedIndices, @@ -459,12 +471,8 @@ async function askSingleQuestion( } : undefined, }; - const startMs = Date.now(); - if (timeoutController && typeof timeout === "number") { - timeoutId = setTimeout(() => { - timeoutTriggered = true; - timeoutController.abort(); - }, timeout); + if (timeoutMs !== undefined) { + armFallbackTimeout(timeoutMs); } try { const choice = dialogSignal @@ -475,7 +483,7 @@ async function askSingleQuestion( // `onTimeout`: their auto-cancel resolves right at the deadline. A // cancel arriving well past the deadline is a deliberate user Esc on // a surface that kept the dialog open — keep treating it as a cancel. - const elapsed = Date.now() - startMs; + const elapsed = Date.now() - timeoutStartedMs; timeoutTriggered = elapsed >= timeout && elapsed <= timeout + TIMEOUT_DETECTION_TOLERANCE_MS; } return { choice, timedOut: timeoutTriggered, navigation: navigationAction }; diff --git a/packages/coding-agent/test/ask-timeout.test.ts b/packages/coding-agent/test/ask-timeout.test.ts index ebbf14f67..34a153207 100644 --- a/packages/coding-agent/test/ask-timeout.test.ts +++ b/packages/coding-agent/test/ask-timeout.test.ts @@ -1,10 +1,18 @@ import { afterEach, beforeAll, describe, expect, it, vi } from "bun:test"; import type { AgentToolContext, AgentToolResult } from "@oh-my-pi/pi-agent-core"; +import type { TUI } from "@oh-my-pi/pi-tui"; +import type { ExtensionUIDialogOptions, ExtensionUISelectItem } from "../src/extensibility/extensions"; +import { HookSelectorComponent } from "../src/modes/components/hook-selector"; import { getThemeByName, setThemeInstance } from "../src/modes/theme/theme"; import type { ToolSession } from "../src/tools"; import { AskTool, type AskToolDetails } from "../src/tools/ask"; type AskExecutionResult = AgentToolResult; +type AskSelect = ( + title: string, + options: ExtensionUISelectItem[], + dialogOptions?: ExtensionUIDialogOptions, +) => Promise; async function drainMicrotasks(): Promise { await Promise.resolve(); @@ -40,7 +48,7 @@ describe("AskTool timeout", () => { it("auto-selects the recommended option when the selector does not settle", async () => { vi.useFakeTimers(); - const select = vi.fn(() => new Promise(() => {})); + const select = vi.fn(() => new Promise(() => {})); const abort = vi.fn(); const context = { hasUI: true, @@ -88,4 +96,83 @@ describe("AskTool timeout", () => { expect(result?.details?.timedOut).toBe(true); expect(abort).not.toHaveBeenCalled(); }); + + it("honors selector timeout resets before using the fallback timeout", async () => { + vi.useFakeTimers(); + let resetTimeout: (() => void) | undefined; + const select = vi.fn((_title, _options, dialogOptions) => { + resetTimeout = dialogOptions?.onTimeoutReset; + return new Promise(() => {}); + }); + const abort = vi.fn(); + const context = { + hasUI: true, + ui: { + select, + editor: vi.fn(), + }, + abort, + } as unknown as AgentToolContext; + let result: AskExecutionResult | undefined; + let rejection: unknown; + + void createAskTool() + .execute( + "ask-timeout", + { + questions: [ + { + id: "db", + question: "Which database?", + options: [{ label: "SQLite" }, { label: "Postgres" }], + recommended: 1, + }, + ], + }, + undefined, + undefined, + context, + ) + .then( + value => { + result = value; + }, + error => { + rejection = error; + }, + ); + + await drainMicrotasks(); + expect(resetTimeout).toBeDefined(); + + vi.advanceTimersByTime(9); + resetTimeout?.(); + vi.advanceTimersByTime(9); + await drainMicrotasks(); + + expect(result).toBeUndefined(); + + vi.advanceTimersByTime(1); + await drainMicrotasks(); + + expect(rejection).toBeUndefined(); + expect(result?.details?.selectedOptions).toEqual(["Postgres"]); + expect(result?.details?.timedOut).toBe(true); + expect(abort).not.toHaveBeenCalled(); + }); + + it("notifies callers when selector input resets the UI countdown", () => { + vi.useFakeTimers(); + const onTimeoutReset = vi.fn(); + const selector = new HookSelectorComponent("Pick one", ["SQLite", "Postgres"], vi.fn(), vi.fn(), { + timeout: 10, + tui: { requestRender: vi.fn() } as unknown as TUI, + onTimeoutReset, + }); + + selector.handleInput("j"); + + expect(onTimeoutReset).toHaveBeenCalledTimes(1); + selector.dispose(); + }); });