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

@@ -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<void> {
 					}
 				}
 				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<typeof setTimeout> | undefined;
+					if (limitMs > 0) {
+						toolTimer = setTimeout(() => {
+							try { toolAc.abort(); } catch { /* ignore */ }
+						}, limitMs);
+					}
+					let r: Awaited<ReturnType<typeof tool.execute>>;
+					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.

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

src/agent/tools/index.ts

@@ -26,6 +26,8 @@ export {
   setSubagentRunner,
   setQuestionAsker,
   disposeShellSession,
+  toolTimeoutMs,
+  withToolTimeout,
   type TodoItem,
 } from "./shared";
 

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<string>((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<string, number> = {};
   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;

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<string, number> = {
+  // 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<T>(p: Promise<T>, ms: number, label: string): Promise<T> {
+  if (!ms || ms <= 0) return p;
+  return new Promise<T>((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<void> {
-  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<void> {
+  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<void> {
+export function waitForShell(sh: BgShell, ms: number, pattern?: RegExp, signal?: AbortSignal): Promise<void> {
   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);
     };

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