From d1c2051f2851724b85bb599ef1acd0fc073f0d19 Mon Sep 17 00:00:00 2001 From: Junmo Kim Date: Mon, 18 May 2026 11:40:51 +0900 Subject: [PATCH] fix(web): aggregate per-response metadata so multi-turn cards show total usage (#637) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(web): extend MessageMetadata to accept aggregated turnCount Add an optional `turnCount` prop to MessageMetadata so the same builder can render an aggregated response-group footer when the caller has already summed usage and dedup-joined model ids. The label set switches to `Models` / `Total` / `N turns` only when `turnCount >= 2`, leaving single-turn footers byte-identical with the existing `Invoke · Model · Usage` output. Also expose `turnCount?: number` on `HappyChatMessageMetadata` so a later commit can inject the aggregated metadata through the library's ThreadMessageLike payload without widening the type at the same time. No call site passes `turnCount` yet, so this commit is behavior-neutral on all existing surfaces (proof-of-invariance test included). * feat(web): aggregate per-response metadata so multi-turn cards show total usage The `@assistant-ui/react` converter joins adjacent assistant messages into one card but only preserves `metadata.custom` from the first block, so multi-turn responses currently show the first turn's usage and model only. Compute response-group aggregates in `useHappyRuntime` and inject the sum on each group's first visible block, where the library will keep them. Per group: usage tokens are summed across distinct turns, model ids are dedup-joined in first-seen order, and the invoke time is the first turn's so the footer keeps showing when the response started (regression-guarded by unit test). `durationMs` is explicitly cleared on aggregated blocks because the first turn's value would otherwise leak through the join. Turn identity prefers the CLI-stamped `localId`. When that is null (claude code spawn sessions today emit `localId=null` on every chunk) the aggregator falls back to a fingerprint built from `model` plus the shared `usage` totals — every block emitted within one Claude SDK message carries an identical usage object, so the fingerprint dedups those chunks without merging distinct turns whose token counts naturally differ. Tool-result chunks with no model or usage are skipped so they cannot inflate the turn count. Single-turn responses get no aggregate entry, so their footers stay byte-identical with the existing behavior. Test plan - `assistant-runtime.test.ts` covers the six grouping scenarios spelled out in the design note (localId-based + null-localId fingerprint fallback) plus two defensive cases for tool_result chunks and cache token preservation. * fix(web): preserve explicit zero sums and count tool-group turns in response aggregator Two correctness gaps in aggregateResponseGroups: - addUsage folded `0 + 0` through `|| undefined`, dropping an explicit-zero cache token sum from the aggregated metadata. Replace the falsy fold with sumOptional(): undefined only when both operands are absent, otherwise (a ?? 0) + (b ?? 0). - turnSourceFromBlock returned null for tool-group blocks, so a card whose visible-first block is a tool-group dropped its turn entirely. Read the first underlying tool-call instead; degrade to null only when the group somehow holds zero tools. Unit tests cover both regressions: tool-group as the first visible block in a response group, explicit-zero cache sums preserved, and the empty-tool-group degrade-to-null path. * fix(web): dedup response-group turns by adjacency rather than set membership The fingerprint fallback (used when localId is null) compared each turn key against a Set of every key seen in the group. A response group whose first and third turns happened to carry the same (model, usage) fingerprint would collapse the third turn into the first, under-counting the visible turn count. Switch to ordering-based dedup: each block's turn key only collides with the immediately previous turn. Adjacent blocks within one SDK message still collapse (their usage object is identical), but non-adjacent fingerprint matches across separate turns stay distinct. Behavior under localId-stamped flows is unchanged because distinct turns always carry distinct localIds. Unit test covers a three-turn group whose first and third turns share a fingerprint with a different middle turn between them. * fix(web): aggregate every tool-call in a tool-group and dedup by createdAt fingerprint `buildVisibleChatBlocks` merges adjacent eligible tool-calls into a single `tool-group` without checking that they share a turn. Reading only the first underlying tool would drop every later tool turn from the aggregate, so each tool-call in the group now contributes its own turn source. The fingerprint fallback (used when the CLI does not stamp `localId`) gains `createdAt` as a third axis. The reducer copies `msg.createdAt` onto every derived ChatBlock, so blocks from one SDK message still collapse to one turn, while two adjacent turns that happen to coincide on `(model, usage)` no longer dedup against each other. Same wall-clock millisecond collisions remain theoretically possible but are bounded by the hub stamp resolution. Helper layer consolidates: `turnSourceFromBlock` (single-or-null) is gone, replaced by `turnSourcesFromBlock` returning the array directly. Test renames clarify the contract — the existing tool-group test now documents the same-turn collapse case — and one new test pins the fingerprint coincidence case. * fix(web): make tool-only response cards expose aggregate metadata `aggregateResponseGroups` keys aggregate metadata onto a response group's first visible block, which can be a `tool-group` when the assistant turn starts with tools. The `toolOnly` render branch did not wire the click toggle that the default/codex branches use, so the new Models/Total/N-turns footer stayed unreachable for those cards. Wrap the toolOnly content with the same cursor-pointer div used in the sibling branches (toggleMetadata, onMetadataKeyDown, role=button, aria-expanded). Carry `min-w-0` on the wrapper so long tool labels keep clipping under the existing `overflow-x-hidden` on MessagePrimitive.Root. The shared `isNestedInteractiveEvent` guard prevents the wrapper toggle from firing when nested tool buttons or disclosures are clicked. --- .../messages/AssistantMessage.tsx | 33 +- .../messages/MessageMetadata.test.ts | 75 +++ .../messages/MessageMetadata.tsx | 28 +- web/src/lib/assistant-runtime.test.ts | 617 ++++++++++++++++++ web/src/lib/assistant-runtime.ts | 263 +++++++- 5 files changed, 1000 insertions(+), 16 deletions(-) create mode 100644 web/src/lib/assistant-runtime.test.ts diff --git a/web/src/components/AssistantChat/messages/AssistantMessage.tsx b/web/src/components/AssistantChat/messages/AssistantMessage.tsx index 0b42ff66..ce922111 100644 --- a/web/src/components/AssistantChat/messages/AssistantMessage.tsx +++ b/web/src/components/AssistantChat/messages/AssistantMessage.tsx @@ -59,6 +59,7 @@ export function HappyAssistantMessage() { const durationMs = useAssistantState(({ message }) => (message.metadata.custom as Partial | undefined)?.durationMs) const usage = useAssistantState(({ message }) => (message.metadata.custom as Partial | undefined)?.usage) const messageModel = useAssistantState(({ message }) => (message.metadata.custom as Partial | undefined)?.model) + const turnCount = useAssistantState(({ message }) => (message.metadata.custom as Partial | undefined)?.turnCount) const hasMetadata = invokedAt != null || (typeof durationMs === 'number' && durationMs >= 0) @@ -100,6 +101,7 @@ export function HappyAssistantMessage() { durationMs={durationMs} usage={usage} model={messageModel ?? null} + turnCount={turnCount} className="mt-1" /> )} @@ -158,16 +160,26 @@ export function HappyAssistantMessage() { id={getConversationMessageAnchorId(messageId)} className={`${rootClass} ${copyText ? 'group/msg' : ''} scroll-mt-4`} > - - {showMetadata && ( - - )} +
+ + {showMetadata && ( + + )} +
) } @@ -193,6 +205,7 @@ export function HappyAssistantMessage() { durationMs={durationMs} usage={usage} model={messageModel ?? null} + turnCount={turnCount} className="mt-1" /> )} diff --git a/web/src/components/AssistantChat/messages/MessageMetadata.test.ts b/web/src/components/AssistantChat/messages/MessageMetadata.test.ts index 8ecefbe8..e06946cd 100644 --- a/web/src/components/AssistantChat/messages/MessageMetadata.test.ts +++ b/web/src/components/AssistantChat/messages/MessageMetadata.test.ts @@ -63,4 +63,79 @@ describe('buildMessageMetadataLabels', () => { expect(buildMessageMetadataLabels({ invokedAt: null }).some(p => p.startsWith('Invoke:'))).toBe(false) expect(buildMessageMetadataLabels({}).some(p => p.startsWith('Invoke:'))).toBe(false) }) + + // Proof of Invariance — single-turn inputs (turnCount omitted, or < 2) + // must produce byte-identical output to the pre-aggregate footer so + // existing single-turn cards do not regress visually. + it('single-turn input is byte-identical with or without turnCount=1', () => { + const base = { + invokedAt: 1700000000000, + durationMs: 1234, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 3, output_tokens: 19, service_tier: 'standard' } + } + const withoutTurn = buildMessageMetadataLabels(base) + const withTurnOne = buildMessageMetadataLabels({ ...base, turnCount: 1 }) + expect(withTurnOne).toEqual(withoutTurn) + }) + + // Byte-for-byte lock on the pre-aggregate label set. PR #555 introduced + // this exact shape; any regression to ordering, label strings, or token + // formatting would surface visually in single-turn cards. + it('pre-aggregate single-turn call produces the exact label sequence', () => { + const parts = buildMessageMetadataLabels({ + invokedAt: 1700000000000, + durationMs: 1234, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 3, output_tokens: 19, service_tier: 'standard' } + }) + // The Invoke value depends on the runner's timezone, so match its + // shape rather than a literal time string. The remaining labels are + // timezone-independent and locked exactly. + expect(parts).toHaveLength(4) + expect(parts[0]).toMatch(/^Invoke: \d{2}:\d{2}:\d{2}$/) + expect(parts.slice(1)).toEqual([ + 'Duration: 1.2s', + 'Model: claude-sonnet-4-6', + 'Usage: 22 billable tokens (3 in / 19 out)' + ]) + }) + + it('switches to Models/Total/N turns labels only when turnCount >= 2', () => { + const parts = buildMessageMetadataLabels({ + invokedAt: 1700000000000, + model: 'claude-sonnet-4-6, claude-haiku-4-5-20251001', + usage: { input_tokens: 100, output_tokens: 200, service_tier: 'standard' }, + turnCount: 3 + }) + expect(parts).toContain('Models: claude-sonnet-4-6, claude-haiku-4-5-20251001') + expect(parts.some(p => p.startsWith('Model:'))).toBe(false) + expect(parts).toContain('Total: 300 billable tokens (100 in / 200 out)') + expect(parts.some(p => p.startsWith('Usage:'))).toBe(false) + expect(parts).toContain('3 turns') + }) + + it('keeps the singular Model label when an aggregated group has only one distinct model', () => { + // Mid-session model switch is rare; the common multi-turn case is one + // model repeated across N turns. The label must stay singular then. + const parts = buildMessageMetadataLabels({ + invokedAt: 1700000000000, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 10, output_tokens: 20, service_tier: 'standard' }, + turnCount: 2 + }) + expect(parts).toContain('Model: claude-sonnet-4-6') + expect(parts.some(p => p.startsWith('Models:'))).toBe(false) + expect(parts).toContain('2 turns') + }) + + it('omits Duration on aggregated footers when durationMs is undefined', () => { + const parts = buildMessageMetadataLabels({ + invokedAt: 1700000000000, + model: 'claude-sonnet-4-6, claude-haiku-4-5-20251001', + usage: { input_tokens: 10, output_tokens: 20, service_tier: 'standard' }, + turnCount: 2 + }) + expect(parts.some(p => p.startsWith('Duration:'))).toBe(false) + }) }) diff --git a/web/src/components/AssistantChat/messages/MessageMetadata.tsx b/web/src/components/AssistantChat/messages/MessageMetadata.tsx index e6f43e6b..1b42becf 100644 --- a/web/src/components/AssistantChat/messages/MessageMetadata.tsx +++ b/web/src/components/AssistantChat/messages/MessageMetadata.tsx @@ -5,11 +5,22 @@ export type MessageMetadataProps = { durationMs?: number usage?: UsageData model?: string | null + /** + * Distinct turn count for the surrounding response group. Single-turn + * footers pass `undefined` (or any value < 2) so the existing + * `Invoke · Model · Usage` output is preserved byte-for-byte. + */ + turnCount?: number className?: string } -export function buildMessageMetadataLabels({ invokedAt, durationMs, usage, model }: Omit): string[] { +export function buildMessageMetadataLabels({ invokedAt, durationMs, usage, model, turnCount }: Omit): string[] { const parts: string[] = [] + // Aggregated footers represent a response group with multiple distinct + // turns. When the caller passes `turnCount >= 2` they have already + // dedup-joined `model` into a comma-separated list and summed `usage` + // across turns; we adjust the labels to reflect that. + const isAggregated = typeof turnCount === 'number' && turnCount >= 2 // Explicit nullish checks — `if (invokedAt)` would drop epoch 0, and // `if (durationMs)` would drop legitimate 0 ms turns. @@ -30,7 +41,9 @@ export function buildMessageMetadataLabels({ invokedAt, durationMs, usage, model const tier = usage?.service_tier const isStandardTier = tier?.toLowerCase() === 'standard' if (model) { - let label = `Model: ${model}` + // Pluralize the label when the caller has joined multiple model ids. + const modelLabel = isAggregated && model.includes(', ') ? 'Models' : 'Model' + let label = `${modelLabel}: ${model}` if (tier && !isStandardTier) label += ` (${tier})` parts.push(label) } else if (tier && !isStandardTier) { @@ -44,14 +57,19 @@ export function buildMessageMetadataLabels({ invokedAt, durationMs, usage, model // add a separate `Cache:` line. const total = usage.input_tokens + usage.output_tokens const formatToken = (n: number) => n.toLocaleString() - parts.push(`Usage: ${formatToken(total)} billable tokens (${formatToken(usage.input_tokens)} in / ${formatToken(usage.output_tokens)} out)`) + const usageLabel = isAggregated ? 'Total' : 'Usage' + parts.push(`${usageLabel}: ${formatToken(total)} billable tokens (${formatToken(usage.input_tokens)} in / ${formatToken(usage.output_tokens)} out)`) + } + + if (isAggregated) { + parts.push(`${turnCount} turns`) } return parts } -export function MessageMetadata({ invokedAt, durationMs, usage, model, className }: MessageMetadataProps) { - const parts = buildMessageMetadataLabels({ invokedAt, durationMs, usage, model }) +export function MessageMetadata({ invokedAt, durationMs, usage, model, turnCount, className }: MessageMetadataProps) { + const parts = buildMessageMetadataLabels({ invokedAt, durationMs, usage, model, turnCount }) if (parts.length === 0) return null return ( diff --git a/web/src/lib/assistant-runtime.test.ts b/web/src/lib/assistant-runtime.test.ts new file mode 100644 index 00000000..c5c01978 --- /dev/null +++ b/web/src/lib/assistant-runtime.test.ts @@ -0,0 +1,617 @@ +import { describe, expect, it } from 'vitest' +import { aggregateResponseGroups } from './assistant-runtime' +import type { AgentEventBlock, AgentTextBlock, CliOutputBlock, ToolCallBlock, UserTextBlock } from '@/chat/types' +import type { ToolGroupBlock, VisibleChatBlock } from '@/chat/toolGroups' + +// Minimal builders for VisibleChatBlock fixtures. Tests focus on metadata +// aggregation behavior across response groups; non-metadata fields default to +// inert values. + +function userText(id: string, overrides: Partial = {}): UserTextBlock { + return { + kind: 'user-text', + id, + localId: null, + createdAt: 0, + text: '', + ...overrides + } +} + +function agentText(id: string, overrides: Partial = {}): AgentTextBlock { + return { + kind: 'agent-text', + id, + localId: null, + createdAt: 0, + text: '', + ...overrides + } +} + +function toolCall(id: string, overrides: Partial = {}): ToolCallBlock { + return { + kind: 'tool-call', + id, + localId: null, + createdAt: 0, + tool: { + id, + name: 'Read', + state: 'completed', + input: {}, + createdAt: 0, + startedAt: null, + completedAt: null, + description: null + }, + children: [], + ...overrides + } +} + +function agentEvent(id: string, event: AgentEventBlock['event']): AgentEventBlock { + return { + kind: 'agent-event', + id, + createdAt: 0, + event + } +} + +function cliOutput(id: string, source: CliOutputBlock['source'], overrides: Partial = {}): CliOutputBlock { + return { + kind: 'cli-output', + id, + localId: null, + createdAt: 0, + text: '', + source, + ...overrides + } +} + +function toolGroup(id: string, tools: ToolCallBlock[], overrides: Partial = {}): ToolGroupBlock { + return { + kind: 'tool-group', + id, + createdAt: 0, + invokedAt: tools[0]?.invokedAt ?? null, + firstToolId: tools[0]?.id ?? id, + lastToolId: tools[tools.length - 1]?.id ?? id, + tools, + defaultOpen: false, + historyState: 'complete', + needsOlderHistory: false, + summary: { + totalTools: tools.length, + countsByKind: { read: 0, search: 0, command: 0, mutation: 0, web: 0, other: 0 }, + fileTargets: [], + commandTargets: [], + searchTargets: [], + urlTargets: [], + otherTargets: [], + errorCount: 0, + runningCount: 0, + pendingCount: 0 + }, + ...overrides + } +} + +describe('aggregateResponseGroups', () => { + it('1. sums usage and dedups model across distinct localIds in a single response group', () => { + // user (no aggregate) → agent-text L1 → tool-call L1 → tool-call L2 → agent-text L3 + // 3 distinct turns. Group's first visible block is the agent-text at L1. + const blocks: VisibleChatBlock[] = [ + userText('u1'), + agentText('a1', { + localId: 'L1', + invokedAt: 100, + durationMs: 1234, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 10, output_tokens: 20, service_tier: 'standard' } + }), + toolCall('t1', { localId: 'L1' }), + toolCall('t2', { + localId: 'L2', + invokedAt: 200, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 5, output_tokens: 7, service_tier: 'standard' } + }), + agentText('a3', { + localId: 'L3', + invokedAt: 300, + durationMs: 5678, + model: 'claude-haiku-4-5-20251001', + usage: { input_tokens: 3, output_tokens: 11, service_tier: 'standard' } + }) + ] + + const aggregates = aggregateResponseGroups(blocks) + const meta = aggregates.get('a1') + expect(meta).toBeDefined() + expect(meta?.turnCount).toBe(3) + // input/output sums across the three distinct localIds. + expect(meta?.usage?.input_tokens).toBe(10 + 5 + 3) + expect(meta?.usage?.output_tokens).toBe(20 + 7 + 11) + // Model dedup preserves first-seen order. "claude-sonnet-4-6" appears + // twice (L1, L2) and must not be duplicated. + expect(meta?.model).toBe('claude-sonnet-4-6, claude-haiku-4-5-20251001') + // Invoke time = first turn (regression guard for the user-reported + // disappearance after PR #555). + expect(meta?.invokedAt).toBe(100) + // Duration is intentionally undefined so the library does not surface + // the first turn's stale duration on the aggregated card. + expect(meta?.durationMs).toBeUndefined() + // Only the group's first visible block carries an aggregate entry. + expect(aggregates.has('u1')).toBe(false) + expect(aggregates.has('t1')).toBe(false) + expect(aggregates.has('t2')).toBe(false) + expect(aggregates.has('a3')).toBe(false) + }) + + it('2. leaves a single-turn group untouched so the existing footer renders unchanged', () => { + // localId 'L1' shared across multiple blocks → still one turn. + const blocks: VisibleChatBlock[] = [ + userText('u1'), + agentText('a1', { + localId: 'L1', + invokedAt: 42, + durationMs: 999, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 3, output_tokens: 19, service_tier: 'standard' } + }), + toolCall('t1', { localId: 'L1' }) + ] + + const aggregates = aggregateResponseGroups(blocks) + // No entry → upstream callback emits the original per-block metadata. + expect(aggregates.size).toBe(0) + }) + + it('3. splits response groups on each user-text boundary', () => { + // user → agent L1 → tool L1 → user → agent L2 → agent L3 + const blocks: VisibleChatBlock[] = [ + userText('u1'), + agentText('a1', { + localId: 'L1', + invokedAt: 100, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 10, output_tokens: 20, service_tier: 'standard' } + }), + toolCall('t1', { localId: 'L1' }), + userText('u2'), + agentText('a2', { + localId: 'L2', + invokedAt: 200, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 4, output_tokens: 8, service_tier: 'standard' } + }), + agentText('a3', { + localId: 'L3', + invokedAt: 300, + model: 'claude-haiku-4-5-20251001', + usage: { input_tokens: 5, output_tokens: 7, service_tier: 'standard' } + }) + ] + + const aggregates = aggregateResponseGroups(blocks) + // First group is a single turn → no entry. + expect(aggregates.has('a1')).toBe(false) + // Second group spans L2 + L3, first visible block is a2. + const meta2 = aggregates.get('a2') + expect(meta2?.turnCount).toBe(2) + expect(meta2?.usage?.input_tokens).toBe(9) + expect(meta2?.usage?.output_tokens).toBe(15) + expect(meta2?.model).toBe('claude-sonnet-4-6, claude-haiku-4-5-20251001') + expect(meta2?.invokedAt).toBe(200) + }) + + it('4. preserves first-seen order when dedup yields two distinct models in one group', () => { + const blocks: VisibleChatBlock[] = [ + userText('u1'), + agentText('a1', { + localId: 'L1', + invokedAt: 100, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 1, output_tokens: 1, service_tier: 'standard' } + }), + agentText('a2', { + localId: 'L2', + invokedAt: 200, + model: 'claude-haiku-4-5-20251001', + usage: { input_tokens: 1, output_tokens: 1, service_tier: 'standard' } + }) + ] + + const aggregates = aggregateResponseGroups(blocks) + const meta = aggregates.get('a1') + expect(meta?.model).toBe('claude-sonnet-4-6, claude-haiku-4-5-20251001') + }) + + it('5. falls back to a (model, usage) fingerprint to count turns when localId is null', () => { + // Claude code spawn sessions today never stamp `localId`; all + // blocks emitted in one Claude SDK message carry an identical + // `usage` object instead. We dedup by that fingerprint so a + // single turn does not over-count when its blocks repeat. + const turn1Usage = { input_tokens: 1, output_tokens: 2, service_tier: 'standard' as const } + const turn2Usage = { input_tokens: 4, output_tokens: 8, service_tier: 'standard' as const } + const turn3Usage = { input_tokens: 16, output_tokens: 32, service_tier: 'standard' as const } + const blocks: VisibleChatBlock[] = [ + userText('u1'), + // turn 1: thinking + tool_use share one usage object + agentText('a1', { localId: null, invokedAt: 100, model: 'claude-sonnet-4-6', usage: turn1Usage }), + toolCall('t1', { localId: null, invokedAt: 105, model: 'claude-sonnet-4-6', usage: turn1Usage }), + // turn 2: a different usage object -> new turn + agentText('a2', { localId: null, invokedAt: 200, model: 'claude-sonnet-4-6', usage: turn2Usage }), + // turn 3: a different model + usage -> new turn + agentText('a3', { localId: null, invokedAt: 300, model: 'claude-haiku-4-5-20251001', usage: turn3Usage }) + ] + + const aggregates = aggregateResponseGroups(blocks) + const meta = aggregates.get('a1') + expect(meta?.turnCount).toBe(3) + // sum across the three distinct turns + expect(meta?.usage?.input_tokens).toBe(21) + expect(meta?.usage?.output_tokens).toBe(42) + expect(meta?.model).toBe('claude-sonnet-4-6, claude-haiku-4-5-20251001') + }) + + it('5a. fingerprint dedup is ordering-based: identical adjacent blocks count as one turn, identical non-adjacent blocks count as separate turns', () => { + // Each Claude SDK message emits multiple blocks that share an + // identical usage object — adjacent blocks must collapse to one + // turn. But two distinct SDK messages occasionally happen to + // produce the same (model, usage) fingerprint when separated by + // a third turn with different totals. A Set-based dedup would + // collapse those non-adjacent matches into one turn and under- + // count; an ordering-based dedup only merges adjacent matches. + const sharedUsage = { input_tokens: 5, output_tokens: 7, service_tier: 'standard' as const } + const middleUsage = { input_tokens: 11, output_tokens: 13, service_tier: 'standard' as const } + const blocks: VisibleChatBlock[] = [ + userText('u1'), + // Turn 1 emits two blocks sharing the same usage fingerprint. + agentText('a1', { localId: null, invokedAt: 100, model: 'claude-sonnet-4-6', usage: sharedUsage }), + toolCall('t1', { localId: null, invokedAt: 101, model: 'claude-sonnet-4-6', usage: sharedUsage }), + // Turn 2: different fingerprint. + agentText('a2', { localId: null, invokedAt: 200, model: 'claude-sonnet-4-6', usage: middleUsage }), + // Turn 3 happens to repeat turn 1's (model, usage) fingerprint. + agentText('a3', { localId: null, invokedAt: 300, model: 'claude-sonnet-4-6', usage: sharedUsage }) + ] + + const aggregates = aggregateResponseGroups(blocks) + const meta = aggregates.get('a1') + // Three distinct turns: adjacent fingerprint match collapses (a1+t1), + // non-adjacent fingerprint match does not (a1 vs a3). + expect(meta?.turnCount).toBe(3) + // Sum: shared + middle + shared. + expect(meta?.usage?.input_tokens).toBe(5 + 11 + 5) + expect(meta?.usage?.output_tokens).toBe(7 + 13 + 7) + }) + + it("5b. skips chunk blocks without model or usage so they do not inflate the turn count", () => { + // hapi's hub stores tool_result chunks as separate agent-role + // messages with no `model`, no `usage`, and `localId=null`. + // They share an SDK turn with the preceding tool_use but the + // fingerprint signal is missing, so the aggregator must skip + // them rather than inflate the turn count. + const turn1Usage = { input_tokens: 3, output_tokens: 8, service_tier: 'standard' as const } + const turn2Usage = { input_tokens: 1, output_tokens: 5, service_tier: 'standard' as const } + const blocks: VisibleChatBlock[] = [ + userText('u1'), + agentText('a1', { localId: null, invokedAt: 100, model: 'claude-sonnet-4-6', usage: turn1Usage }), + toolCall('t1', { localId: null, invokedAt: 101, model: 'claude-sonnet-4-6', usage: turn1Usage }), + // tool_result chunk: no model, no usage + agentText('a2_result', { localId: null, invokedAt: 102 }), + // final turn with a different usage + agentText('a3_final', { localId: null, invokedAt: 200, model: 'claude-sonnet-4-6', usage: turn2Usage }) + ] + + const aggregates = aggregateResponseGroups(blocks) + const meta = aggregates.get('a1') + expect(meta?.turnCount).toBe(2) + expect(meta?.usage?.input_tokens).toBe(4) + expect(meta?.usage?.output_tokens).toBe(13) + }) + + it('6. ends a response group at an agent-event boundary (library chunk flush)', () => { + const blocks: VisibleChatBlock[] = [ + userText('u1'), + agentText('a1', { + localId: 'L1', + invokedAt: 100, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 10, output_tokens: 20, service_tier: 'standard' } + }), + // limit-reached splits the library's chunk; the next assistant + // block starts a new card and therefore a new response group. + agentEvent('e1', { type: 'limit-reached', endsAt: 0, limitType: '5h' }), + agentText('a2', { + localId: 'L2', + invokedAt: 200, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 4, output_tokens: 8, service_tier: 'standard' } + }), + agentText('a3', { + localId: 'L3', + invokedAt: 300, + model: 'claude-haiku-4-5-20251001', + usage: { input_tokens: 1, output_tokens: 1, service_tier: 'standard' } + }) + ] + + const aggregates = aggregateResponseGroups(blocks) + // Pre-event group is a single turn → no entry. + expect(aggregates.has('a1')).toBe(false) + // Post-event group has two turns starting at a2. + const meta2 = aggregates.get('a2') + expect(meta2?.turnCount).toBe(2) + expect(meta2?.invokedAt).toBe(200) + expect(meta2?.usage?.input_tokens).toBe(5) + expect(meta2?.usage?.output_tokens).toBe(9) + }) + + it('does not aggregate user-role cli-output blocks (they do not belong to a response group)', () => { + // Defensive: a cli-output with source='user' is rendered as a user + // role message by the converter, so it must not be folded into an + // assistant response group nor act as the group's first block. + const blocks: VisibleChatBlock[] = [ + cliOutput('c1', 'user'), + agentText('a1', { + localId: 'L1', + invokedAt: 100, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 1, output_tokens: 2, service_tier: 'standard' } + }), + agentText('a2', { + localId: 'L2', + invokedAt: 200, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 3, output_tokens: 4, service_tier: 'standard' } + }) + ] + + const aggregates = aggregateResponseGroups(blocks) + expect(aggregates.has('c1')).toBe(false) + // a1 is the first visible block of the response group spanning L1+L2. + const meta = aggregates.get('a1') + expect(meta?.turnCount).toBe(2) + }) + + it('counts a tool-group whose tools all belong to a single turn as one turn', () => { + // `buildVisibleChatBlocks` merges adjacent eligible tool-calls into a + // single `tool-group` block. When every underlying tool shares one + // turn (same localId, same usage), the aggregator must collapse them + // to one turn — the group is not its own turn boundary. + const turn1Usage = { input_tokens: 7, output_tokens: 11, service_tier: 'standard' as const } + const turn2Usage = { input_tokens: 2, output_tokens: 9, service_tier: 'standard' as const } + const tool1 = toolCall('t1', { localId: 'L1', invokedAt: 100, model: 'claude-sonnet-4-6', usage: turn1Usage }) + const tool2 = toolCall('t2', { localId: 'L1', invokedAt: 105, model: 'claude-sonnet-4-6', usage: turn1Usage }) + const blocks: VisibleChatBlock[] = [ + userText('u1'), + toolGroup('g1', [tool1, tool2]), + agentText('a1', { + localId: 'L2', + invokedAt: 200, + model: 'claude-sonnet-4-6', + usage: turn2Usage + }) + ] + + const aggregates = aggregateResponseGroups(blocks) + // g1 is the group's first visible block. + const meta = aggregates.get('g1') + expect(meta?.turnCount).toBe(2) + // input/output sums across the tool-group turn and the agent-text turn. + expect(meta?.usage?.input_tokens).toBe(7 + 2) + expect(meta?.usage?.output_tokens).toBe(11 + 9) + expect(meta?.invokedAt).toBe(100) + }) + + it('skips an empty tool-group block without throwing or inflating the turn count', () => { + // Defensive: a malformed tool-group with zero tools must degrade to + // null rather than crash. The surrounding response group continues. + const blocks: VisibleChatBlock[] = [ + userText('u1'), + toolGroup('g0', []), + agentText('a1', { + localId: 'L1', + invokedAt: 100, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 1, output_tokens: 2, service_tier: 'standard' } + }), + agentText('a2', { + localId: 'L2', + invokedAt: 200, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 3, output_tokens: 4, service_tier: 'standard' } + }) + ] + + const aggregates = aggregateResponseGroups(blocks) + // g0 is the first visible block of the response group (no entry — the + // empty group contributes no turn, and the boundary remains at the + // user block). + const meta = aggregates.get('g0') + expect(meta?.turnCount).toBe(2) + expect(meta?.usage?.input_tokens).toBe(4) + expect(meta?.usage?.output_tokens).toBe(6) + }) + + it('counts every turn inside a tool-group that spans multiple assistant turns', () => { + // Regression guard for the case the helper expands: a single + // `tool-group` block may wrap tool-calls from two distinct turns. + const turn1Usage = { input_tokens: 7, output_tokens: 11, service_tier: 'standard' as const } + const turn2Usage = { input_tokens: 2, output_tokens: 9, service_tier: 'standard' as const } + const tool1 = toolCall('t1', { localId: 'L1', invokedAt: 100, model: 'claude-sonnet-4-6', usage: turn1Usage }) + const tool2 = toolCall('t2', { localId: 'L2', invokedAt: 110, model: 'claude-haiku-4-5-20251001', usage: turn2Usage }) + const blocks: VisibleChatBlock[] = [ + userText('u1'), + toolGroup('g1', [tool1, tool2]) + ] + + const aggregates = aggregateResponseGroups(blocks) + const meta = aggregates.get('g1') + expect(meta?.turnCount).toBe(2) + expect(meta?.usage?.input_tokens).toBe(7 + 2) + expect(meta?.usage?.output_tokens).toBe(11 + 9) + expect(meta?.model).toBe('claude-sonnet-4-6, claude-haiku-4-5-20251001') + expect(meta?.invokedAt).toBe(100) + }) + + it('keeps two fingerprint-mode turns distinct when only token totals coincide (different createdAt)', () => { + // Two consecutive SDK messages happen to report identical + // `(model, usage)` (rare but possible on very short turns). Different + // createdAt values keep their fingerprints distinct so the aggregator + // counts both turns instead of dedupping them. + const usage = { input_tokens: 3, output_tokens: 5, service_tier: 'standard' as const } + const blocks: VisibleChatBlock[] = [ + userText('u1'), + agentText('a1', { localId: null, createdAt: 1000, invokedAt: 100, model: 'claude-sonnet-4-6', usage }), + agentText('a2', { localId: null, createdAt: 2000, invokedAt: 200, model: 'claude-sonnet-4-6', usage }) + ] + + const aggregates = aggregateResponseGroups(blocks) + const meta = aggregates.get('a1') + expect(meta?.turnCount).toBe(2) + expect(meta?.usage?.input_tokens).toBe(3 + 3) + expect(meta?.usage?.output_tokens).toBe(5 + 5) + }) + + it('preserves a 0 sum on cache token fields instead of folding it to undefined', () => { + // Both turns omit cache_creation/cache_read entirely. The aggregator + // must not invent a 0 either — the field remains undefined when no + // turn carried a value. Conversely, if one turn reports 0 explicitly + // we keep 0 in the sum (covered by the next case). + const blocks: VisibleChatBlock[] = [ + userText('u1'), + agentText('a1', { + localId: 'L1', + invokedAt: 100, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 1, output_tokens: 2, service_tier: 'standard' } + }), + agentText('a2', { + localId: 'L2', + invokedAt: 200, + model: 'claude-sonnet-4-6', + usage: { input_tokens: 3, output_tokens: 4, service_tier: 'standard' } + }) + ] + + const aggregates = aggregateResponseGroups(blocks) + const meta = aggregates.get('a1') + expect(meta?.usage?.cache_creation_input_tokens).toBeUndefined() + expect(meta?.usage?.cache_read_input_tokens).toBeUndefined() + }) + + it('keeps an explicit 0 in a cache token sum (does not coerce 0 → undefined via ||)', () => { + // Regression for the `(a ?? 0) + (b ?? 0) || undefined` pattern: + // 0 + 0 must remain 0 when at least one turn carried the field + // explicitly, otherwise downstream surfaces lose the signal. + const blocks: VisibleChatBlock[] = [ + userText('u1'), + agentText('a1', { + localId: 'L1', + invokedAt: 100, + model: 'claude-sonnet-4-6', + usage: { + input_tokens: 1, + output_tokens: 2, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + service_tier: 'standard' + } + }), + agentText('a2', { + localId: 'L2', + invokedAt: 200, + model: 'claude-sonnet-4-6', + usage: { + input_tokens: 3, + output_tokens: 4, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + service_tier: 'standard' + } + }) + ] + + const aggregates = aggregateResponseGroups(blocks) + const meta = aggregates.get('a1') + expect(meta?.usage?.cache_creation_input_tokens).toBe(0) + expect(meta?.usage?.cache_read_input_tokens).toBe(0) + }) + + it('keeps a partial cache token value when only one turn carries it', () => { + // The other turn contributes 0 (treated as the missing-side default) + // and the sum equals the side that did carry a value. + const blocks: VisibleChatBlock[] = [ + userText('u1'), + agentText('a1', { + localId: 'L1', + invokedAt: 100, + model: 'claude-sonnet-4-6', + usage: { + input_tokens: 1, + output_tokens: 2, + cache_creation_input_tokens: 50, + service_tier: 'standard' + } + }), + agentText('a2', { + localId: 'L2', + invokedAt: 200, + model: 'claude-sonnet-4-6', + usage: { + input_tokens: 3, + output_tokens: 4, + service_tier: 'standard' + } + }) + ] + + const aggregates = aggregateResponseGroups(blocks) + const meta = aggregates.get('a1') + expect(meta?.usage?.cache_creation_input_tokens).toBe(50) + expect(meta?.usage?.cache_read_input_tokens).toBeUndefined() + }) + + it('does not surface cache_read/cache_creation tokens via aggregation (sums them but display ignores)', () => { + // We still sum every UsageData field so the aggregate is structurally + // complete, but the visible label only consumes input/output. This + // lets future surfaces decide independently. + const blocks: VisibleChatBlock[] = [ + userText('u1'), + agentText('a1', { + localId: 'L1', + invokedAt: 100, + model: 'claude-sonnet-4-6', + usage: { + input_tokens: 1, + output_tokens: 2, + cache_creation_input_tokens: 100, + cache_read_input_tokens: 50, + service_tier: 'standard' + } + }), + agentText('a2', { + localId: 'L2', + invokedAt: 200, + model: 'claude-sonnet-4-6', + usage: { + input_tokens: 3, + output_tokens: 4, + cache_creation_input_tokens: 200, + cache_read_input_tokens: 50, + service_tier: 'standard' + } + }) + ] + + const aggregates = aggregateResponseGroups(blocks) + const meta = aggregates.get('a1') + expect(meta?.usage?.cache_creation_input_tokens).toBe(300) + expect(meta?.usage?.cache_read_input_tokens).toBe(100) + }) +}) diff --git a/web/src/lib/assistant-runtime.ts b/web/src/lib/assistant-runtime.ts index 71ccb710..9c3d542e 100644 --- a/web/src/lib/assistant-runtime.ts +++ b/web/src/lib/assistant-runtime.ts @@ -11,6 +11,19 @@ import type { AgentEvent, ToolCallBlock } from '@/chat/types' import type { ToolGroupBlock, VisibleChatBlock } from '@/chat/toolGroups' import type { AttachmentMetadata, MessageStatus as HappyMessageStatus, Session } from '@/types/api' +/** + * Aggregated metadata for a multi-turn response group, surfaced on the + * group's first visible block so the `@assistant-ui/react` converter + * preserves it after joining adjacent assistant messages. + */ +export type AggregatedAssistantMeta = { + usage?: UsageData + model: string | null + invokedAt: number | null + durationMs: undefined + turnCount: number +} + export type HappyChatMessageMetadata = { kind: 'user' | 'assistant' | 'tool' | 'event' | 'cli-output' | 'codex-review' status?: HappyMessageStatus @@ -25,6 +38,12 @@ export type HappyChatMessageMetadata = { usage?: UsageData model?: string | null review?: CodexReview + /** + * Distinct turn count when this block carries an aggregated response + * group footer. Single-turn blocks omit this field so the existing + * per-message footer is rendered unchanged. + */ + turnCount?: number } function formatCodexReviewText(review: CodexReview): string { @@ -49,6 +68,214 @@ function formatCodexReviewText(review: CodexReview): string { return lines.join('\n') } +type VisibleChatBlockRole = 'user' | 'assistant' | 'system' + +/** + * Mirror the role assignment used by `toThreadMessageLike` so response + * group boundaries (the `@assistant-ui/react` converter joins adjacent + * assistant-role messages only) stay consistent with what the library + * actually flushes as one card. + */ +function visibleBlockRole(block: VisibleChatBlock): VisibleChatBlockRole { + if (block.kind === 'user-text') return 'user' + if (block.kind === 'agent-event') return 'system' + if (block.kind === 'cli-output') return block.source === 'user' ? 'user' : 'assistant' + return 'assistant' +} + +type TurnSource = { + localId: string | null + invokedAt: number | null + durationMs: number | undefined + model: string | null + usage: UsageData | undefined + createdAt: number +} + +// Return one turn source per claude-SDK message that the visible block +// represents. `tool-group` is a derived view: `buildVisibleChatBlocks` merges +// adjacent eligible tool-calls without checking that they share a turn, so +// each underlying tool-call contributes its own source. Other assistant-role +// kinds map to a single source. +function turnSourcesFromBlock(block: VisibleChatBlock): TurnSource[] { + if (block.kind === 'tool-group') { + return block.tools.map((tool) => ({ + localId: tool.localId, + invokedAt: tool.invokedAt ?? null, + durationMs: tool.durationMs, + model: tool.model ?? null, + usage: tool.usage, + createdAt: tool.createdAt + })) + } + if ( + block.kind === 'agent-text' + || block.kind === 'agent-reasoning' + || block.kind === 'cli-output' + || block.kind === 'tool-call' + ) { + return [{ + localId: block.localId, + invokedAt: block.invokedAt ?? null, + durationMs: block.durationMs, + model: block.model ?? null, + usage: block.usage, + createdAt: block.createdAt + }] + } + return [] +} + +/** + * Fallback turn key for environments where the CLI does not stamp a + * non-null `localId` on each turn's blocks (e.g. claude code spawn + * sessions today). The key combines `(model, usage totals, createdAt)`: + * the reducer copies `msg.createdAt` onto every derived `ChatBlock`, so + * blocks from the same SDK message share createdAt and collapse to one + * turn, while blocks from different SDK messages stay distinct even + * when their `(model, usage)` happens to coincide. Limitation: two SDK + * messages stamped at the same wall-clock millisecond would still + * collide, but the hub stamp resolution makes that vanishingly rare. + */ +function turnFingerprint( + model: string | null, + usage: UsageData | undefined, + createdAt: number +): string { + if (!usage) return `m=${model ?? ''}|u=|c=${createdAt}` + return [ + `m=${model ?? ''}`, + `i=${usage.input_tokens}`, + `o=${usage.output_tokens}`, + `cc=${usage.cache_creation_input_tokens ?? ''}`, + `cr=${usage.cache_read_input_tokens ?? ''}`, + `t=${usage.service_tier ?? ''}`, + `c=${createdAt}` + ].join('|') +} + +/** + * Sum two optional token counters. Returns `undefined` only when both + * operands are absent; an explicit `0` on either side participates in + * the sum (so a `0 + 0` total stays `0` instead of collapsing to + * `undefined` via JavaScript's falsy `||`). + */ +function sumOptional(a: number | undefined, b: number | undefined): number | undefined { + if (a === undefined && b === undefined) return undefined + return (a ?? 0) + (b ?? 0) +} + +function addUsage(target: UsageData, addend: UsageData): UsageData { + return { + input_tokens: target.input_tokens + addend.input_tokens, + output_tokens: target.output_tokens + addend.output_tokens, + cache_creation_input_tokens: sumOptional( + target.cache_creation_input_tokens, + addend.cache_creation_input_tokens + ), + cache_read_input_tokens: sumOptional( + target.cache_read_input_tokens, + addend.cache_read_input_tokens + ), + // service_tier dedup follows the rule documented in the plan: pick + // the first turn's tier when the group is mixed. We keep it on the + // target so downstream label logic sees a stable value. + service_tier: target.service_tier + } +} + +/** + * Walk the visible block list, identify response groups (runs of + * assistant-role blocks separated by user-text / agent-event / + * user-source cli-output boundaries), and return a map keyed by the + * id of each group's first visible block whose value is the summed + * metadata for that group. Only groups spanning two or more distinct + * turns produce an entry so single-turn cards remain byte-for-byte + * unchanged at the call site. + */ +export function aggregateResponseGroups( + blocks: readonly VisibleChatBlock[] +): Map { + const aggregates = new Map() + let groupFirstBlockId: string | null = null + // Ordering-based turn dedup: we compare each block's turn key against + // the immediately previous turn's key. A `Set` of all seen keys would + // collapse a third turn whose `(model, usage)` fingerprint happens to + // match a non-adjacent earlier turn — claude code spawn sessions + // legitimately produce repeated fingerprints when token totals + // coincide, and merging them under-counts the visible turn count. + let prevTurnKey: string | null = null + const seenModels: string[] = [] + let groupInvokedAt: number | null = null + let groupUsage: UsageData | undefined + let groupTurnCount = 0 + + const flush = () => { + if (groupFirstBlockId !== null && groupTurnCount >= 2) { + const joinedModel = seenModels.length > 0 ? seenModels.join(', ') : null + aggregates.set(groupFirstBlockId, { + usage: groupUsage, + model: joinedModel, + invokedAt: groupInvokedAt, + durationMs: undefined, + turnCount: groupTurnCount + }) + } + groupFirstBlockId = null + prevTurnKey = null + seenModels.length = 0 + groupInvokedAt = null + groupUsage = undefined + groupTurnCount = 0 + } + + for (const block of blocks) { + const role = visibleBlockRole(block) + if (role !== 'assistant') { + // Boundary: close the open group, if any. + flush() + continue + } + + if (groupFirstBlockId === null) { + groupFirstBlockId = block.id + } + + for (const turn of turnSourcesFromBlock(block)) { + // Prefer the CLI-stamped `localId` when present. When it is null + // (today's claude code spawn flow) fall back to a fingerprint + // built from `(model, usage totals, createdAt)` — see + // `turnFingerprint` for the createdAt rationale. + const turnKey = turn.localId !== null + ? `id:${turn.localId}` + : `fp:${turnFingerprint(turn.model, turn.usage, turn.createdAt)}` + + // Skip blocks with no turn signal at all (no localId, no usage, + // no model) so they don't inflate the turn count. + if (turn.localId === null && !turn.usage && !turn.model) continue + + // Only the immediately previous turn's key dedups — see prevTurnKey + // comment above. Non-adjacent matches keep their separate counts. + if (turnKey === prevTurnKey) continue + prevTurnKey = turnKey + + groupTurnCount += 1 + if (turn.invokedAt != null && (groupInvokedAt === null || turn.invokedAt < groupInvokedAt)) { + groupInvokedAt = turn.invokedAt + } + if (turn.model && !seenModels.includes(turn.model)) { + seenModels.push(turn.model) + } + if (turn.usage) { + groupUsage = groupUsage ? addUsage(groupUsage, turn.usage) : { ...turn.usage } + } + } + } + + flush() + return aggregates +} + function toThreadMessageLike(block: VisibleChatBlock): ThreadMessageLike { if (block.kind === 'user-text') { const messageId = `user:${block.id}` @@ -309,10 +536,44 @@ export function useHappyRuntime(props: { }) { const isRunning = props.isRunning ?? props.session.thinking + // Compute response-group aggregates once per block list so we can + // inject the summed metadata onto each group's first visible block. + // The library's `joinExternalMessages` only preserves + // `metadata.custom` from the first block of a joined chunk, so this + // is the surface that survives the join. + const aggregates = useMemo( + () => aggregateResponseGroups(props.blocks), + [props.blocks] + ) + + const convertBlock = useCallback( + (block: VisibleChatBlock): ThreadMessageLike => { + const message = toThreadMessageLike(block) + const aggregate = aggregates.get(block.id) + if (!aggregate) return message + const existing = message.metadata?.custom as HappyChatMessageMetadata | undefined + return { + ...message, + metadata: { + ...message.metadata, + custom: { + ...(existing ?? { kind: 'assistant' }), + usage: aggregate.usage, + model: aggregate.model, + invokedAt: aggregate.invokedAt, + durationMs: aggregate.durationMs, + turnCount: aggregate.turnCount + } satisfies HappyChatMessageMetadata + } + } + }, + [aggregates] + ) + // Use cached message converter for performance optimization // This prevents re-converting all messages on every render const convertedMessages = useExternalMessageConverter({ - callback: toThreadMessageLike, + callback: convertBlock, messages: props.blocks as VisibleChatBlock[], isRunning, })