Repository navigation
fix: token usage tracking in AI models and UI components - #710
Conversation
…ining context token calculations in chat transport and UI components.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThis PR enables token usage tracking across model clients and refactors how token usage is persisted and displayed. It introduces ChangesToken Usage Tracking and Display
Sequence DiagramsequenceDiagram
participant ToolLoopAgent
participant StreamHandler
participant MessageMetadata
ToolLoopAgent->>StreamHandler: emit finish-step event with inputTokens
StreamHandler->>StreamHandler: record inputTokens to lastStepInputTokens
ToolLoopAgent->>StreamHandler: emit finish event with totalUsage
StreamHandler->>MessageMetadata: compute final usage (prefer totalUsage)
StreamHandler->>MessageMetadata: attach lastStepInputTokens
MessageMetadata->>MessageMetadata: store in session token usage
Estimated Code Review Effort🎯 3 (Moderate) | ⏱️ ~25 minutes Suggested Reviewers
Poem
🚥 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)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 ESLint
ESLint install failed. For unrecoverable errors, disable the tool in CodeRabbit configuration. 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 |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4659889b41
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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.
Inline comments:
In `@packages/ai-core/src/chatTransport.ts`:
- Around line 389-401: The finish branch only reads (part as {totalUsage?:
LanguageModelUsage}).totalUsage and ignores a possible (part as {usage?:
LanguageModelUsage}).usage payload, which can contain final usage when
totalUsage is missing; update the logic in chatTransport.ts around
toMessageTokenUsage/finishUsage so it checks both totalUsage and usage (e.g.,
prefer totalUsage, else use usage) before falling back to accumulatedUsage, then
set finalUsage.lastStepInputTokens and call rememberSessionTokenUsage(sessionId,
finalUsage) as before; reference toMessageTokenUsage, finishUsage,
accumulatedUsage, finalUsage, lastStepInputTokens, and rememberSessionTokenUsage
to locate the change.
In `@packages/ai-core/src/components/ContextUsageIndicator.tsx`:
- Around line 71-104: The loop currently mixes reported assistant tokenUsage
with per-message estimates, inflating cumulativeTokens; change the logic so
reported usage and fallback estimates are mutually exclusive: when an assistant
message has metadata.tokenUsage (reported.totalTokens), add only
reported.totalTokens to cumulativeTokens and set lastContextTokens from
reported.lastStepInputTokens/inputTokens as before, and set
hasAnyReportedUsage=true; when hasAnyReportedUsage is true, do NOT call
estimateMessageTokens for any other messages (skip estimating user/assistant
messages without usage); only when no reported usage at all fall back to
estimating every message and include systemInstructions via
estimateTokenCount(systemInstructions). Ensure contextTokens uses
lastContextTokens when hasAnyReportedUsage, otherwise uses cumulativeTokens.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: c4fad06d-7318-44a7-bcf8-9c7052fe1b2e
📒 Files selected for processing (8)
apps/sqlrooms-cli-ui/src/createDashboardAgent.tsexamples/ai-agent/src/skills/getModel.tsexamples/ai-nextjs/src/app/api/chat/route.tsexamples/ai-subagents/src/agents/TravelPlannerAgent.tspackages/ai-core/src/AiSlice.tspackages/ai-core/src/chatTransport.tspackages/ai-core/src/components/ContextUsageIndicator.tsxpackages/ai-core/src/types.ts
improves the accuracy and consistency of token usage tracking and reporting across the AI agent framework and UI
onFinishcallback from the agent creation options increateLocalChatTransportFactory, as token usage is now handled centrally.Summary by CodeRabbit