feat: track and display input/output tokens separately#1251
feat: track and display input/output tokens separately#1251shivamashtikar wants to merge 1 commit intomainfrom
Conversation
- Add inputTokens and outputTokens columns to messages table - Update all insertMessage calls to store tokens separately - Update aggregation queries in sharedAgentUsage, admin, and message modules - Add input/output token display to Dashboard, AdminChatsTable components - Maintain backward compatibility with existing totalTokens field
shivamashtikar
requested review from
devesh-juspay,
junaid-shirur,
kalpadhwaryu and
zereraz
as code owners
📝 WalkthroughWalkthroughThis pull request extends token tracking across the system to separately capture input and output tokens. Changes include database schema updates, SQL aggregation logic, backend API responses, and frontend UI displays to surface per-direction token breakdowns alongside existing total token metrics. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Tip Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord. 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 |
Summary of ChangesHello @shivamashtikar, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a significant enhancement to how token usage is tracked and displayed within the application. By separating input and output tokens at the database level and propagating this change through the API and frontend, users and administrators can now gain a more precise understanding of LLM interaction costs and resource consumption. This provides greater transparency and improved analytics without disrupting existing data. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request effectively implements the tracking and separate display of input and output tokens, which is a great enhancement for usage monitoring. The changes are consistently applied across the database, backend APIs, and frontend components. I've found one issue in the Dashboard.tsx component where a copy-paste error has led to duplicated token information and the removal of the message count display. My review includes a specific comment with a suggested fix for this.
| <div className="flex flex-col items-center"> | ||
| <span className="text-sm font-medium"> | ||
| {(agent.totalTokens || 0).toLocaleString()} | ||
| </span> | ||
| <span className="text-xs text-muted-foreground"> | ||
| total tokens | ||
| </span> | ||
| <span className="text-xs text-muted-foreground"> | ||
| ({(agent.inputTokens || 0).toLocaleString()} in / {(agent.outputTokens || 0).toLocaleString()} out) | ||
| </span> | ||
| </div> |
There was a problem hiding this comment.
This block incorrectly replaces the 'total messages' display with a 'total tokens' display. This seems to be a copy-paste error, as there's another block that correctly updates the token display. This results in duplicated token information and the message count being removed from the SharedAgentUsageCard.
To fix this, this block should be reverted to display the agent.totalMessages as it did before.
<div className="flex flex-col items-center">
<span className="text-sm font-medium">
{agent.totalMessages}
</span>
<span className="text-xs text-muted-foreground">
messages
</span>
</div>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
server/db/message.ts (1)
166-200:⚠️ Potential issue | 🟡 MinorReturn type doesn’t include
inputTokens/outputTokens.
getMessageFeedbackStatsnow returns these fields, but the Promise return type omits them, so callers can’t access them without type assertions. Please align the signature with the new payload.✅ Suggested type update
export async function getMessageFeedbackStats({ ... }): Promise<{ totalLikes: number totalDislikes: number totalCost: number totalTokens: number + inputTokens: number + outputTokens: number feedbackByChat: Record< string, { likes: number; dislikes: number; cost: number; tokens: number } > feedbackMessages: Array<{ messageId: string chatExternalId: string type: "like" | "dislike" feedbackText: string[] timestamp: string }> }> {server/api/chat/message-agents.ts (1)
4985-5001:⚠️ Potential issue | 🟠 MajorUpdate
agentContext.tokenUsagein thetoken_usageevent handler to populate token counts.The
token_usageevent (lines 4846–4855) only logs telemetry and does not updateagentContext.tokenUsage. SinceinputTokensandoutputTokens(lines 4999–5000) are persisted directly from the unmodified context, they remain0for all MessageAgents runs. Similarly,totalTokens(line 4982) and all aggregated token counts will be0. Update the event handler to accumulate token values into the context, confirming first whether JAF'sevt.datavalues are deltas or cumulative totals.Suggested fix
case "token_usage": { const tokenUsageSpan = jafStreamingSpan.startSpan("token_usage") tokenUsageSpan.setAttribute("prompt_tokens", evt.data.prompt ?? 0) tokenUsageSpan.setAttribute( "completion_tokens", evt.data.completion ?? 0 ) tokenUsageSpan.setAttribute("total_tokens", evt.data.total ?? 0) + // TODO: confirm whether evt.data.* are deltas or totals from JAF + agentContext.tokenUsage.input += evt.data.prompt ?? 0 + agentContext.tokenUsage.output += evt.data.completion ?? 0 tokenUsageSpan.end() break }server/api/chat/agents.ts (1)
705-731:⚠️ Potential issue | 🟠 MajorPrevent NaN token totals by defaulting missing usage values to 0.
If any chunk’s usage omits
inputTokensoroutputTokens, the reduction yieldsNaN, which then propagates totokensUsed/inputTokens/outputTokensand can break inserts or analytics. Default to 0 during aggregation (or normalize when pushing).🛠️ Suggested defensive aggregation (apply to all token reductions)
- const totalTokens = tokenArr.reduce( - (acc, tokens) => ({ - inputTokens: acc.inputTokens + tokens.inputTokens, - outputTokens: acc.outputTokens + tokens.outputTokens, - }), - { inputTokens: 0, outputTokens: 0 }, - ) + const totalTokens = tokenArr.reduce( + (acc, tokens) => ({ + inputTokens: acc.inputTokens + (tokens.inputTokens ?? 0), + outputTokens: acc.outputTokens + (tokens.outputTokens ?? 0), + }), + { inputTokens: 0, outputTokens: 0 }, + )Also applies to: 742-768, 774-800, 925-951, 1676-1703, 1926-1954, 2561-2588, 2843-2870
frontend/src/components/AdminChatsTable.tsx (1)
619-627:⚠️ Potential issue | 🟡 MinorHide input/output breakdown when legacy data has zeros.
For older chats,
inputTokens/outputTokenswill be 0 whiletotalTokensis non-zero, and the UI will misleadingly show “0 in / 0 out.” Render the breakdown only when either value is present.✅ Conditional display to avoid misleading zeros
- <span className="text-xs text-muted-foreground"> - ({(chat.inputTokens || 0).toLocaleString()} in / {(chat.outputTokens || 0).toLocaleString()} out) - </span> + {(chat.inputTokens || 0) + (chat.outputTokens || 0) > 0 && ( + <span className="text-xs text-muted-foreground"> + ({(chat.inputTokens || 0).toLocaleString()} in / {(chat.outputTokens || 0).toLocaleString()} out) + </span> + )}
🤖 Fix all issues with AI agents
In `@frontend/src/components/Dashboard.tsx`:
- Around line 4631-4635: processAdminChatsData currently returns an
AdminDashboardStats and per-user AdminUserUsage objects without populating
inputTokens/outputTokens, causing "0 in / 0 out" everywhere; update
processAdminChatsData to accumulate chat-level token counts (inputTokens and
outputTokens) into both the top-level AdminDashboardStats return object and into
each per-user AdminUserUsage entry as you iterate chats (ensure you sum
chat.inputTokens/chat.outputTokens or equivalent fields when building totals and
user aggregates), and return those populated fields so MetricCard and
AdminUsersLeaderboard consume the real token values.
- Around line 3178-3182: processUserChatsData currently returns totalTokens but
never computes inputTokens or outputTokens, causing MetricCard to always show "0
in / 0 out"; update processUserChatsData to accumulate inputTokens and
outputTokens from each chat/message (e.g., sum request prompt tokens into
inputTokens and model response tokens into outputTokens) and include these
fields on the returned userStats object, or alternatively remove the in/out
breakdown from where MetricCard renders (reference MetricCard props using
userStats.inputTokens and userStats.outputTokens) until those fields are
populated; ensure the returned object from processUserChatsData contains numeric
inputTokens and outputTokens to avoid relying on the || 0 fallback.
- Around line 976-986: There are two identical JSX blocks rendering token info
(a div with className "flex flex-col items-center" that displays
(agent.totalTokens || 0).toLocaleString() and the in/out tokens using
agent.inputTokens and agent.outputTokens) inside the agent row; remove the
duplicate block so each agent card only renders this token column once, leaving
the single remaining token-info div and adjusting any surrounding layout/spacing
if needed.
| <div className="flex flex-col items-center"> | ||
| <span className="text-sm font-medium"> | ||
| {(agent.totalTokens || 0).toLocaleString()} | ||
| </span> | ||
| <span className="text-xs text-muted-foreground"> | ||
| total tokens | ||
| </span> | ||
| <span className="text-xs text-muted-foreground"> | ||
| ({(agent.inputTokens || 0).toLocaleString()} in / {(agent.outputTokens || 0).toLocaleString()} out) | ||
| </span> | ||
| </div> |
There was a problem hiding this comment.
Bug: Duplicate token column in the shared agent list row.
Lines 976–986 and 1011–1021 render identical "total tokens" + "(in / out)" blocks within the same agent card. This results in the token information appearing twice per row. One of these blocks should be removed.
🐛 Remove the duplicate block (lines 1011–1021)
</div>
- <div className="flex flex-col items-center">
- <span className="text-sm font-medium">
- {(agent.totalTokens || 0).toLocaleString()}
- </span>
- <span className="text-xs text-muted-foreground">
- total tokens
- </span>
- <span className="text-xs text-muted-foreground">
- ({(agent.inputTokens || 0).toLocaleString()} in / {(agent.outputTokens || 0).toLocaleString()} out)
- </span>
- </div>
<div className="flex items-center gap-2 text-xs">Also applies to: 1011-1021
🤖 Prompt for AI Agents
In `@frontend/src/components/Dashboard.tsx` around lines 976 - 986, There are two
identical JSX blocks rendering token info (a div with className "flex flex-col
items-center" that displays (agent.totalTokens || 0).toLocaleString() and the
in/out tokens using agent.inputTokens and agent.outputTokens) inside the agent
row; remove the duplicate block so each agent card only renders this token
column once, leaving the single remaining token-info div and adjusting any
surrounding layout/spacing if needed.
| <MetricCard | ||
| title="Total Tokens" | ||
| value={(userStats.totalTokens || 0).toLocaleString()} | ||
| description="Tokens processed" | ||
| description={`${(userStats.inputTokens || 0).toLocaleString()} in / ${(userStats.outputTokens || 0).toLocaleString()} out`} | ||
| icon={Activity} |
There was a problem hiding this comment.
Bug: userStats.inputTokens / outputTokens are always undefined → displays "0 in / 0 out".
processUserChatsData (line 2893) returns an object with totalTokens but never computes inputTokens or outputTokens. The fallback || 0 prevents a crash, but the UI will always show zeros here, which is misleading.
Either accumulate inputTokens/outputTokens from the chat data inside processUserChatsData, or don't display the in/out breakdown on this page until the data is available.
🐛 Sketch: propagate input/output tokens in processUserChatsData
const processUserChatsData = (chats: any[]) => {
...
let totalTokensCalculated = 0
+ let inputTokensCalculated = 0
+ let outputTokensCalculated = 0
...
chats.forEach((chat) => {
...
totalTokensCalculated += chat.totalTokens || 0
+ inputTokensCalculated += chat.inputTokens || 0
+ outputTokensCalculated += chat.outputTokens || 0
...
})
return {
...
totalTokens: totalTokensCalculated,
+ inputTokens: inputTokensCalculated,
+ outputTokens: outputTokensCalculated,
...
}
}🤖 Prompt for AI Agents
In `@frontend/src/components/Dashboard.tsx` around lines 3178 - 3182,
processUserChatsData currently returns totalTokens but never computes
inputTokens or outputTokens, causing MetricCard to always show "0 in / 0 out";
update processUserChatsData to accumulate inputTokens and outputTokens from each
chat/message (e.g., sum request prompt tokens into inputTokens and model
response tokens into outputTokens) and include these fields on the returned
userStats object, or alternatively remove the in/out breakdown from where
MetricCard renders (reference MetricCard props using userStats.inputTokens and
userStats.outputTokens) until those fields are populated; ensure the returned
object from processUserChatsData contains numeric inputTokens and outputTokens
to avoid relying on the || 0 fallback.
| <MetricCard | ||
| title="Total Tokens" | ||
| value={(adminStats.totalTokens || 0).toLocaleString()} | ||
| description="Total tokens processed" | ||
| description={`${(adminStats.inputTokens || 0).toLocaleString()} in / ${(adminStats.outputTokens || 0).toLocaleString()} out`} | ||
| icon={Activity} |
There was a problem hiding this comment.
Bug: adminStats.inputTokens / outputTokens are never populated → always "0 in / 0 out".
processAdminChatsData (line 3977) returns an AdminDashboardStats object but never sets inputTokens or outputTokens in the return value (line 4268). The same issue applies to the per-user objects built at line 4112 — AdminUserUsage entries lack these fields, so the AdminUsersLeaderboard at line 2241 will also always show zeros.
You need to accumulate inputTokens and outputTokens in both the top-level stats and per-user objects inside processAdminChatsData, or source these values from a server-side aggregation endpoint instead of computing them client-side from chat-level data that may not carry the breakdown.
🤖 Prompt for AI Agents
In `@frontend/src/components/Dashboard.tsx` around lines 4631 - 4635,
processAdminChatsData currently returns an AdminDashboardStats and per-user
AdminUserUsage objects without populating inputTokens/outputTokens, causing "0
in / 0 out" everywhere; update processAdminChatsData to accumulate chat-level
token counts (inputTokens and outputTokens) into both the top-level
AdminDashboardStats return object and into each per-user AdminUserUsage entry as
you iterate chats (ensure you sum chat.inputTokens/chat.outputTokens or
equivalent fields when building totals and user aggregates), and return those
populated fields so MetricCard and AdminUsersLeaderboard consume the real token
values.
1 participant
Description
Testing
Additional Notes
Summary by CodeRabbit
New Features
Chores