zedcode/PawanOsman-ChatGPT
1
0
Fork 0
mirror of https://github.com/PawanOsman/ChatGPT.git synced 2026-07-18 08:05:57 +02:00
Code Issues Packages Projects Releases Wiki Activity

Enhance tool timeout management and UI integration

Browse Source 
- Updated the `runAgent` function to improve timeout handling for MCP tools, ensuring shared budget management and better execution limits.
- Introduced a countdown mechanism for tool calls, allowing for real-time UI feedback on execution status and tool cancellations.
- Enhanced the UI components to display timeout badges and countdowns for running tools, improving user awareness and responsiveness.
- Implemented a silent countdown feature for collapsed tool groups, ensuring timely cancellations without UI clutter.
- Refactored the `Tool` component to integrate timeout badges, providing clearer visual indicators of tool status.
This commit is contained in:
Pawan Osman
2026-07-15 18:04:43 +03:00
parent 6b8301a28e
commit b63ac10f8f
4 changed files with 121 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent/loop.ts
+48 -4
View File
@@ -652,7 +652,10 @@ export async function runAgent(opts: RunAgentOptions): Promise<void> {
// 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<void> {
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<void> {
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];
src/shared/turns.ts
+3 -1
View File
@@ -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;
webview-ui/sidebar/App.tsx
+11 -2
View File
@@ -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 (
<div className={"explore-section" + (open ? " open" : "")}>
{/* Always watch every timed tool so countdown-0 kills even when collapsed. */}
{timed.map((t) => (
<ToolTimeoutWatch key={`watch-${t.callId}`} block={t} />
))}
<div className="explore-head" onClick={() => setOpen((o) => !o)}>
<span className={"tchev" + (open ? " open" : "")}>
<Icon name="chevD" size={12} />
</span>
<Icon name="search" size={12} className="explore-icon" />
<span className="explore-title">{running ? "Exploring" : exploreSummary(tools)}</span>
{headTimed ? <TimeoutBadge block={headTimed} /> : null}
{running ? <span className="spinner" /> : <span className="explore-count">{tools.length}</span>}
</div>
{!open && running && <div className="explore-subtitle">{subtitle}</div>}
@@ -425,8 +433,9 @@ function SubagentChat({ block, onBack }: { block: import("./types").ToolBlock; o
<Icon name="chevD" size={12} /> Back to chat
</button>
<span className="sub-readonly">{isReadonlySubagent(block.input) ? "read-only" : "agent"}</span>
<TimeoutBadge block={block} />
{running && (
<button className="sub-stop" onClick={() => post({ type: "cancelSubagent", callId: block.callId })}>
<button className="sub-stop" onClick={() => post({ type: "cancelSubagent", callId: block.callId, reason: "user" })}>
<Icon name="close" size={12} /> Stop
</button>
)}
webview-ui/sidebar/components/Tool.tsx
+59 -24
View File
@@ -183,6 +183,7 @@ function TodoList({ block }: { block: ToolBlock }) {
</span>
<span className="label">Todos</span>
<span className="right">
<TimeoutBadge block={block} />
<StatusIcon status={block.status} />
</span>
</div>
@@ -291,6 +292,7 @@ function PlanCard({ block, onImplement }: { block: ToolBlock; onImplement?: (pat
<Icon name="todo" />
</span>
<span className="plan-title">{title}</span>
<TimeoutBadge block={block} />
<span className="badge badge-plan">Plan</span>
<StatusIcon status={block.status} />
</div>
@@ -324,41 +326,64 @@ function StatusIcon({ status }: { status: ToolBlock["status"] }) {
return status === "completed" ? <Icon name="check" className="ok-icon" /> : <Icon name="close" className="err-icon" />;
}
/** 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<number | null>(() => {
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<string>();
/**
* 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<number | null>(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 (
<span className={"tool-timeout" + (urgent ? " urgent" : "") + (left === 0 ? " zero" : "")} title="Timeout remaining">
<span
className={"tool-timeout" + (urgent ? " urgent" : "") + (left === 0 ? " zero" : "")}
title="Timeout remaining — tool is killed at 0"
>
{left === 0 ? "timeout" : formatCountdown(left)}
</span>
);
}
/** 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;