From 09a433ab917aec3775778ebe6d494a53d4af980e Mon Sep 17 00:00:00 2001 From: Eric Lee Date: Wed, 19 Aug 2026 07:54:40 -0700 Subject: [PATCH] fix(tokens): count a subagent's tokens as the session's, and cover the fold MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The token read-outs undercounted every run that fanned out. A child's `assistant/message` usage went onto the child's own row — the agents overlay, and the `Done (3 tool uses · 1.2k tokens)` a delegation settles with — and stopped there. The session odometer behind `tokens: N in / M out`, `/status` and the context read-out never heard about it, so a turn that delegated most of its work reported the sliver the parent did itself: tokens: 1.1k in / 84 out ← parent only tokens: 42.6k in / 3.1k out ← what the session actually spent It is one bill. The two halves of the bus are already disjoint — `bindAgent` routes by session identity and `onChildSessionEvent` drops anything the lifecycle never announced as a child — so folding the child's step in counts it exactly once, and a sibling top-level session still cannot leak in. Also lands the test coverage for the fold that #23 carried in without it: `statsFromUsage` and its precedence behind a CostSnapshot, the price-free `usage` path through `message.complete`, and the rider a resume replays out of the log. Co-Authored-By: Claude Opus 5 --- src/__tests__/harnessClient.test.ts | 40 +++++++++++++++++++ src/__tests__/sessionStats.test.ts | 62 +++++++++++++++++++++++++++-- src/harness/client.ts | 12 ++++++ 3 files changed, 111 insertions(+), 3 deletions(-) diff --git a/src/__tests__/harnessClient.test.ts b/src/__tests__/harnessClient.test.ts index c692be7..c8e6032 100644 --- a/src/__tests__/harnessClient.test.ts +++ b/src/__tests__/harnessClient.test.ts @@ -564,6 +564,46 @@ describe('HarnessGatewayClient', () => { }) }) + it('bills the child\'s tokens to the session odometer, not just its own row', async () => { + const w = await delegating() + + // The parent's own step… + w.fire('assistant/message', { + message: { content: [{ text: 'delegating', type: 'text' }], id: 'p1', role: 'assistant', source: { kind: 'model' } }, + step: 1, + turn: 1, + usage: { inputTokens: 100, outputTokens: 10 } + }) + // …and the child's, which used far more than the parent did. Both are + // the same bill; only the child's row used to hear about the second. + w.fireFrom(CHILD, 'assistant/message', { + message: { content: [{ text: 'done', type: 'text' }], id: 'm1', role: 'assistant', source: { kind: 'model' } }, + step: 1, + turn: 1, + usage: { cacheReadTokens: 400, inputTokens: 500, outputTokens: 60 } + }) + + await expect(w.client.request('session.usage', {})).resolves.toMatchObject({ + calls: 2, + input: 1000, + output: 70, + total: 1070 + }) + }) + + it('ignores a sibling session that was never announced as a child', async () => { + const w = await delegating() + + w.fireFrom('some-other-session', 'assistant/message', { + message: { content: [{ text: 'not ours', type: 'text' }], id: 'x1', role: 'assistant', source: { kind: 'model' } }, + step: 1, + turn: 1, + usage: { inputTokens: 999, outputTokens: 999 } + }) + + await expect(w.client.request('session.usage', {})).resolves.toMatchObject({ calls: 0, input: 0, output: 0 }) + }) + const delegate = async ( w: ReturnType, { callId = 'd1', resultText = 'the child answer' }: { callId?: string; resultText?: string } = {} diff --git a/src/__tests__/sessionStats.test.ts b/src/__tests__/sessionStats.test.ts index 88d0f31..0eb0180 100644 --- a/src/__tests__/sessionStats.test.ts +++ b/src/__tests__/sessionStats.test.ts @@ -2,8 +2,10 @@ * Session-stats line under the composer — the deleted REPL's bottom toolbar: * `provider · model · cwd · turns: N · tokens: X in / Y out · cost $C`. * - * Covers the CostSnapshot fold, the width-aware line builder, the - * message.complete → ui.sessionStats patch, and the component render gate. + * Covers both folds — the priced CostSnapshot and the price-free `usage` + * odometer a backend like the dsh harness sends instead — plus their + * precedence, the width-aware line builder, the message.complete → + * ui.sessionStats patch, and the component render gate. */ import { PassThrough } from 'node:stream' @@ -18,7 +20,7 @@ import { resetTurnState } from '../app/turnStore.js' import { getUiState, patchUiState, resetUiState } from '../app/uiStore.js' import { SessionStatsLine } from '../components/sessionStatsLine.js' import type { CostSnapshot } from '../gatewayTypes.js' -import { buildSessionStatsLine, statsFromCostSnapshot, ZERO_SESSION_STATS } from '../lib/sessionStats.js' +import { buildSessionStatsLine, statsFromCostSnapshot, statsFromUsage, ZERO_SESSION_STATS } from '../lib/sessionStats.js' import { stripAnsi } from '../lib/text.js' import type { Msg } from '../types.js' @@ -51,6 +53,23 @@ describe('statsFromCostSnapshot', () => { }) }) +describe('statsFromUsage', () => { + it('reads the cumulative odometer a price-free backend sends instead', () => { + // The harness folds cache reads/writes into `input` before it ships them, + // which is already what inputTokens means. + expect(statsFromUsage({ calls: 4, input: 33189, output: 622, total: 33811 }, 3)).toEqual({ + costUsd: 0, + inputTokens: 33189, + outputTokens: 622, + turns: 3 + }) + }) + + it('reports spend only when the backend priced it', () => { + expect(statsFromUsage({ calls: 1, cost_usd: 0.0048, input: 10, output: 2, total: 12 }, 1).costUsd).toBe(0.0048) + }) +}) + describe('buildSessionStatsLine', () => { const HOME = process.env.HOME @@ -176,6 +195,43 @@ describe('message.complete → ui.sessionStats', () => { }) }) + it('folds a price-free backend\'s usage odometer when there is no snapshot', () => { + const onEvent = createGatewayEventHandler(buildCtx()) + + // The dsh harness has no price table, so it publishes token totals here + // and never a CostSnapshot. This read `tokens: 0 in / 0 out` all session. + onEvent({ + payload: { session_turns: 1, text: 'done', usage: { calls: 2, input: 33189, output: 622, total: 33811 } }, + type: 'message.complete' + } as any) + + expect(getUiState().sessionStats).toEqual({ costUsd: 0, inputTokens: 33189, outputTokens: 622, turns: 1 }) + }) + + it('prefers the snapshot over usage when a backend sends both', () => { + const onEvent = createGatewayEventHandler(buildCtx()) + + // On a snapshot-carrying backend `usage` describes the last call, not the + // session — the subagent-inclusive snapshot is the one to believe. + onEvent({ + payload: { cost: SNAPSHOT, session_turns: 1, usage: { calls: 1, input: 12, output: 3, total: 15 } }, + type: 'message.complete' + } as any) + + expect(getUiState().sessionStats).toMatchObject({ inputTokens: 33189, outputTokens: 622 }) + }) + + it('folds the usage rider a resume replays out of the session log', () => { + const onEvent = createGatewayEventHandler(buildCtx()) + + onEvent({ + payload: { session_turns: 1, usage: { calls: 1, input: 1000, output: 120, total: 1120 } }, + type: 'session.stats' + } as any) + + expect(getUiState().sessionStats).toEqual({ costUsd: 0, inputTokens: 1000, outputTokens: 120, turns: 1 }) + }) + it('leaves sessionStats untouched when the payload carries neither field', () => { const onEvent = createGatewayEventHandler(buildCtx()) const before = getUiState().sessionStats diff --git a/src/harness/client.ts b/src/harness/client.ts index 1fc8e15..78bad23 100644 --- a/src/harness/client.ts +++ b/src/harness/client.ts @@ -455,6 +455,10 @@ export class HarnessGatewayClient extends GatewayClient { * `assistant/message` carries the `usage` the adapter reported — so a * resumed session has no reason to restart its counters at zero while * `turns:` keeps counting; that mismatch was visible on the stats line. + * + * Only this session's own steps: a subagent's live in the child's log, and + * this session never held those events. Live delegations still fold in + * through `onChildSessionEvent`. */ private replayUsage(events: readonly SessionEvent[]): Usage { const totals: Usage = { calls: 0, input: 0, output: 0, total: 0 } @@ -753,6 +757,14 @@ export class HarnessGatewayClient extends GatewayClient { if (usage) { child.inputTokens += usage.inputTokens + (usage.cacheReadTokens ?? 0) + (usage.cacheWriteTokens ?? 0) child.outputTokens += usage.outputTokens + // …and onto the session odometer, because a delegated step is spend + // on this session's behalf: it is the same bill, and the stats line + // and /status have always claimed to count subagents. Only the + // child's own row counted it before, so a run that fanned out + // reported a fraction of what it actually used. The two sides of + // the bus are disjoint (bindAgent routes by session identity), so + // nothing is counted twice. + foldUsage(this.usageTotals, usage) } // One event per step rather than one per delta: a child's reasoning