feat: support richer token cost tracking for ask#1353
Merged
Conversation
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.
Record the provider-reported input/output token usage of each agent step in the chat message metadata and display it per step group in the thinking steps view (joined to UI step groups via the step index now tagged on each tool token usage entry). Also fix the tool output estimator to measure the model-visible payload: tools with a toModelOutput mapping (all builtins) send only their output text to the model, so estimating the raw ToolResult object was counting UI-only metadata the model never sees. The bytes-per-token ratio is now a uniform ~2 chars/token, calibrated against provider-reported per-step usage of code-heavy tool results.
Collect usage from researchStream.steps and response.messages after the stream completes (covers approval-gated and failed tool calls, off the hot path), nest tool estimates under their step in a single stepTokenUsage array, and join UI steps to entries by stepIndex.
This comment has been minimized.
This comment has been minimized.
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (3)
💤 Files with no reviewable changes (2)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughAdds per-step token usage tracking for the Ask Sourcebot chat agent. A new offline estimation utility computes token footprints from tool outputs. The message metadata schema gains a ChangesPer-step tool token usage estimation and display
Sequence Diagram(s)sequenceDiagram
participant Agent as agent.ts
participant ResearchStream as researchStream
participant Metadata as message-metadata
participant DetailsCard as DetailsCard
participant ThinkingSteps as ThinkingSteps
participant ToolTokenBadge as ToolTokenBadge
Agent->>ResearchStream: await stream completion
ResearchStream-->>Agent: response.messages, steps
Agent->>Agent: scan tool-result parts → toolUsageByToolCallId
Agent->>Agent: align to steps → stepTokenUsage[]
Agent->>Metadata: emit stepTokenUsage (concat with priorMetadata)
Metadata-->>DetailsCard: metadata.stepTokenUsage
DetailsCard->>DetailsCard: build toolTokenUsageMap
DetailsCard->>ThinkingSteps: ThinkingStep[], toolTokenUsageMap, answerStepIndex
ThinkingSteps->>ThinkingSteps: lookup usage[stepIndex] per step
ThinkingSteps->>ToolTokenBadge: estimatedOutputTokens per tool call
ToolTokenBadge-->>ThinkingSteps: badge UI with tooltip
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
packages/web/src/ee/features/chat/agent.ts (1)
257-263:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPrevent computed token metadata from being overwritten by caller metadata
On Line 262, spreading
...metadatalast lets externalmetadatareplace derived fields (stepTokenUsage, totals,modelName,traceId), which can break the step-index join contract used by the UI.Suggested fix
writer.write({ type: 'message-metadata', messageMetadata: { + ...metadata, totalTokens: (priorMetadata?.totalTokens ?? 0) + (totalUsage.totalTokens ?? 0), totalInputTokens: (priorMetadata?.totalInputTokens ?? 0) + (totalUsage.inputTokens ?? 0), totalOutputTokens: (priorMetadata?.totalOutputTokens ?? 0) + (totalUsage.outputTokens ?? 0), 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()), // Concatenated (not summed) across approval-continuation // phases so earlier phases' steps are preserved in order. stepTokenUsage: [...(priorMetadata?.stepTokenUsage ?? []), ...stepTokenUsage], modelName, traceId, - ...metadata, } });🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/web/src/ee/features/chat/agent.ts` around lines 257 - 263, The spread operator `...metadata` is placed last in the object literal, which allows caller-provided metadata to overwrite the computed derived fields (stepTokenUsage, modelName, traceId, and totals). Reorder the object properties by moving `...metadata` to the beginning of the object literal before the computed fields, so that the carefully derived values (stepTokenUsage, modelName, traceId, and any totals fields) are spread after the external metadata and cannot be accidentally overwritten by caller data.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@packages/web/src/ee/features/chat/agent.ts`:
- Around line 257-263: The spread operator `...metadata` is placed last in the
object literal, which allows caller-provided metadata to overwrite the computed
derived fields (stepTokenUsage, modelName, traceId, and totals). Reorder the
object properties by moving `...metadata` to the beginning of the object literal
before the computed fields, so that the carefully derived values
(stepTokenUsage, modelName, traceId, and any totals fields) are spread after the
external metadata and cannot be accidentally overwritten by caller data.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 181d5be9-92d5-4715-9a94-791f21b19000
📒 Files selected for processing (12)
packages/web/src/ee/features/chat/agent.test.tspackages/web/src/ee/features/chat/agent.tspackages/web/src/ee/features/chat/components/chatThread/chatThreadListItem.tsxpackages/web/src/ee/features/chat/components/chatThread/detailsCard.test.tsxpackages/web/src/ee/features/chat/components/chatThread/detailsCard.tsxpackages/web/src/ee/features/chat/components/chatThread/tools/mcpToolComponent.tsxpackages/web/src/ee/features/chat/components/chatThread/tools/toolOutputGuard.tsxpackages/web/src/ee/features/chat/components/chatThread/tools/toolSearchToolComponent.tsxpackages/web/src/ee/features/chat/components/chatThread/tools/toolTokenBadge.tsxpackages/web/src/features/chat/tokenEstimation.test.tspackages/web/src/features/chat/tokenEstimation.tspackages/web/src/features/chat/types.ts
jsourcebot
changed the title
Spread caller-supplied metadata before the derived token fields so stepTokenUsage and the totals can't be clobbered, which would desync the UI's index-based step join.
brendan-kellam
previously approved these changes
Jun 19, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds tracking of token costs per step, also adds estimates of tool call token usage. This information is embedded in the chat history. Tool call token usages are estimates because a single step can run multiple tool calls and there is no mechanism to discern which part of the input token cost came from which tool call.
Summary by CodeRabbit
New Features
Refactor
Tests