From dee3c778895d57ce952a97d945fc2e88b935b4ff Mon Sep 17 00:00:00 2001 From: Jack Minnetian <270441393+BlueBottleLatte@users.noreply.github.com> Date: Thu, 11 Jun 2026 17:21:25 -0700 Subject: [PATCH 1/7] feat(web): track estimated per-tool output token usage in chat Estimate the input-token footprint of each tool call's output (the cost the result imposes when fed back to the model on subsequent steps) using a local length-based estimator, persist it per tool call in the chat message metadata, and surface it inline in each tool call row next to the Details toggle. Estimates are ~-prefixed to keep them distinct from the authoritative billed token totals. --- packages/web/src/ee/features/chat/agent.ts | 25 +++++++++++-- .../components/chatThread/detailsCard.tsx | 27 +++++++++++--- .../chatThread/tools/mcpToolComponent.tsx | 10 +++++- .../chatThread/tools/toolOutputGuard.tsx | 10 ++++++ .../chatThread/tools/toolTokenBadge.tsx | 19 ++++++++++ .../web/src/features/chat/tokenEstimation.ts | 35 +++++++++++++++++++ packages/web/src/features/chat/types.ts | 12 +++++++ 7 files changed, 131 insertions(+), 7 deletions(-) create mode 100644 packages/web/src/ee/features/chat/components/chatThread/tools/toolTokenBadge.tsx create mode 100644 packages/web/src/features/chat/tokenEstimation.ts diff --git a/packages/web/src/ee/features/chat/agent.ts b/packages/web/src/ee/features/chat/agent.ts index 32c9befbf..336b04579 100644 --- a/packages/web/src/ee/features/chat/agent.ts +++ b/packages/web/src/ee/features/chat/agent.ts @@ -1,4 +1,5 @@ -import { SBChatMessage, SBChatMessageMetadata } from "@/features/chat/types"; +import { SBChatMessage, SBChatMessageMetadata, ToolTokenUsageEntry } from "@/features/chat/types"; +import { estimateToolOutputTokens } from "@/features/chat/tokenEstimation"; import { getFileSource } from '@/features/git'; import { isServiceError } from "@/lib/utils"; import { LanguageModelV3 as AISDKLanguageModelV3 } from "@ai-sdk/provider"; @@ -149,6 +150,8 @@ export const createMessageStream = async ({ const startTime = new Date(); + const collectedToolTokenUsage: ToolTokenUsageEntry[] = []; + const researchStream = await createAgentStream({ model, providerOptions: modelProviderOptions, @@ -163,6 +166,9 @@ export const createMessageStream = async ({ data: source, }); }, + onToolTokenUsage: (entry) => { + collectedToolTokenUsage.push(entry); + }, onMcpServerDiscovered: (sanitizedName, faviconUrl) => { writer.write({ type: 'data-mcp-server', @@ -200,6 +206,10 @@ export const createMessageStream = async ({ totalCacheReadTokens: (priorMetadata?.totalCacheReadTokens ?? 0) + (totalUsage.inputTokenDetails?.cacheReadTokens ?? 0), totalCacheWriteTokens: (priorMetadata?.totalCacheWriteTokens ?? 0) + (totalUsage.inputTokenDetails?.cacheWriteTokens ?? 0), totalResponseTimeMs: (priorMetadata?.totalResponseTimeMs ?? 0) + (new Date().getTime() - startTime.getTime()), + // Unlike the token totals above, this is concatenated (not + // summed) across approval-continuation phases so tool calls + // from the pre-approval phase are preserved. + toolTokenUsage: [...(priorMetadata?.toolTokenUsage ?? []), ...collectedToolTokenUsage], modelName, traceId, ...metadata, @@ -227,6 +237,7 @@ interface AgentOptions { inputMessages: ModelMessage[]; inputSources: Source[]; onWriteSource: (source: Source) => void; + onToolTokenUsage?: (entry: ToolTokenUsageEntry) => void; onMcpServerDiscovered: (sanitizedName: string, faviconUrl: string) => void; onMcpServerFailed: (serverName: string) => void; traceId: string; @@ -245,6 +256,7 @@ const createAgentStream = async ({ selectedRepos, disabledMcpServerIds, onWriteSource, + onToolTokenUsage, onMcpServerDiscovered, onMcpServerFailed, traceId, @@ -431,7 +443,16 @@ const createAgentStream = async ({ return null; }, onStepFinish: ({ toolResults }) => { - toolResults.forEach(({ output, dynamic }) => { + toolResults.forEach(({ toolCallId, toolName, output, dynamic }) => { + // Token estimation runs for every tool result — including + // dynamic (MCP) tools and error outputs — since they all + // re-enter the model's context on the next step. + onToolTokenUsage?.({ + toolCallId, + toolName, + estimatedOutputTokens: estimateToolOutputTokens(output), + }); + if (dynamic || isServiceError(output)) { return; } diff --git a/packages/web/src/ee/features/chat/components/chatThread/detailsCard.tsx b/packages/web/src/ee/features/chat/components/chatThread/detailsCard.tsx index 63bd1525e..442ec4645 100644 --- a/packages/web/src/ee/features/chat/components/chatThread/detailsCard.tsx +++ b/packages/web/src/ee/features/chat/components/chatThread/detailsCard.tsx @@ -59,6 +59,12 @@ const DetailsCardComponent = ({ (part.type === 'dynamic-tool' && part.toolName.startsWith('mcp_')) ).length, [thinkingSteps]); + // Lookup of estimated output tokens by tool call id, used to badge + // individual tool calls in the thinking steps. + const toolTokenUsageMap = useMemo(() => new Map( + (metadata?.toolTokenUsage ?? []).map(({ toolCallId, estimatedOutputTokens }) => [toolCallId, estimatedOutputTokens]) + ), [metadata?.toolTokenUsage]); + const cacheReadTokens = metadata?.totalCacheReadTokens ?? 0; const inputTokens = metadata?.totalInputTokens ?? 0; const cachedInputPercent = inputTokens > 0 @@ -202,6 +208,7 @@ const DetailsCardComponent = ({ thinkingSteps={thinkingSteps} isNetworkActive={isNetworkActive} isThinking={isThinking} + toolTokenUsageMap={toolTokenUsageMap} /> @@ -213,7 +220,7 @@ const DetailsCardComponent = ({ export const DetailsCard = memo(DetailsCardComponent, isEqual); -const ThinkingSteps = ({ thinkingSteps, isNetworkActive, isThinking }: { thinkingSteps: SBChatMessagePart[][], isNetworkActive: boolean, isThinking: boolean }) => { +const ThinkingSteps = ({ thinkingSteps, isNetworkActive, isThinking, toolTokenUsageMap }: { thinkingSteps: SBChatMessagePart[][], isNetworkActive: boolean, isThinking: boolean, toolTokenUsageMap?: Map }) => { const { scrollRef, contentRef, scrollToBottom } = useStickToBottom(); const [shouldStick, setShouldStick] = useState(isThinking); const prevIsThinking = usePrevious(isThinking); @@ -240,7 +247,10 @@ const ThinkingSteps = ({ thinkingSteps, isNetworkActive, isThinking }: { thinkin
{step.map((part, index) => (
- +
))}
@@ -251,7 +261,7 @@ const ThinkingSteps = ({ thinkingSteps, isNetworkActive, isThinking }: { thinkin } -export const StepPartRenderer = ({ part }: { part: SBChatMessagePart }) => { +export const StepPartRenderer = ({ part, estimatedOutputTokens }: { part: SBChatMessagePart, estimatedOutputTokens?: number }) => { switch (part.type) { case 'reasoning': case 'text': @@ -265,6 +275,7 @@ export const StepPartRenderer = ({ part }: { part: SBChatMessagePart }) => { return ( {(output) => } @@ -274,6 +285,7 @@ export const StepPartRenderer = ({ part }: { part: SBChatMessagePart }) => { return ( {(output) => } @@ -283,6 +295,7 @@ export const StepPartRenderer = ({ part }: { part: SBChatMessagePart }) => { return ( {(output) => } @@ -292,6 +305,7 @@ export const StepPartRenderer = ({ part }: { part: SBChatMessagePart }) => { return ( {(output) => } @@ -301,6 +315,7 @@ export const StepPartRenderer = ({ part }: { part: SBChatMessagePart }) => { return ( {(output) => } @@ -310,6 +325,7 @@ export const StepPartRenderer = ({ part }: { part: SBChatMessagePart }) => { return ( {(output) => } @@ -319,6 +335,7 @@ export const StepPartRenderer = ({ part }: { part: SBChatMessagePart }) => { return ( {(output) => } @@ -328,6 +345,7 @@ export const StepPartRenderer = ({ part }: { part: SBChatMessagePart }) => { return ( {(output) => } @@ -337,6 +355,7 @@ export const StepPartRenderer = ({ part }: { part: SBChatMessagePart }) => { return ( {(output) => } @@ -352,7 +371,7 @@ export const StepPartRenderer = ({ part }: { part: SBChatMessagePart }) => { return ; case 'dynamic-tool': if (part.toolName.startsWith('mcp_')) { - return ; + return ; } return null; case 'data-source': diff --git a/packages/web/src/ee/features/chat/components/chatThread/tools/mcpToolComponent.tsx b/packages/web/src/ee/features/chat/components/chatThread/tools/mcpToolComponent.tsx index c162d2841..2b4cc840f 100644 --- a/packages/web/src/ee/features/chat/components/chatThread/tools/mcpToolComponent.tsx +++ b/packages/web/src/ee/features/chat/components/chatThread/tools/mcpToolComponent.tsx @@ -4,10 +4,12 @@ import { CopyIconButton } from "@/app/(app)/components/copyIconButton"; import { McpFavicon } from "@/ee/features/chat/mcp/components/mcpFavicon"; import { useMcpServerIconMap } from "@/ee/features/chat/mcpServerIconContext"; import { cn } from "@/lib/utils"; +import { Separator } from "@/components/ui/separator"; import { DynamicToolUIPart } from "ai"; import { CheckCircle, ChevronDown, XCircle } from "lucide-react"; import { useCallback, useMemo, useState } from "react"; import { JsonHighlighter, unescapeJsonStrings } from "./jsonHighlighter"; +import { ToolTokenBadge } from "./toolTokenBadge"; export function parseMcpToolName(toolName: string): { serverName: string; toolName: string } | null { if (!toolName.startsWith('mcp_')) { @@ -24,7 +26,7 @@ export function parseMcpToolName(toolName: string): { serverName: string; toolNa }; } -export const McpToolComponent = ({ part }: { part: DynamicToolUIPart }) => { +export const McpToolComponent = ({ part, estimatedOutputTokens }: { part: DynamicToolUIPart, estimatedOutputTokens?: number }) => { const needsApproval = part.state === 'approval-requested'; const [isExpanded, setIsExpanded] = useState(needsApproval); const onToggle = useCallback(() => setIsExpanded(v => !v), []); @@ -128,6 +130,12 @@ export const McpToolComponent = ({ part }: { part: DynamicToolUIPart }) => {
{renderStatus()}
+ {estimatedOutputTokens !== undefined && ( + <> + + + + )} {hasInput && (