fix: use backend thread token usage for header total (#2800)

* fix: use backend thread token usage for header total

* Refactor thread token usage fetch
This commit is contained in:
YuJitang
2026-05-09 19:40:32 +08:00
committed by GitHub
parent 881ff71252
commit 417416087b
16 changed files with 540 additions and 35 deletions
@@ -1,7 +1,7 @@
import type { Message } from "@langchain/langgraph-sdk";
import { expect, test } from "vitest";
import { accumulateUsage } from "@/core/messages/usage";
import { accumulateUsage, selectHeaderTokenUsage } from "@/core/messages/usage";
import {
getAssistantTurnUsageMessages,
getMessageGroups,
@@ -79,3 +79,86 @@ test("keeps header and per-turn aggregation consistent for duplicated UI groups"
totalTokens: 27,
});
});
test("prefers backend thread usage for header totals", () => {
const messages = [
{
id: "ai-visible",
type: "ai",
content: "Visible answer",
usage_metadata: { input_tokens: 10, output_tokens: 5, total_tokens: 15 },
},
] as Message[];
expect(
selectHeaderTokenUsage({
backendUsage: { inputTokens: 100, outputTokens: 50, totalTokens: 150 },
messages,
}),
).toEqual({
inputTokens: 100,
outputTokens: 50,
totalTokens: 150,
});
});
test("adds current in-flight message usage to backend header totals", () => {
const completedMessages = [
{
id: "ai-completed",
type: "ai",
content: "Completed answer",
usage_metadata: { input_tokens: 10, output_tokens: 5, total_tokens: 15 },
},
{
id: "ai-pending",
type: "ai",
content: "Streaming answer",
usage_metadata: { input_tokens: 4, output_tokens: 6, total_tokens: 10 },
},
] as Message[];
expect(
selectHeaderTokenUsage({
backendUsage: { inputTokens: 100, outputTokens: 50, totalTokens: 150 },
messages: completedMessages,
pendingMessages: [completedMessages[1]!],
}),
).toEqual({
inputTokens: 104,
outputTokens: 56,
totalTokens: 160,
});
});
test("falls back to visible messages when backend usage is unavailable or zero", () => {
const messages = [
{
id: "ai-visible",
type: "ai",
content: "Visible answer",
usage_metadata: { input_tokens: 10, output_tokens: 5, total_tokens: 15 },
},
] as Message[];
expect(
selectHeaderTokenUsage({
backendUsage: null,
messages,
}),
).toEqual({
inputTokens: 10,
outputTokens: 5,
totalTokens: 15,
});
expect(
selectHeaderTokenUsage({
backendUsage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
messages,
}),
).toEqual({
inputTokens: 10,
outputTokens: 5,
totalTokens: 15,
});
});