smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(extensions): pause tool-call timeout during human dialogs

Browse Source
This commit is contained in:
Seljuke
2026-08-12 00:24:09 +02:00
parent 326d24bd40
commit 7e384fbc4f
5 changed files with 295 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/src/config/settings-schema.ts
+12
View File
@@ -5494,6 +5494,18 @@ export const SETTINGS_SCHEMA = {
"commit.changelogMaxDiffChars": { type: "number", default: 120000 },
"extensionHandlers.toolCallTimeoutMs": {
type: "number",
default: 30_000,
ui: {
tab: "tools",
group: "Extensions",
label: "Tool Call Handler Timeout (ms)",
description:
"Active-work timeout for extension tool_call handlers; time awaiting OMP-owned dialogs does not count",
},
},
"dev.autoqa": {
type: "boolean",
default: true,
packages/coding-agent/src/extensibility/extensions/runner.ts
+102 -21
View File
@@ -116,6 +116,11 @@ function handlerTimeoutForEvent(eventType: string): number {
const EXTENSION_HANDLER_TIMEOUT = Symbol("extensionHandlerTimeout");
const EXTENSION_HANDLER_ABORTED = Symbol("extensionHandlerAborted");
interface HandlerTimeoutBudget {
pause(): void;
resume(): void;
}
function attachHandlerSignal(
dialogOptions: ExtensionUIDialogOptions | undefined,
handlerSignal: AbortSignal,
@@ -126,22 +131,35 @@ function attachHandlerSignal(
return { ...dialogOptions, signal: AbortSignal.any([dialogOptions.signal, handlerSignal]) };
}
function createHandlerUIContext(ui: ExtensionUIContext, handlerSignal: AbortSignal): ExtensionUIContext {
function createHandlerUIContext(
ui: ExtensionUIContext,
handlerSignal: AbortSignal,
timeoutBudget?: HandlerTimeoutBudget,
): ExtensionUIContext {
const askDialog = ui.askDialog;
const runDialog = async <T>(dialog: () => Promise<T>): Promise<T> => {
timeoutBudget?.pause();
try {
return await dialog();
} finally {
timeoutBudget?.resume();
}
};
const dialogMethods = {
select: (title, options, dialogOptions) =>
ui.select(title, options, attachHandlerSignal(dialogOptions, handlerSignal)),
runDialog(() => ui.select(title, options, attachHandlerSignal(dialogOptions, handlerSignal))),
confirm: (title, message, dialogOptions) =>
ui.confirm(title, message, attachHandlerSignal(dialogOptions, handlerSignal)),
runDialog(() => ui.confirm(title, message, attachHandlerSignal(dialogOptions, handlerSignal))),
input: (title, placeholder, dialogOptions) =>
ui.input(title, placeholder, attachHandlerSignal(dialogOptions, handlerSignal)),
runDialog(() => ui.input(title, placeholder, attachHandlerSignal(dialogOptions, handlerSignal))),
askDialog: askDialog
? (questions, dialogOptions) =>
askDialog.call(ui, questions, attachHandlerSignal(dialogOptions, handlerSignal))
runDialog(() => askDialog.call(ui, questions, attachHandlerSignal(dialogOptions, handlerSignal)))
: undefined,
custom: (factory, options) => runDialog(() => ui.custom(factory, options)),
editor: (title, prefill, dialogOptions, editorOptions) =>
ui.editor(title, prefill, attachHandlerSignal(dialogOptions, handlerSignal), editorOptions),
} satisfies Pick<ExtensionUIContext, "select" | "confirm" | "input" | "askDialog" | "editor">;
runDialog(() => ui.editor(title, prefill, attachHandlerSignal(dialogOptions, handlerSignal), editorOptions)),
} satisfies Pick<ExtensionUIContext, "select" | "confirm" | "input" | "askDialog" | "custom" | "editor">;
const delegatedMethods = new Map<PropertyKey, unknown>();
return new Proxy(ui, {
@@ -166,10 +184,14 @@ function createHandlerUIContext(ui: ExtensionUIContext, handlerSignal: AbortSign
* `pi.setModel()` and then reading `ctx.model` would see a stale model.
* Prototype delegation keeps every getter live while overriding `ui`.
*/
function createHandlerContext(ctx: ExtensionContext, handlerSignal: AbortSignal): ExtensionContext {
function createHandlerContext(
ctx: ExtensionContext,
handlerSignal: AbortSignal,
timeoutBudget?: HandlerTimeoutBudget,
): ExtensionContext {
const scoped: ExtensionContext = Object.create(ctx);
Object.defineProperty(scoped, "ui", {
value: createHandlerUIContext(ctx.ui, handlerSignal),
value: createHandlerUIContext(ctx.ui, handlerSignal, timeoutBudget),
enumerable: true,
configurable: true,
});
@@ -189,7 +211,7 @@ function createHandlerContext(ctx: ExtensionContext, handlerSignal: AbortSignal)
* can `clearTimeout` on the winning branch.
*/
async function raceHandlerWithTimeout<T>(
work: (handlerSignal: AbortSignal) => Promise<T> | T,
work: (handlerSignal: AbortSignal, timeoutBudget: HandlerTimeoutBudget) => Promise<T> | T,
timeoutMs: number,
signal?: AbortSignal,
): Promise<T | typeof EXTENSION_HANDLER_TIMEOUT | typeof EXTENSION_HANDLER_ABORTED> {
@@ -202,13 +224,52 @@ async function raceHandlerWithTimeout<T>(
>();
const onAbort = () => resolveInterrupt(EXTENSION_HANDLER_ABORTED);
signal?.addEventListener("abort", onAbort, { once: true });
const timer = setTimeout(() => {
let timer: Timer | undefined;
let remainingMs = timeoutMs;
let activeSince = performance.now();
let pauseDepth = 0;
let settled = false;
const clearTimer = () => {
if (timer === undefined) return;
clearTimeout(timer);
timer = undefined;
};
const expire = () => {
if (settled) return;
settled = true;
clearTimer();
timeoutController.abort(new DOMException(`Handler timed out after ${timeoutMs}ms`, "TimeoutError"));
resolveInterrupt(EXTENSION_HANDLER_TIMEOUT);
}, timeoutMs);
};
const armTimer = () => {
if (settled || pauseDepth > 0) return;
activeSince = performance.now();
timer = setTimeout(expire, Math.max(0, remainingMs));
};
const settle = () => {
if (settled) return;
settled = true;
clearTimer();
};
const timeoutBudget: HandlerTimeoutBudget = {
pause: () => {
if (settled) return;
pauseDepth++;
if (pauseDepth !== 1) return;
remainingMs = Math.max(0, remainingMs - (performance.now() - activeSince));
clearTimer();
if (remainingMs <= 0) expire();
},
resume: () => {
if (settled || pauseDepth === 0) return;
pauseDepth--;
if (pauseDepth === 0) armTimer();
},
};
armTimer();
try {
if (signal?.aborted) return EXTENSION_HANDLER_ABORTED;
const workPromise = Promise.resolve(work(handlerSignal));
const workPromise = Promise.resolve(work(handlerSignal, timeoutBudget));
const result = await Promise.race([workPromise, interruptPromise]);
if (result === EXTENSION_HANDLER_TIMEOUT) {
await Promise.race([
@@ -221,7 +282,7 @@ async function raceHandlerWithTimeout<T>(
}
return result;
} finally {
clearTimeout(timer);
settle();
signal?.removeEventListener("abort", onAbort);
}
}
@@ -1038,23 +1099,33 @@ export class ExtensionRunner {
ext: Extension,
timeoutMs: number,
onFailure?: (kind: "timeout" | "error", message: string) => TResult,
outerSignal?: AbortSignal,
): Promise<TResult | undefined> {
const signal =
// `session_stop` carries its own signal on the event; `tool_call` receives
// the outer dispatch signal (loop request or wrapper execute) so an abort
// while a handler awaits a human dialog cancels the dialog and settles the
// gate without executing the underlying tool. Compose whichever apply.
const sessionStopSignal =
event.type === "session_stop" && "signal" in event && event.signal instanceof AbortSignal
? event.signal
: undefined;
const signals = [outerSignal, sessionStopSignal].filter((s): s is AbortSignal => s !== undefined);
const signal = signals.length === 0 ? undefined : signals.length === 1 ? signals[0] : AbortSignal.any(signals);
if (signal?.aborted) return undefined;
const registrationScope: ToolRegistrationScope = { pending: new Set(), closed: false };
let handlerResult: TResult | typeof EXTENSION_HANDLER_TIMEOUT | typeof EXTENSION_HANDLER_ABORTED | undefined;
let handlerFailure: { error: unknown } | undefined;
try {
handlerResult = await raceHandlerWithTimeout(
async handlerSignal => {
async (handlerSignal, budget) => {
registrationScope.signal = handlerSignal;
let result: TResult | undefined;
try {
result = await this.#toolRegistrationScope.run(registrationScope, () =>
handler(event, createHandlerContext(ctx, handlerSignal)),
handler(
event,
createHandlerContext(ctx, handlerSignal, event.type === "tool_call" ? budget : undefined),
),
);
} catch (error) {
handlerFailure = { error };
@@ -1215,8 +1286,8 @@ export class ExtensionRunner {
/**
* Emit a `tool_call` event to every subscribed extension before the tool executes.
*
* Each handler is bounded by `extensionHandlerTimeoutMs` (default 30s). This
* matches the timeout policy already applied to `emitToolResult` and every
* Each handler is bounded by `extensionHandlers.toolCallTimeoutMs` (default
* 30s). This matches the timeout policy already applied to `emitToolResult` and every
* other handler routed through `#runHandlerWithTimeout`; without it a single
* hung extension (unresolved `await`, network call with no timeout) would
* park `ExtensionToolWrapper.execute` indefinitely and freeze tool
@@ -1227,9 +1298,9 @@ export class ExtensionRunner {
* pre-execution gate — an unresponsive extension MUST NOT be treated as
* silent consent to run the tool.
*/
async emitToolCall(event: ToolCallEvent): Promise<ToolCallEventResult | undefined> {
async emitToolCall(event: ToolCallEvent, signal?: AbortSignal): Promise<ToolCallEventResult | undefined> {
const ctx = this.createContext();
const timeoutMs = extensionHandlerTimeoutMs;
const timeoutMs = this.settings?.get("extensionHandlers.toolCallTimeoutMs") ?? extensionHandlerTimeoutMs;
let result: ToolCallEventResult | undefined;
for (const ext of this.extensions) {
@@ -1250,6 +1321,7 @@ export class ExtensionRunner {
? `Extension ${ext.path} timed out after ${timeoutMs}ms`
: `Extension ${ext.path} failed: ${message}`,
}),
signal,
);
if (handlerResult) {
@@ -1258,9 +1330,18 @@ export class ExtensionRunner {
return result;
}
}
// Fail closed when the outer dispatch aborted while a handler was
// pending: an aborted gate MUST NOT become silent consent to run the
// underlying tool. Symmetric with the timeout policy above.
if (signal?.aborted) {
return { block: true, reason: `Tool execution was cancelled while an extension handler was pending` };
}
}
}
if (signal?.aborted) {
return { block: true, reason: `Tool execution was cancelled while an extension handler was pending` };
}
return result;
}
packages/coding-agent/src/extensibility/extensions/wrapper.ts
+12 -9
View File
@@ -205,15 +205,18 @@ export class ExtensionToolWrapper<TParameters extends TSchema = TSchema, TDetail
let effectiveParams = params;
if (!loopEmittedToolCall && this.runner.hasHandlers("tool_call")) {
try {
const callResult = (await this.runner.emitToolCall({
type: "tool_call",
toolName: this.tool.name,
toolCallId,
input: normalizeToolEventInput(
this.tool.name,
resolveToolEventInput(this.tool, toolEventArgs(params, context)),
),
})) as ToolCallEventResult | undefined;
const callResult = (await this.runner.emitToolCall(
{
type: "tool_call",
toolName: this.tool.name,
toolCallId,
input: normalizeToolEventInput(
this.tool.name,
resolveToolEventInput(this.tool, toolEventArgs(params, context)),
),
},
signal,
)) as ToolCallEventResult | undefined;
if (callResult?.block) {
const reason = callResult.reason || "Tool execution was blocked by an extension";
packages/coding-agent/src/session/agent-session.ts
+11 -8
View File
@@ -1370,7 +1370,7 @@ export class AgentSession {
// Pre-scheduling tool_call wiring: extension handlers run at arg-prep
// time so a block/revision lands before concurrency resolution,
// tool_execution_start, and the wrapper's approval gate.
this.agent.beforeToolCall = ctx => this.#beforeToolCall(ctx);
this.agent.beforeToolCall = (ctx, signal) => this.#beforeToolCall(ctx, signal);
this.agent.providerSessionState = this.#providerSessionState;
this.#syncAgentSessionId();
this.#todo.syncFromBranch();
@@ -3276,7 +3276,7 @@ export class AgentSession {
* emit a second event (nested xd:// device dispatches and direct non-loop
* execution still emit there).
*/
async #beforeToolCall(ctx: BeforeToolCallContext): Promise<BeforeToolCallResult | undefined> {
async #beforeToolCall(ctx: BeforeToolCallContext, signal?: AbortSignal): Promise<BeforeToolCallResult | undefined> {
const runner = this.#extensionRunner;
if (!runner?.hasHandlers("tool_call")) return undefined;
const metadata = ctx.toolCall.providerMetadata;
@@ -3294,12 +3294,15 @@ export class AgentSession {
? { actions: computer.actions, pendingSafetyChecks: computer.pendingSafetyChecks }
: ctx.args;
runner.markToolCallEmitted(ctx.toolCall.id, ctx.tool.name);
const callResult = await runner.emitToolCall({
type: "tool_call",
toolName: ctx.tool.name,
toolCallId: ctx.toolCall.id,
input: normalizeToolEventInput(ctx.tool.name, resolveToolEventInput(ctx.tool, eventArgs)),
});
const callResult = await runner.emitToolCall(
{
type: "tool_call",
toolName: ctx.tool.name,
toolCallId: ctx.toolCall.id,
input: normalizeToolEventInput(ctx.tool.name, resolveToolEventInput(ctx.tool, eventArgs)),
},
signal,
);
if (callResult?.block) {
return { block: true, reason: callResult.reason || "Tool execution was blocked by an extension" };
}
packages/coding-agent/test/extensions-runner.test.ts
+158 -12
View File
@@ -10,6 +10,7 @@ import type { AgentMessage, AgentTool } from "@oh-my-pi/pi-agent-core";
import type { ImageContent, TextContent } from "@oh-my-pi/pi-ai";
import { getBundledModel } from "@oh-my-pi/pi-catalog/models";
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
import { discoverAndLoadExtensions, ExtensionRuntime } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/loader";
import {
EXTENSION_HANDLER_TIMEOUT_MS,
@@ -1301,7 +1302,7 @@ describe("ExtensionRunner", () => {
});
});
it("times out tool_call handlers with fail-closed policy so a hung extension cannot indefinitely block tool execution (#3948)", async () => {
it("uses the configured tool_call timeout and fails closed so a hung extension cannot block execution (#3948)", async () => {
const hangExtensionPath = path.join(tempDir.path(), "hang-tool-call.ts");
fs.writeFileSync(
hangExtensionPath,
@@ -1321,14 +1322,14 @@ describe("ExtensionRunner", () => {
tempDir.path(),
sessionManager,
modelRegistry,
undefined,
Settings.isolated({ "extensionHandlers.toolCallTimeoutMs": 10 }),
);
const warnSpy = vi.spyOn(logger, "warn").mockImplementation(() => {});
const errors: Array<{ extensionPath: string; event: string; error: string }> = [];
runner.onError(err => {
errors.push(err);
});
testSetExtensionHandlerTimeoutMs(10);
const executeCalls: unknown[] = [];
const tool: AgentTool = {
name: "sleepy",
@@ -1481,7 +1482,7 @@ describe("ExtensionRunner", () => {
expect(errors).toEqual([]);
});
it("aborts a tool_call handler's confirmation before returning its timeout block", async () => {
it("pauses a tool_call handler timeout during standard and custom dialogs, then resumes its budget", async () => {
const extensionPath = path.join(tempDir.path(), "confirm-tool-call.ts");
const markerPath = path.join(tempDir.path(), "confirm-settled.txt");
fs.writeFileSync(
@@ -1492,8 +1493,11 @@ describe("ExtensionRunner", () => {
export default function(pi) {
pi.on("tool_call", async (_event, ctx) => {
ctx.ui.notify("Waiting for confirmation");
await new Promise(resolve => setTimeout(resolve, 8));
await ctx.ui.confirm("High-risk command", "Allow this command?");
await ctx.ui.custom(() => ({}));
fs.writeFileSync(${JSON.stringify(markerPath)}, "settled");
await Promise.withResolvers().promise;
});
}
`,
@@ -1515,8 +1519,16 @@ describe("ExtensionRunner", () => {
dialogSignal?.addEventListener("abort", () => dialog.resolve(false), { once: true });
return await dialog.promise;
};
const customDialog = Promise.withResolvers<void>();
let customPending = false;
const custom: ExtensionUIContext["custom"] = async <T>() => {
customPending = true;
await customDialog.promise;
return undefined as T;
};
const uiPrototype = Object.create(runner.getUIContext(), {
confirm: { value: confirm },
custom: { value: custom },
notify: { value: notify },
});
const uiContext: ExtensionUIContext = Object.create(uiPrototype);
@@ -1549,25 +1561,159 @@ describe("ExtensionRunner", () => {
undefined,
uiContext,
);
testSetExtensionHandlerTimeoutMs(10);
vi.useFakeTimers();
let now = 0;
const performanceNow = vi.spyOn(performance, "now").mockImplementation(() => now);
try {
testSetExtensionHandlerTimeoutMs(25);
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);
const flush = async () => {
for (let attempts = 0; attempts < 10; attempts++) await Promise.resolve();
};
const execution = wrapped.execute("tool-call-id", {});
await flush();
expect(notify).toHaveBeenCalledWith("Waiting for confirmation");
expect(dialogSignal).toBeUndefined();
now = 8;
vi.advanceTimersByTime(8);
await flush();
expect(dialogSignal).toBeDefined();
now = 108;
vi.advanceTimersByTime(100);
await flush();
expect(dialogSignal?.aborted).toBe(false);
dialog.resolve(true);
await flush();
expect(customPending).toBe(true);
expect(fs.existsSync(markerPath)).toBe(false);
now = 208;
vi.advanceTimersByTime(100);
await flush();
expect(dialogSignal?.aborted).toBe(false);
expect(fs.existsSync(markerPath)).toBe(false);
customDialog.resolve();
await flush();
expect(fs.readFileSync(markerPath, "utf8")).toBe("settled");
now = 225;
vi.advanceTimersByTime(17);
await flush();
vi.advanceTimersByTime(0);
await flush();
await expect(execution).rejects.toThrow(`Extension ${extensionPath} timed out after 25ms`);
} finally {
performanceNow.mockRestore();
vi.useRealTimers();
}
});
it("cancels a pending confirmation and blocks tool execution when the outer dispatch aborts (#4223)", async () => {
const extensionPath = path.join(tempDir.path(), "confirm-abort-tool-call.ts");
const recordPath = path.join(tempDir.path(), "confirm-abort-executed.jsonl");
fs.writeFileSync(
extensionPath,
`
export default function(pi) {
pi.on("tool_call", async (_event, ctx) => {
await ctx.ui.confirm("High-risk command", "Allow this command?");
});
}
`,
);
const result = await loadTestExtensions([extensionPath]);
const runner = new ExtensionRunner(
result.extensions,
result.runtime,
tempDir.path(),
sessionManager,
modelRegistry,
);
let dialogSignal: AbortSignal | undefined;
const dialog = Promise.withResolvers<boolean>();
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 },
});
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,
);
const tool: AgentTool = {
name: "guarded",
label: "Guarded",
description: "must not execute after the extension gate times out",
description: "must not execute after the dispatch aborts",
parameters: Type.Object({}),
strict: true,
execute: async () => ({ content: [{ type: "text", text: "ran" }] }),
execute: async () => {
fs.appendFileSync(recordPath, "ran\n");
return { content: [{ type: "text", text: "ran" }] };
},
};
const wrapped = new ExtensionToolWrapper(tool, runner);
const flush = async () => {
for (let attempts = 0; attempts < 10; attempts++) await Promise.resolve();
};
await expect(wrapped.execute("tool-call-id", {})).rejects.toThrow(
`Extension ${extensionPath} timed out after 10ms`,
);
expect(notify).toHaveBeenCalledWith("Waiting for confirmation");
const controller = new AbortController();
const execution = wrapped.execute("tool-call-id", {} as never, controller.signal);
await flush();
expect(dialogSignal).toBeDefined();
expect(dialogSignal?.aborted).toBe(false);
controller.abort();
await flush();
expect(dialogSignal?.aborted).toBe(true);
expect(fs.readFileSync(markerPath, "utf8")).toBe("settled");
await expect(execution).rejects.toThrow();
expect(fs.existsSync(recordPath)).toBe(false);
});
});