0ebb1adce8
Tokens were only counted post-hoc (onFinish) and the header badge updated only on
chat open/switch; reasoning wasn't requested or shown. Now a counter ticks LIVE
during generation and surfaces reasoning ("thinking") tokens separately, like
Claude Code's `Thinking… · N tokens`.
Architecture (AI SDK v6): no provider gives exact per-token usage mid-stream, so
the live number is a cheap client estimate (chars/≈4) reconciled to AUTHORITATIVE
provider usage at step boundaries and turn end. The useChat per-delta re-render is
the existing realtime engine.
- server: `chatStreamMetadata` now also forwards usage on `finish-step` + `finish`;
`sendReasoning: true`; persisted `metadata.usage` carries `reasoningTokens`
(normalized from `outputTokenDetails` or the deprecated field).
- client: pure `count-stream-tokens` (estimateTokens / liveTurnTokens, prefers
authoritative usage else estimate); `Thinking… · N tokens` in the typing
indicator; collapsible "Thinking" reasoning block; throttled (~8 Hz) live
turn-token header badge; `reasoningTokens` in types + Markdown export.
Review fixes folded in:
- v6 `finish-step.usage` is PER-STEP, not cumulative — the server now ACCUMULATES
a running sum (new pure `accumulateStepUsage`) and sends the cumulative, which
converges to `finish.totalUsage`, so the live counter never jumps DOWN on a
multi-step agent turn.
- reasoning double-count: the authoritative turn-total is attributed to a block
ONLY for a single-reasoning-part (one-step) turn; multi-step blocks each show
their own estimate (the authoritative total stays in the header).
- no "0" badge flash at turn start (require live > 0, else show context size).
- comment refreshed (finish-step trigger).
Tests: server `accumulateStepUsage` + updated `chatStreamMetadata` (34 in the
suite); client pure-fn tests. Both tsc clean; 162 client ai-chat + the ai-chat
server suite pass.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
51 lines
1.6 KiB
TypeScript
51 lines
1.6 KiB
TypeScript
import { describe, expect, it } from "vitest";
|
|
import type { UIMessage } from "@ai-sdk/react";
|
|
import { tailThinkingTokens } from "@/features/ai-chat/components/message-list.tsx";
|
|
|
|
/**
|
|
* Pure-helper tests for `tailThinkingTokens`: the live thinking-token count the
|
|
* standalone typing indicator shows. It is the reasoning split of the tail
|
|
* assistant message (estimate while streaming, authoritative once usage arrives).
|
|
*/
|
|
const msg = (
|
|
role: "user" | "assistant",
|
|
parts: unknown[],
|
|
metadata?: unknown,
|
|
): UIMessage =>
|
|
({ id: Math.random().toString(), role, parts, metadata }) as UIMessage;
|
|
|
|
describe("tailThinkingTokens", () => {
|
|
it("is 0 when there are no messages", () => {
|
|
expect(tailThinkingTokens([])).toBe(0);
|
|
});
|
|
|
|
it("is 0 when the tail message is the user's", () => {
|
|
expect(tailThinkingTokens([msg("user", [{ type: "text", text: "q" }])])).toBe(0);
|
|
});
|
|
|
|
it("is 0 when the assistant has produced no reasoning yet", () => {
|
|
expect(
|
|
tailThinkingTokens([msg("assistant", [{ type: "text", text: "answer" }])]),
|
|
).toBe(0);
|
|
});
|
|
|
|
it("estimates reasoning tokens from streamed reasoning text", () => {
|
|
// 8 chars -> 2 tokens.
|
|
expect(
|
|
tailThinkingTokens([
|
|
msg("assistant", [{ type: "reasoning", text: "12345678" }]),
|
|
]),
|
|
).toBe(2);
|
|
});
|
|
|
|
it("uses authoritative usage.reasoningTokens once the server attaches it", () => {
|
|
expect(
|
|
tailThinkingTokens([
|
|
msg("assistant", [{ type: "reasoning", text: "x" }], {
|
|
usage: { outputTokens: 100, reasoningTokens: 42 },
|
|
}),
|
|
]),
|
|
).toBe(42);
|
|
});
|
|
});
|