From 3b246ec78703f8443749552af3caabec417846de Mon Sep 17 00:00:00 2001 From: Pawan Osman Date: Wed, 15 Jul 2026 15:19:38 +0300 Subject: [PATCH] Implement per-tool timeout and abort signal handling for agent tools - Introduced a timeout mechanism for various tools (e.g., Grep, Glob, Shell) to prevent blocking the agent loop due to hung processes. - Enhanced the `walk` function to accept an optional AbortSignal, allowing for graceful cancellation of directory traversals. - Updated the `runAgent` function to utilize the new timeout and abort handling, ensuring tools can be aborted without affecting the overall agent execution. - Refactored the `grepTool`, `globTool`, and `fileSearchTool` to incorporate abort signal handling, improving responsiveness during long-running operations. - Added utility functions for managing tool timeouts and signal handling across the agent's tools. --- src/agent/loop.ts | 37 +++++++++++++- src/agent/tools/files.ts | 8 +-- src/agent/tools/index.ts | 2 + src/agent/tools/search.ts | 24 +++++++-- src/agent/tools/shared.ts | 101 ++++++++++++++++++++++++++++++++++++-- src/agent/tools/shell.ts | 37 +++++++++++--- 6 files changed, 186 insertions(+), 23 deletions(-) diff --git a/src/agent/loop.ts b/src/agent/loop.ts index 8aa92f2..3d68c90 100644 --- a/src/agent/loop.ts +++ b/src/agent/loop.ts @@ -9,7 +9,7 @@ import { streamChat, SamplingParams, ModelParams } from "./provider"; import type { OAuthKind } from "./oauth"; -import { TOOLS, schemasForMode, toolsForMode, resetTodos, getTodos, disposeShellSession, EDIT_TOOLS, type AskQuestionItem, type ToolContext } from "./tools"; +import { TOOLS, schemasForMode, toolsForMode, resetTodos, getTodos, disposeShellSession, EDIT_TOOLS, toolTimeoutMs, withToolTimeout, type AskQuestionItem, type ToolContext } from "./tools"; import { actionTypeForCall } from "./approvalPolicy"; import { getWorkspaceRoot } from "../context/workspaceUtils"; import { systemPrompt } from "./prompt"; @@ -612,7 +612,40 @@ export async function runAgent(opts: RunAgentOptions): Promise { } } try { - const r = await tool.execute(input, signal, call.id, toolCtx); + // Per-tool hard timeout + linked abort so hung Shell/Grep/Glob + // cannot block the agent loop. Task/AskQuestion skip the timeout. + const limitMs = toolTimeoutMs(call.name); + const toolAc = new AbortController(); + const onParentAbort = () => { + try { toolAc.abort(); } catch { /* ignore */ } + }; + if (signal.aborted) onParentAbort(); + else signal.addEventListener("abort", onParentAbort, { once: true }); + // Abort the tool slightly before the race rejects so cleanup can run. + let toolTimer: ReturnType | undefined; + if (limitMs > 0) { + toolTimer = setTimeout(() => { + try { toolAc.abort(); } catch { /* ignore */ } + }, limitMs); + } + let r: Awaited>; + try { + r = await withToolTimeout( + tool.execute(input, toolAc.signal, call.id, toolCtx), + limitMs, + call.name, + ); + } catch (e) { + const msg = e instanceof Error ? e.message : String(e); + r = { + output: msg.startsWith("timeout:") + ? `error: ${msg}. The tool was aborted so the agent can continue — retry with a narrower scope or shorter command.` + : `error: ${msg}`, + }; + } finally { + if (toolTimer) clearTimeout(toolTimer); + signal.removeEventListener("abort", onParentAbort); + } const status: "completed" | "error" = r.output.startsWith("error:") ? "error" : "completed"; results[i] = { status, output: r.output, diff: r.diff, startLine: r.startLine, endLine: r.endLine, image: r.image }; // afterEdit hook on successful edits. diff --git a/src/agent/tools/files.ts b/src/agent/tools/files.ts index 0cedce8..60b6167 100644 --- a/src/agent/tools/files.ts +++ b/src/agent/tools/files.ts @@ -93,11 +93,11 @@ export const listDirTool = defineTool("ListDir", false, async (input) => { }); // ---- Glob ---- -export const globTool = defineTool("Glob", false, async (input) => { +export const globTool = defineTool("Glob", false, async (input, abortSignal) => { const root = input.target_directory ? safePath(input.target_directory) : getWorkspaceRoot(); // Include ignored dirs so patterns like "**/node_modules/**" can match. const all: string[] = []; - await walk(root, all, 0, true); + await walk(root, all, 0, true, abortSignal); // Prepend "**/" when the pattern isn't already rooted (schema behavior). let pattern: string = String(input.glob_pattern ?? ""); if (pattern && !pattern.startsWith("**/")) pattern = "**/" + pattern; @@ -111,10 +111,10 @@ export const globTool = defineTool("Glob", false, async (input) => { }); // ---- FileSearch (fuzzy filename search) ---- -export const fileSearchTool = defineTool("FileSearch", false, async (input) => { +export const fileSearchTool = defineTool("FileSearch", false, async (input, abortSignal) => { const root = getWorkspaceRoot(); const all: string[] = []; - await walk(root, all, 0); + await walk(root, all, 0, false, abortSignal); const q = String(input.query || "").toLowerCase(); const rel = all.map((f) => path.relative(root, f).split(path.sep).join("/")); const scored = rel diff --git a/src/agent/tools/index.ts b/src/agent/tools/index.ts index abc5a40..66f1321 100644 --- a/src/agent/tools/index.ts +++ b/src/agent/tools/index.ts @@ -26,6 +26,8 @@ export { setSubagentRunner, setQuestionAsker, disposeShellSession, + toolTimeoutMs, + withToolTimeout, type TodoItem, } from "./shared"; diff --git a/src/agent/tools/search.ts b/src/agent/tools/search.ts index 732d4ab..7b22bc5 100644 --- a/src/agent/tools/search.ts +++ b/src/agent/tools/search.ts @@ -64,14 +64,27 @@ export const grepTool = defineTool("Grep", false, async (input, abortSignal) => args.push("--", input.pattern, target); const out = await new Promise((res) => { - const c = spawn("rg", args, { cwd: root, signal: abortSignal }); + let settled = false; let o = ""; - c.stdout.on("data", (d) => (o += d)); - c.on("error", () => res("(grep failed)")); + const c = spawn("rg", args, { cwd: root, signal: abortSignal }); + const finish = (v: string) => { + if (settled) return; + settled = true; + clearTimeout(timer); + res(v); + }; + // Hard kill hung rg even if AbortSignal is missing/ignored. + const timer = setTimeout(() => { + try { c.kill("SIGTERM"); } catch { /* ignore */ } + finish(o ? o.slice(0, 50_000) + "\n(grep timed out)" : "(grep timed out)"); + }, 25_000); + c.stdout?.on("data", (d) => (o += d)); + c.stderr?.on("data", () => { /* ignore */ }); + c.on("error", () => finish("(grep failed)")); c.on("close", () => { let lines = o.split("\n").filter(Boolean); if (skip) lines = lines.slice(skip); - res(lines.slice(0, cap).join("\n") || "(no matches)"); + finish(lines.slice(0, cap).join("\n") || "(no matches)"); }); }); return { output: out }; @@ -80,7 +93,7 @@ export const grepTool = defineTool("Grep", false, async (input, abortSignal) => // Node fallback (no ripgrep available). Honor path/glob/type/-A/-B/-C/multiline. const scopeRoot = input.path ? safePath(input.path) : root; const all: string[] = []; - await walk(scopeRoot, all, 0); + await walk(scopeRoot, all, 0, false, abortSignal); const flags = input["-i"] ? "i" : ""; const lineRe = new RegExp(input.pattern, flags); @@ -94,6 +107,7 @@ export const grepTool = defineTool("Grep", false, async (input, abortSignal) => const countByFile: Record = {}; const order: string[] = []; for (const f of all.slice(0, 5000)) { + if (abortSignal?.aborted) return { output: "(grep aborted)" }; const rel = path.relative(root, f).split(path.sep).join("/"); if (globRe && !globRe.test(rel)) continue; if (typeExts && !typeExts.includes(path.extname(f).toLowerCase())) continue; diff --git a/src/agent/tools/shared.ts b/src/agent/tools/shared.ts index 9148b59..ce22855 100644 --- a/src/agent/tools/shared.ts +++ b/src/agent/tools/shared.ts @@ -16,6 +16,80 @@ import type { SubagentRunner, QuestionAsker } from "./types"; // Directories never walked/listed. export const IGNORE = new Set([".git", "node_modules", "dist", "out"]); +// --------------------------------------------------------------------------- +// Per-tool hard timeouts (ms). Prevents a hung Grep/Glob/Shell/etc. from +// blocking the agent loop forever. Task/AskQuestion are excluded (own budgets). +// --------------------------------------------------------------------------- +export const TOOL_TIMEOUT_MS: Record = { + // Shell/AwaitShell cap their own foreground wait (90s/120s); the outer + // timeout is a safety net above that so it never fires before the tool's. + Shell: 120_000, + AwaitShell: 150_000, + Grep: 30_000, + Glob: 30_000, + FileSearch: 20_000, + SemanticSearch: 45_000, + SearchDocs: 30_000, + ListDir: 15_000, + Read: 30_000, + ReadLints: 20_000, + WebSearch: 25_000, + WebFetch: 30_000, + StrReplace: 30_000, + Write: 30_000, + Delete: 15_000, + EditNotebook: 30_000, + CallMcpTool: 60_000, + FetchMcpResource: 45_000, + ListMcpResources: 20_000, + TodoWrite: 5_000, + TodoRead: 5_000, + WritePlan: 15_000, + SwitchMode: 5_000, +}; +/** Default when a tool has no explicit entry. */ +export const DEFAULT_TOOL_TIMEOUT_MS = 60_000; +/** Tools that manage their own lifetime (user wait / nested agent). */ +export const NO_TOOL_TIMEOUT = new Set(["Task", "AskQuestion"]); + +/** + * Race a tool promise against a hard timeout. On timeout rejects with an Error + * whose message starts with "timeout:" so the loop can surface it cleanly. + * Does not cancel the underlying work by itself — pass a linked AbortSignal + * into the tool when possible (Shell/Grep honor it). + */ +export function withToolTimeout(p: Promise, ms: number, label: string): Promise { + if (!ms || ms <= 0) return p; + return new Promise((resolve, reject) => { + let settled = false; + const timer = setTimeout(() => { + if (settled) return; + settled = true; + reject(new Error(`timeout: ${label} exceeded ${Math.round(ms / 1000)}s`)); + }, ms); + p.then( + (v) => { + if (settled) return; + settled = true; + clearTimeout(timer); + resolve(v); + }, + (e) => { + if (settled) return; + settled = true; + clearTimeout(timer); + reject(e); + }, + ); + }); +} + +/** Resolve the hard timeout for a tool name (0 = none). */ +export function toolTimeoutMs(name: string): number { + if (NO_TOOL_TIMEOUT.has(name)) return 0; + return TOOL_TIMEOUT_MS[name] ?? DEFAULT_TOOL_TIMEOUT_MS; +} + // Stopwords for the keyword-based SemanticSearch fallback. export const STOP = new Set([ "where", "what", "which", "does", "with", "this", "that", "have", "from", @@ -101,8 +175,17 @@ export function firstDiffLine(before: string, after: string): number { * Recursively collect file paths under `dir` (depth-capped). IGNORE dirs * (.git/node_modules/dist/out) are skipped unless `includeIgnored` is true. */ -export async function walk(dir: string, out: string[], depth: number, includeIgnored = false): Promise { - if (depth > 12) return; +export async function walk( + dir: string, + out: string[], + depth: number, + includeIgnored = false, + signal?: AbortSignal, + /** Soft cap so huge trees cannot hang the tool forever. */ + maxFiles = 50_000, +): Promise { + if (depth > 12 || out.length >= maxFiles) return; + if (signal?.aborted) return; let entries; try { entries = await fs.readdir(dir, { withFileTypes: true }); @@ -110,10 +193,11 @@ export async function walk(dir: string, out: string[], depth: number, includeIgn return; } for (const e of entries) { + if (signal?.aborted || out.length >= maxFiles) return; if (!includeIgnored && IGNORE.has(e.name)) continue; const full = path.join(dir, e.name); if (e.isDirectory()) { - await walk(full, out, depth + 1, includeIgnored); + await walk(full, out, depth + 1, includeIgnored, signal, maxFiles); } else { out.push(full); } @@ -273,11 +357,18 @@ export function disposeShellSession(key: string): void { } /** Wait until the shell finishes, `pattern` matches its output, or `ms` elapses. */ -export function waitForShell(sh: BgShell, ms: number, pattern?: RegExp): Promise { +export function waitForShell(sh: BgShell, ms: number, pattern?: RegExp, signal?: AbortSignal): Promise { return new Promise((resolve) => { const deadline = Date.now() + Math.max(0, ms); const tick = () => { - sh.pump?.(); + try { sh.pump?.(); } catch { /* ignore */ } + if (signal?.aborted) { + if (!sh.done) { + sh.output += "\n(aborted)"; + sh.done = true; + } + return resolve(); + } if (sh.done || (pattern && pattern.test(sh.output)) || Date.now() >= deadline) return resolve(); setTimeout(tick, 100); }; diff --git a/src/agent/tools/shell.ts b/src/agent/tools/shell.ts index 8e72ed2..85c2015 100644 --- a/src/agent/tools/shell.ts +++ b/src/agent/tools/shell.ts @@ -47,7 +47,10 @@ function buildNotify(input: any, ctx: any): ShellNotify | undefined { // by a sentinel echo so we can detect completion and capture the exit code. export const runTerminalTool = defineTool("Shell", true, async (input, abortSignal, _callId, ctx) => { const root = getWorkspaceRoot(); - const blockMs = typeof input.block_until_ms === "number" ? input.block_until_ms : 30_000; + // Cap foreground wait: default 30s, hard max 90s so a hung command cannot + // block the agent loop. block_until_ms=0 still backgrounds immediately. + const rawBlock = typeof input.block_until_ms === "number" ? input.block_until_ms : 30_000; + const blockMs = rawBlock <= 0 ? 0 : Math.min(Math.max(0, rawBlock), 90_000); const command = String(input.command ?? ""); // Prune finished shells older than 10 minutes to bound the registry. @@ -92,14 +95,19 @@ export const runTerminalTool = defineTool("Shell", true, async (input, abortSign }; const pumpTimer = setInterval(() => sh.pump?.(), 100); - const onAbort = () => { - sh.output += "\n(aborted)"; - sh.done = true; + const killSession = () => { try { if (process.platform === "win32") session.proc.kill(); else session.proc.kill("SIGTERM"); } catch { /* ignore */ } }; + const onAbort = () => { + if (!sh.done) { + sh.output += "\n(aborted)"; + sh.done = true; + } + killSession(); + }; abortSignal?.addEventListener("abort", onAbort); // Optional per-command working directory (a `cd` that does not persist). @@ -107,13 +115,26 @@ export const runTerminalTool = defineTool("Shell", true, async (input, abortSign const wrapped = isWin ? `${cd ? `Push-Location -LiteralPath '${cd.replace(/'/g, "''")}'; ` : ""}${command}${cd ? "; Pop-Location" : ""}\nWrite-Output "${SENTINEL}:$LASTEXITCODE"\n` : `${cd ? `pushd '${cd.replace(/'/g, "'\\''")}' && ` : ""}${command}\n__oc_rc=$?\n${cd ? "popd >/dev/null 2>&1\n" : ""}echo "${SENTINEL}:$__oc_rc"\n`; - session.proc.stdin?.write(wrapped); + try { + session.proc.stdin?.write(wrapped); + } catch (e) { + clearInterval(pumpTimer); + abortSignal?.removeEventListener("abort", onAbort); + sh.done = true; + return { output: `error: shell write failed: ${e instanceof Error ? e.message : String(e)}` }; + } - await waitForShell(sh, blockMs); + await waitForShell(sh, blockMs, undefined, abortSignal); sh.pump?.(); clearInterval(pumpTimer); abortSignal?.removeEventListener("abort", onAbort); + // Hard timeout with no sentinel: mark done so AwaitShell does not hang, and + // note that the process may still be running in the session. + if (!sh.done && blockMs > 0) { + sh.output += `\n(timeout after ${blockMs}ms — command still running in background shell ${sh.id}; use AwaitShell or cancel)`; + } + // Strip the sentinel line from the rendered body. sh.output = sh.output.replace(new RegExp("\\n?" + SENTINEL + ":-?\\d+\\s*"), ""); return { output: renderShell(sh) }; @@ -151,6 +172,8 @@ export const awaitShellTool = defineTool("AwaitShell", false, async (input) => { return { output: `error: invalid pattern: ${e instanceof Error ? e.message : String(e)}` }; } } - await waitForShell(sh, blockMs, pattern); + // Cap AwaitShell polls too so a stuck bg job cannot hang the loop forever. + const capped = blockMs <= 0 ? 0 : Math.min(blockMs, 120_000); + await waitForShell(sh, capped, pattern); return { output: renderShell(sh) }; });