diff --git a/src/agent/loop.ts b/src/agent/loop.ts index 4a9795a..099e856 100644 --- a/src/agent/loop.ts +++ b/src/agent/loop.ts @@ -652,7 +652,10 @@ export async function runAgent(opts: RunAgentOptions): Promise { // with {} would call tools with missing params — fail the call instead. badArgs = true; } - const tMs = toolTimeoutMs(call.name); + // MCP tools share CallMcpTool budget when no per-name override. + const tMs = call.name.startsWith("mcp__") + ? toolTimeoutMs("CallMcpTool") + : toolTimeoutMs(call.name); // Shell: countdown uses block_until_ms when shorter than tool budget. let timeoutMs = tMs > 0 ? tMs : undefined; if ((call.name === "Shell" || call.name === "AwaitShell") && !badArgs) { @@ -711,7 +714,7 @@ export async function runAgent(opts: RunAgentOptions): Promise { finishUi(i); return; } - // MCP tool dispatch. + // MCP tool dispatch (same hard timeout + countdown as built-ins). if (call.name.startsWith("mcp__")) { if (!isAgentic()) { // MCP tools may mutate; only allow in agentic modes. @@ -732,8 +735,49 @@ export async function runAgent(opts: RunAgentOptions): Promise { results[i] = { status: "error", output: `blocked by hook: ${mcpVeto}` }; return; } - const out = await mcpManager.callTool(call.name, input); - results[i] = { status: out.startsWith("error:") ? "error" : "completed", output: out }; + const limitMs = parsed[i].timeoutMs ?? toolTimeoutMs("CallMcpTool"); + const toolAc = new AbortController(); + const killTool = () => { try { toolAc.abort(); } catch { /* ignore */ } }; + if (registerSubagentAbort) registerSubagentAbort(call.id, killTool); + emit({ + type: "tool-call-started", + callId: call.id, + name: call.name, + input, + timeoutMs: limitMs > 0 ? limitMs : undefined, + startedAt: Date.now(), + }); + const onParentAbort = () => killTool(); + if (signal.aborted) onParentAbort(); + else signal.addEventListener("abort", onParentAbort, { once: true }); + try { + const out = await withToolTimeout( + Promise.resolve().then(() => mcpManager.callTool(call.name, input)), + limitMs, + call.name, + () => { + killTool(); + results[i] = { + status: "error", + output: `error: timeout: ${call.name} exceeded ${Math.round((limitMs || 0) / 1000)}s. Tool aborted.`, + }; + finishUi(i); + }, + ); + results[i] = { status: out.startsWith("error:") ? "error" : "completed", output: out }; + finishUi(i); + } catch (e) { + const msg = e instanceof Error ? e.message : String(e); + results[i] = { + status: "error", + output: msg.startsWith("timeout:") + ? `error: ${msg}. Tool aborted.` + : `error: ${msg}`, + }; + finishUi(i); + } finally { + signal.removeEventListener("abort", onParentAbort); + } return; } const tool = TOOLS[call.name]; diff --git a/src/shared/turns.ts b/src/shared/turns.ts index 3321530..f17a229 100644 --- a/src/shared/turns.ts +++ b/src/shared/turns.ts @@ -306,11 +306,13 @@ export function applyEvent(turns: Turn[], ev: AgentEvent): Turn[] { const startedAt = ev.startedAt; if (existing >= 0) { const prev = turn.blocks[existing] as ToolBlock; + // Never reopen a settled tool (timeout/cancel may race a late start). + if (prev.status !== "running") return list; turn.blocks[existing] = { ...prev, name: ev.name, input: Object.keys(ev.input || {}).length ? ev.input : prev.input, - status: prev.status === "running" ? "running" : prev.status, + status: "running", timeoutMs: timeoutMs ?? prev.timeoutMs, startedAt: startedAt ?? prev.startedAt, } as AssistantBlock; diff --git a/webview-ui/sidebar/App.tsx b/webview-ui/sidebar/App.tsx index 55a5c2c..434150d 100644 --- a/webview-ui/sidebar/App.tsx +++ b/webview-ui/sidebar/App.tsx @@ -12,7 +12,7 @@ import { Icon } from "../shared/icons"; import { renderMarkdown } from "../shared/markdown"; import { vscode } from "../shared/vscode"; import { Composer, KIND_SVG, applyFileIconTo } from "./components/Composer"; -import { ToolCard, isReadonlySubagent } from "./components/Tool"; +import { ToolCard, isReadonlySubagent, TimeoutBadge, ToolTimeoutWatch, isToolCountdownActive } from "./components/Tool"; import { History } from "./components/History"; import type { AgentEvent, ApprovalMode, ApprovalRequestInfo, AssistantBlock, AssistantTurn, Attachment, ConversationSummary, ErrorBlock, InMessage, MentionItem, Mode, ModelDef, ModelOption, OutMessage, PendingChangeInfo, PersonaInfo, ThinkingBlock, ToolBlock, Turn, UserTurn } from "./types"; import { applyEvent, applyToBlocks, closeTrailingThinking, forceSettleOpenWork, parsePartialArgs, renderMentionTokens } from "./types"; @@ -213,15 +213,23 @@ function ExploringSection({ const running = tools.some((t) => t.status === "running") || !!live; const current = [...tools].reverse().find((t) => t.status === "running") ?? tools[tools.length - 1]; const subtitle = running ? capitalize(toolLabel(current.name)) : exploreSummary(tools); + // Keep kill-at-zero active even when the group is collapsed (no ToolCard mount). + const timed = tools.filter((t) => isToolCountdownActive(t) && t.timeoutMs && t.timeoutMs > 0); + const headTimed = timed[0] ?? (current && isToolCountdownActive(current) ? current : null); return (
+ {/* Always watch every timed tool so countdown-0 kills even when collapsed. */} + {timed.map((t) => ( + + ))}
setOpen((o) => !o)}> {running ? "Exploring" : exploreSummary(tools)} + {headTimed ? : null} {running ? : {tools.length}}
{!open && running &&
{subtitle}
} @@ -425,8 +433,9 @@ function SubagentChat({ block, onBack }: { block: import("./types").ToolBlock; o Back to chat {isReadonlySubagent(block.input) ? "read-only" : "agent"} + {running && ( - )} diff --git a/webview-ui/sidebar/components/Tool.tsx b/webview-ui/sidebar/components/Tool.tsx index b7fcab7..ec174ef 100644 --- a/webview-ui/sidebar/components/Tool.tsx +++ b/webview-ui/sidebar/components/Tool.tsx @@ -183,6 +183,7 @@ function TodoList({ block }: { block: ToolBlock }) { Todos +
@@ -291,6 +292,7 @@ function PlanCard({ block, onImplement }: { block: ToolBlock; onImplement?: (pat {title} + Plan @@ -324,41 +326,64 @@ function StatusIcon({ status }: { status: ToolBlock["status"] }) { return status === "completed" ? : ; } -/** Live countdown for a running tool. At 0: cancel/kill via host. */ -function useToolCountdown(block: ToolBlock): number | null { - const budget = block.timeoutMs && block.timeoutMs > 0 ? block.timeoutMs : 0; - const started = block.startedAt || 0; +/** True while this tool/task should show a live timeout countdown. */ +export function isToolCountdownActive(block: ToolBlock): boolean { + if (block.name === "AskQuestion" || block.name === "ask_question") return false; const isTask = block.name === "Task" || block.name === "task"; - const subDone = block.subStatus === "finished" || block.subStatus === "cancelled" || block.subStatus === "error"; - // Task: keep countdown while nested work is open (incl. bg after parent tool completed). - const running = isTask - ? !subDone && (block.status === "running" || block.subStatus === "running" || (!!block.subBlocks?.length && !subDone)) - : block.status === "running"; - const [left, setLeft] = React.useState(() => { - if (!budget || !started || !running) return null; - return Math.max(0, Math.ceil((started + budget - Date.now()) / 1000)); - }); - const fired = React.useRef(false); + if (isTask) { + const subDone = block.subStatus === "finished" || block.subStatus === "cancelled" || block.subStatus === "error"; + if (subDone) return false; + // Parent tool may complete early (bg Task); keep counting while nested work open. + return block.status === "running" || block.subStatus === "running" || !!block.subBlocks?.length; + } + return block.status === "running"; +} + +/** Dedup kill-at-zero across multiple countdown mounts (explore head + card). */ +const firedTimeouts = new Set(); + +/** + * Live countdown for a running tool/task. At 0: cancel/kill via host. + * Uses host `startedAt` when present; otherwise starts the clock on first + * observation so every timed tool always shows a countdown. + */ +export function useToolCountdown(block: ToolBlock): number | null { + const budget = block.timeoutMs && block.timeoutMs > 0 ? block.timeoutMs : 0; + const hostStarted = block.startedAt && block.startedAt > 0 ? block.startedAt : 0; + const running = isToolCountdownActive(block); + const localStart = React.useRef(0); + const [left, setLeft] = React.useState(null); + React.useEffect(() => { - fired.current = false; - // Wait until execute has stamped startedAt (budget alone is not enough). - if (!budget || !started || !running) { + if (!running || !budget || !block.callId) { + localStart.current = 0; setLeft(null); return; } + // Prefer host clock; fall back to first UI observation of this run. + if (hostStarted) localStart.current = hostStarted; + else if (!localStart.current) localStart.current = Date.now(); + // Allow re-fire only on a brand-new callId (set already cleared on settle). + if (block.status !== "running" && !(block.name === "Task" || block.name === "task")) { + firedTimeouts.delete(block.callId); + } + const tick = () => { - const sec = Math.max(0, Math.ceil((started + budget - Date.now()) / 1000)); + const start = hostStarted || localStart.current; + if (!start) return; + const sec = Math.max(0, Math.ceil((start + budget - Date.now()) / 1000)); setLeft(sec); - if (sec <= 0 && !fired.current) { - fired.current = true; + if (sec <= 0 && !firedTimeouts.has(block.callId)) { + firedTimeouts.add(block.callId); post({ type: "cancelSubagent", callId: block.callId, reason: "timeout" }); } }; tick(); const id = window.setInterval(tick, 250); return () => window.clearInterval(id); - }, [budget, started, running, block.callId, block.status, block.subStatus]); - if (!running || left == null) return null; + }, [budget, hostStarted, running, block.callId, block.status, block.subStatus, block.timeoutMs, block.name]); + + if (!running || !budget || left == null) return null; return left; } @@ -371,17 +396,27 @@ function formatCountdown(sec: number): string { return `${sec}s`; } -function TimeoutBadge({ block }: { block: ToolBlock }) { +/** Visible countdown badge; also drives kill-at-zero via useToolCountdown. */ +export function TimeoutBadge({ block }: { block: ToolBlock }) { const left = useToolCountdown(block); if (left == null) return null; const urgent = left <= 5; return ( - + {left === 0 ? "timeout" : formatCountdown(left)} ); } +/** Silent countdown (kill at 0) without rendering — for collapsed explore groups. */ +export function ToolTimeoutWatch({ block }: { block: ToolBlock }) { + useToolCountdown(block); + return null; +} + function Diff({ diff }: { diff: string }) { const lines = diff.split("\n"); const needsExpand = lines.length > 6;