From 669120a1fe2a9d2ad09f07db9fd7924d1116f406 Mon Sep 17 00:00:00 2001 From: Savas Date: Wed, 9 Sep 2026 16:57:39 -0300 Subject: [PATCH] fix(agent): show real token counts in subagent progress Streaming providers (Verboo/OpenAI shim) emit message_start with zeroed usage; the real usage arrives in message_delta, which mutates the last yielded message in place after the consumer already read it. Subagent progress therefore counted tool calls but kept tokens at 0. - Add syncProgressUsageFromMessages() to re-derive token counters from accumulated assistant messages (input cumulative, output summed) - Apply it in runAsyncAgentLifecycle, sync and backgrounded AgentTool paths, both per-message and after the loop ends - Show tokens/tools in the background tasks list (local_agent rows) - Log when a model override is ignored in coordinator mode instead of silently swallowing it - Fix indentation in the backgrounded error block - Sanitize description angle brackets in task notification summary - Add regression test for token convergence after in-place usage mutation --- src/components/tasks/BackgroundTask.tsx | 46 +++++-- src/tasks/LocalAgentTask/LocalAgentTask.tsx | 42 +++++- .../LocalAgentTask/tokenProgress.test.ts | 120 ++++++++++++++++++ src/tools/AgentTool/AgentTool.tsx | 18 ++- src/tools/AgentTool/agentToolUtils.ts | 9 ++ 5 files changed, 219 insertions(+), 16 deletions(-) create mode 100644 src/tasks/LocalAgentTask/tokenProgress.test.ts diff --git a/src/components/tasks/BackgroundTask.tsx b/src/components/tasks/BackgroundTask.tsx index 7402270d75..7cb79195a0 100644 --- a/src/components/tasks/BackgroundTask.tsx +++ b/src/components/tasks/BackgroundTask.tsx @@ -3,7 +3,7 @@ import * as React from 'react'; import { Text } from 'src/ink.js'; import type { BackgroundTaskState } from 'src/tasks/types.js'; import type { DeepImmutable } from 'src/types/utils.js'; -import { truncate } from 'src/utils/format.js'; +import { truncate, formatNumber } from 'src/utils/format.js'; import { toInkColor } from 'src/utils/ink.js'; import { plural } from 'src/utils/stringUtils.js'; import { DIAMOND_FILLED, DIAMOND_OPEN } from '../../constants/figures.js'; @@ -123,26 +123,46 @@ export function BackgroundTask(t0) { } else { t1 = $[24]; } + const tokenCount = "result" in task ? task.result?.totalTokens : task.progress?.tokenCount; + const toolUseCount = "result" in task ? task.result?.totalToolUseCount : task.progress?.toolUseCount; + let tToken; + if ($[25] !== tokenCount) { + tToken = tokenCount !== undefined && tokenCount > 0 && · {formatNumber(tokenCount)} tokens; + $[25] = tokenCount; + $[26] = tToken; + } else { + tToken = $[26]; + } + let tTool; + if ($[27] !== toolUseCount) { + tTool = toolUseCount !== undefined && toolUseCount > 0 && · {toolUseCount} {toolUseCount === 1 ? "tool" : "tools"}; + $[27] = toolUseCount; + $[28] = tTool; + } else { + tTool = $[28]; + } const t2 = task.status === "completed" ? "done" : undefined; const t3 = task.status === "completed" && !task.notified ? ", unread" : undefined; let t4; - if ($[25] !== t2 || $[26] !== t3 || $[27] !== task.status) { + if ($[29] !== t2 || $[30] !== t3 || $[31] !== task.status) { t4 = ; - $[25] = t2; - $[26] = t3; - $[27] = task.status; - $[28] = t4; + $[29] = t2; + $[30] = t3; + $[31] = task.status; + $[32] = t4; } else { - t4 = $[28]; + t4 = $[32]; } let t5; - if ($[29] !== t1 || $[30] !== t4) { - t5 = {t1}{" "}{t4}; - $[29] = t1; - $[30] = t4; - $[31] = t5; + if ($[33] !== t1 || $[34] !== t4 || $[35] !== tToken || $[36] !== tTool) { + t5 = {t1}{" "}{tToken}{tTool}{" "}{t4}; + $[33] = t1; + $[34] = t4; + $[35] = tToken; + $[36] = tTool; + $[37] = t5; } else { - t5 = $[31]; + t5 = $[37]; } return t5; } diff --git a/src/tasks/LocalAgentTask/LocalAgentTask.tsx b/src/tasks/LocalAgentTask/LocalAgentTask.tsx index a09e8168ad..39c20f6c63 100644 --- a/src/tasks/LocalAgentTask/LocalAgentTask.tsx +++ b/src/tasks/LocalAgentTask/LocalAgentTask.tsx @@ -59,6 +59,43 @@ export function getTokenCountFromTracker(tracker: ProgressTracker): number { return tracker.latestInputTokens + tracker.cumulativeOutputTokens; } +/** + * Recompute token counters from accumulated assistant messages. + * + * Streaming providers (Verboo/OpenAI shim) emit message_start with zeroed + * usage; the real usage arrives in message_delta, which mutates the LAST + * yielded message in place (see claude.ts message_delta handler). Reading + * usage at yield time therefore undercounts tokens — tool calls count but + * tokens stay 0. Because lifecycle loops keep references to every yielded + * message, this recompute converges to the real usage once each request's + * message_delta lands on the accumulated references. + * + * Input tokens are cumulative per request (latest wins); output tokens are + * per-request and summed. + */ +export function syncProgressUsageFromMessages( + tracker: ProgressTracker, + messages: Message[], +): void { + let latestInputTokens = 0; + let cumulativeOutputTokens = 0; + for (const message of messages) { + if (message.type !== 'assistant') continue; + const usage = message.message.usage; + if (!usage) continue; + const inputTokens = + usage.input_tokens + + (usage.cache_creation_input_tokens ?? 0) + + (usage.cache_read_input_tokens ?? 0); + if (inputTokens > latestInputTokens) { + latestInputTokens = inputTokens; + } + cumulativeOutputTokens += usage.output_tokens; + } + tracker.latestInputTokens = latestInputTokens; + tracker.cumulativeOutputTokens = cumulativeOutputTokens; +} + /** * Resolver function that returns a human-readable activity description * for a given tool name and input. Used to pre-compute descriptions @@ -245,7 +282,10 @@ export function enqueueAgentNotification({ // results may reference stale task output. The prompt suggestion text is // preserved; only the pre-computed response is discarded. abortSpeculation(setAppState); - const summary = status === 'completed' ? `Agent "${description}" completed` : status === 'failed' ? `Agent "${description}" failed: ${error || 'Unknown error'}` : `Agent "${description}" was stopped`; + // Strip angle brackets from user/model-supplied description so the summary + // never breaks the envelope (BUG-4: parser regex truncation / malformed XML). + const safeDescription = description.replace(/[<>]/g, ''); + const summary = status === 'completed' ? `Agent "${safeDescription}" completed` : status === 'failed' ? `Agent "${safeDescription}" failed: ${error || 'Unknown error'}` : `Agent "${safeDescription}" was stopped`; const outputPath = getTaskOutputPath(taskId); const toolUseIdLine = toolUseId ? `\n<${TOOL_USE_ID_TAG}>${toolUseId}` : ''; const resultSection = finalMessage ? `\n${finalMessage}` : ''; diff --git a/src/tasks/LocalAgentTask/tokenProgress.test.ts b/src/tasks/LocalAgentTask/tokenProgress.test.ts new file mode 100644 index 0000000000..4db3d3d0c0 --- /dev/null +++ b/src/tasks/LocalAgentTask/tokenProgress.test.ts @@ -0,0 +1,120 @@ +import { expect, test } from 'bun:test' +import { + createProgressTracker, + getProgressUpdate, + syncProgressUsageFromMessages, + updateProgressFromMessage, +} from './LocalAgentTask.js' +import { createAssistantMessage } from '../../utils/messages.js' + +// BUG-1 regression: streaming (Verboo/OpenAI shim) yields assistant messages +// with zeroed usage; the real usage arrives in message_delta which mutates the +// yielded message in place AFTER the consumer already read it. Tool calls +// counted, tokens stuck at 0. syncProgressUsageFromMessages re-derives tokens +// from the accumulated references (which message_delta mutated). +test('syncProgressUsageFromMessages converges tokens after in-place usage mutation', () => { + const tracker = createProgressTracker() + const agentMessages: Parameters[1] = [] + + // Request 1: text + Bash tool_use, both yielded with usage 0 + const m1 = createAssistantMessage({ + content: 'ok', + usage: { + input_tokens: 0, + output_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + server_tool_use: { web_search_requests: 0, web_fetch_requests: 0 }, + service_tier: null, + cache_creation: { ephemeral_1h_input_tokens: 0, ephemeral_5m_input_tokens: 0 }, + inference_geo: null, + iterations: null, + speed: null, + }, + }) + const m2 = createAssistantMessage({ + content: [ + { + type: 'tool_use', + id: 'tu1', + name: 'Bash', + input: { command: 'ls' }, + }, + ], + usage: { + input_tokens: 0, + output_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + server_tool_use: { web_search_requests: 0, web_fetch_requests: 0 }, + service_tier: null, + cache_creation: { ephemeral_1h_input_tokens: 0, ephemeral_5m_input_tokens: 0 }, + inference_geo: null, + iterations: null, + speed: null, + }, + }) + + // Consumer processes m1 then m2 as they arrive (usage still 0) + agentMessages.push(m1, m2) + for (const m of agentMessages) updateProgressFromMessage(tracker, m) + syncProgressUsageFromMessages(tracker, agentMessages) + expect(tracker.toolUseCount).toBe(1) + expect(getProgressUpdate(tracker).tokenCount).toBe(0) + + // message_delta mutates ONLY the last message in place with the real usage + ;(m2.message as { usage: unknown }).usage = { + input_tokens: 1500, + output_tokens: 40, + cache_creation_input_tokens: 5000, + cache_read_input_tokens: 0, + server_tool_use: { web_search_requests: 0, web_fetch_requests: 0 }, + service_tier: null, + cache_creation: { ephemeral_1h_input_tokens: 0, ephemeral_5m_input_tokens: 0 }, + inference_geo: null, + iterations: null, + speed: null, + } + + // Next message arrives (request 2) → recompute picks up the mutation + const m3 = createAssistantMessage({ + content: [ + { type: 'tool_use', id: 'tu2', name: 'Read', input: { file_path: 'x' } }, + ], + usage: { + input_tokens: 0, + output_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + server_tool_use: { web_search_requests: 0, web_fetch_requests: 0 }, + service_tier: null, + cache_creation: { ephemeral_1h_input_tokens: 0, ephemeral_5m_input_tokens: 0 }, + inference_geo: null, + iterations: null, + speed: null, + }, + }) + agentMessages.push(m3) + updateProgressFromMessage(tracker, m3) + syncProgressUsageFromMessages(tracker, agentMessages) + expect(tracker.toolUseCount).toBe(2) + // input 1500 + cache 5000 (request 1, accumulated) + output 40; request 2 still pending + expect(getProgressUpdate(tracker).tokenCount).toBe(6540) + + // Request 2's message_delta lands (mutates m3), agent ends → final recompute + ;(m3.message as { usage: unknown }).usage = { + input_tokens: 1700, + output_tokens: 25, + cache_creation_input_tokens: 5000, + cache_read_input_tokens: 0, + server_tool_use: { web_search_requests: 0, web_fetch_requests: 0 }, + service_tier: null, + cache_creation: { ephemeral_1h_input_tokens: 0, ephemeral_5m_input_tokens: 0 }, + inference_geo: null, + iterations: null, + speed: null, + } + syncProgressUsageFromMessages(tracker, agentMessages) + // latest input 1700 + cache 5000, outputs 40 + 25 = 65 + expect(getProgressUpdate(tracker).tokenCount).toBe(6765) +}) \ No newline at end of file diff --git a/src/tools/AgentTool/AgentTool.tsx b/src/tools/AgentTool/AgentTool.tsx index 7f5e6b74f9..34dacd5305 100644 --- a/src/tools/AgentTool/AgentTool.tsx +++ b/src/tools/AgentTool/AgentTool.tsx @@ -11,7 +11,7 @@ import { startAgentSummarization } from '../../services/AgentSummary/agentSummar import { getFeatureValue_CACHED_MAY_BE_STALE } from '../../services/analytics/growthbook.js'; import { type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS, logEvent } from '../../services/analytics/index.js'; import { clearDumpState } from '../../services/api/dumpPrompts.js'; -import { completeAgentTask as completeAsyncAgent, createActivityDescriptionResolver, createProgressTracker, enqueueAgentNotification, failAgentTask as failAsyncAgent, getProgressUpdate, getTokenCountFromTracker, isLocalAgentTask, killAsyncAgent, registerAgentForeground, registerAsyncAgent, unregisterAgentForeground, updateAgentProgress as updateAsyncAgentProgress, updateProgressFromMessage } from '../../tasks/LocalAgentTask/LocalAgentTask.js'; +import { completeAgentTask as completeAsyncAgent, createActivityDescriptionResolver, createProgressTracker, enqueueAgentNotification, failAgentTask as failAsyncAgent, getProgressUpdate, getTokenCountFromTracker, isLocalAgentTask, killAsyncAgent, registerAgentForeground, registerAsyncAgent, syncProgressUsageFromMessages, unregisterAgentForeground, updateAgentProgress as updateAsyncAgentProgress, updateProgressFromMessage } from '../../tasks/LocalAgentTask/LocalAgentTask.js'; import { checkRemoteAgentEligibility, formatPreconditionError, getRemoteTaskSessionUrl, registerRemoteAgentTask } from '../../tasks/RemoteAgentTask/RemoteAgentTask.js'; import { assembleToolPool } from '../../tools.js'; import { asAgentId } from '../../types/ids.js'; @@ -252,6 +252,13 @@ export const AgentTool = buildTool({ cwd }: AgentToolInput, toolUseContext, canUseTool, assistantMessage, onProgress?) { const startTime = Date.now(); + // Coordinator workers must use the default model — the coordinator system + // prompt instructs this and model routing assumes it. Silently swallowing + // the param hides mistakes (the model may believe it selected a model). + // Log so the mismatch is visible instead of silent. + if (isCoordinatorMode() && modelParam !== undefined) { + logForDebugging(`[Coordinator] Agent model override ignored (coordinator mode): "${modelParam}"`); + } const model = isCoordinatorMode() ? undefined : modelParam; // Get app state for permission mode and agent filtering @@ -1003,6 +1010,7 @@ export const AgentTool = buildTool({ // Track progress for backgrounded agents updateProgressFromMessage(tracker, msg, resolveActivity2, toolUseContext.options.tools); + syncProgressUsageFromMessages(tracker, agentMessages); updateAsyncAgentProgress(backgroundedTaskId, getProgressUpdate(tracker), rootSetAppState); const lastToolName = getLastToolUseName(msg); if (lastToolName) { @@ -1011,6 +1019,7 @@ export const AgentTool = buildTool({ } if (!isCurrentBackground()) return; if (backgroundController.signal.aborted) throw new AbortError(); + syncProgressUsageFromMessages(tracker, agentMessages); const agentResult = finalizeAgentTool(agentMessages, backgroundedTaskId, metadata); // Mark task completed FIRST so TaskOutput(block=true) @@ -1068,7 +1077,7 @@ export const AgentTool = buildTool({ reason: 'user_cancel_background' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS }); const worktreeResult = await cleanupWorktreeIfNeeded(); - if (!isCurrentBackground()) return; + if (!isCurrentBackground()) return; const partialResult = extractPartialResult(agentMessages); enqueueAgentNotification({ taskId: backgroundedTaskId, @@ -1138,6 +1147,7 @@ export const AgentTool = buildTool({ // Emit task_progress for the VS Code subagent panel updateProgressFromMessage(syncTracker, message, syncResolveActivity, toolUseContext.options.tools); + syncProgressUsageFromMessages(syncTracker, agentMessages); if (foregroundTaskId) { const lastToolName = getLastToolUseName(message); if (lastToolName) { @@ -1232,6 +1242,10 @@ export const AgentTool = buildTool({ // closure owns a separate stop function (stopBackgroundedSummarization). stopForegroundSummarization?.(); + // BUG-1 fix: last request's message_delta lands after the final yield; + // recompute so progress/SDK totals carry real token counts. + syncProgressUsageFromMessages(syncTracker, agentMessages); + // Unregister foreground task if agent completed without being backgrounded if (foregroundTaskId && !wasReplaced) { unregisterAgentForeground(foregroundTaskId, rootSetAppState, foregroundExecutionId); diff --git a/src/tools/AgentTool/agentToolUtils.ts b/src/tools/AgentTool/agentToolUtils.ts index bfddbde60f..221a6bfe99 100644 --- a/src/tools/AgentTool/agentToolUtils.ts +++ b/src/tools/AgentTool/agentToolUtils.ts @@ -31,6 +31,7 @@ import { getTokenCountFromTracker, isLocalAgentTask, killAsyncAgent, + syncProgressUsageFromMessages, type ProgressTracker, updateAgentProgress as updateAsyncAgentProgress, updateProgressFromMessage, @@ -625,6 +626,10 @@ export async function runAsyncAgentLifecycle({ resolveActivity, toolUseContext.options.tools, ) + // BUG-1 fix: message_delta mutates usage in-place after the message is + // yielded; recompute from the accumulated references so tokens converge + // to the real values once each request's usage lands. + syncProgressUsageFromMessages(tracker, agentMessages) updateAsyncAgentProgress( taskId, getProgressUpdate(tracker), @@ -648,6 +653,10 @@ export async function runAsyncAgentLifecycle({ if (!isCurrentExecution()) return if (abortController.signal.aborted) throw new AbortError() + // BUG-1 fix: last request's message_delta lands after the final yield — + // recompute so the completion notification carries real token counts. + syncProgressUsageFromMessages(tracker, agentMessages) + const agentResult = finalizeAgentTool(agentMessages, taskId, metadata) // Mark task completed FIRST so TaskOutput(block=true) unblocks