fix: dedupe token usage aggregation by message id (#2770)

This commit is contained in:
YuJitang
2026-05-08 09:54:20 +08:00
committed by GitHub
parent 6c220a9aef
commit 530bda7107
2 changed files with 103 additions and 6 deletions
+22 -6
View File
@@ -28,7 +28,12 @@ export function getUsageMetadata(message: Message): TokenUsage | null {
}
/**
* Accumulate token usage across all AI messages in a thread.
* Accumulate token usage across AI messages.
*
* UI rendering may place the same AI message in more than one group, such as
* when a message contains both reasoning and final answer content. Token usage
* is attached to the AI message itself, so a message id should only contribute
* once to any aggregate.
*/
export function accumulateUsage(messages: Message[]): TokenUsage | null {
const cumulative: TokenUsage = {
@@ -37,14 +42,25 @@ export function accumulateUsage(messages: Message[]): TokenUsage | null {
totalTokens: 0,
};
let hasUsage = false;
const countedMessageIds = new Set<string>();
for (const message of messages) {
const usage = getUsageMetadata(message);
if (usage) {
hasUsage = true;
cumulative.inputTokens += usage.inputTokens;
cumulative.outputTokens += usage.outputTokens;
cumulative.totalTokens += usage.totalTokens;
if (!usage) {
continue;
}
if (message.id) {
if (countedMessageIds.has(message.id)) {
continue;
}
countedMessageIds.add(message.id);
}
hasUsage = true;
cumulative.inputTokens += usage.inputTokens;
cumulative.outputTokens += usage.outputTokens;
cumulative.totalTokens += usage.totalTokens;
}
return hasUsage ? cumulative : null;
}