Skip to content

fix(claude-code): count cached and subagent tokens in usage - #5227

Open
gwarmstrong wants to merge 8 commits into
NVIDIA-NeMo:mainfrom
gwarmstrong:georgea/claude-code-usage-cache
Open

gwarmstrong wants to merge 8 commits into
NVIDIA-NeMo:mainfrom
gwarmstrong:georgea/claude-code-usage-cache

Conversation

@gwarmstrong

@gwarmstrong gwarmstrong commented Oct 28, 2026 •

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes the token usage that claude_code_agent reports. #5158 removed a real double count (stream-json repeats a message's usage on every content block) and that stays. But its totals were still wrong, because it read usage from places that hold only part of it:

  • Cached input was dropped. It took input from the result event's usage.input_tokens. In Anthropic's usage that is only the uncached part; cache writes and reads are in cache_creation_input_tokens and cache_read_input_tokens. A five-call DeepSeek session reported 341 input tokens for the 99,541 the server processed. cached_tokens was always 0.
  • Subagent calls were dropped. The result event's usage covers the main conversation only. The subagents' calls are in its modelUsage. In a session with two subagents, usage held 40,016 of 101,969 input tokens.
  • A run cut off at the timeout reported 0 output. Such a run has no result event, and the per-message usage in stream-json is the one each message started with. vLLM sends output only at the end of a message, so every capped run reported no output.

#5158's tests used hand-written events with no cache fields and with output in the per-message usage, so they built in the same assumptions.

The changes:

  • A run that ends on its own takes usage from the result event's modelUsage, summed over models. Input is inputTokens plus the cache writes and reads, and cached_tokens is the cache reads. It falls back to usage when modelUsage is absent.
  • A run cut off at the timeout counts each call once, by message id, from Claude Code's session transcripts and the stream:
    • A call whose transcript shows it finished counts its final usage, output included.
    • A call cut off in flight counts the input it was sent with (from its transcript entry, else from the stream's message_start usage) and no output, since the server reports output only at the end. A call cut off before Claude Code recorded it, seconds after it started, is in neither and is not counted.
    • Context compaction is not counted: Claude Code leaves it out of its transcripts.
    • When no transcript could be read, the stream's per-message input is used as before.
    • Transcripts are read once per run and shared with the observation extraction. The host path reads the run's CLAUDE_CONFIG_DIR before removing it; the sandbox path extracts the transcript archive the runner already collects.
  • Usage records are validated. A record must report input or output tokens, and every reported count must be a non-negative integer. An invalid modelUsage falls back to usage.
  • Unknown is kept apart from zero. The response's cached_tokens is the cache reads, or None when no cache count was reported; reasoning_tokens is None, since Anthropic reports none.
  • Gym's Anthropic converter keeps the cache split in both directions, so claude_code_agent behind a Gym model server reports it too. Responses' cached_tokens becomes cache_read_input_tokens, with input_tokens excluding it. Anthropic usage converted to Responses counts cache reads and writes in input_tokens, with the reads as cached_tokens. Cache writes have no Responses field, so on the way back they are part of input_tokens. Any caller that read anthropic_to_responses's input_tokens as the uncached part now gets the total.
  • README: what the usage counts.

One gap that this PR does not correct, documented in the README: the count is Claude Code's own, and against vLLM Claude Code counts a call's prompt twice when that call's final usage reports 0 uncached input tokens.

  • Why: vLLM sends the whole prompt as message_start's input_tokens and the split only in the final message_delta. Claude Code keeps the start value when the final one is 0, then adds the cache fields on top.
  • How often: in the 2026-10-17 DeepSeek V4.1 Flash run that was 321 of 43,473 calls, about 0.7% of input.
  • Where it happens: never with the Kimi or GLM servers we checked, and not with Anthropic's API, which sends the split at message_start.

Follow-up to #5158.

Validation

  • At d8aee9c (main merged 2026-11-04): claude_code_agent: 222 passed. Anthropic converter tests: 56 passed (7 new, 1 updated); the other tests that use the converter: 1,242 passed. pre-commit clean.
  • Ground truth: a recording proxy between Claude Code 2.1.207 and a self-hosted DeepSeek V4.1 Flash on vLLM logged, for every call, the server's message_start usage, each message_delta usage, and whether message_stop arrived. "Server" below is the sum of the calls that finished (input = the final split's uncached + cache writes + cache reads).
  • New tests, built from those sessions: seven of them plus one GLM 5.3 session whose server sends no cache fields. Each fixture keeps the event structure, ids and usage, with all text replaced by placeholders. Every expected value is computed in the test from the proxy records, not from the parser: finished calls as the server sent them, plus the message_start input of each unfinished call that the stream or a transcript names.
    • The table's "Main" column is main's parser on the same sessions. Every case in the validation and converter tests below, and the killed-subagent session, also gives a different result at 5f52948, this PR's previous head.
    • Further tests cover the transcript readers, including unreadable transcript files, and check that each run reads its transcripts once.
    • Validation tests: empty, cache-only, boolean, fractional, negative and string counts in modelUsage fall back to usage; a usage without cache fields reports cached_tokens None.
    • Converter tests: both directions, usage without cache counts, round trips that keep the total input, and a served Anthropic usage (200 uncached + 800 cache reads, 50 output) that Claude Code receives through the converter and Gym parses back to (1,000 / 800 / 50).
    • Two existing tests changed only to expect the new cached_tokens key, and one converter egress test now expects input 13 (10 uncached + 3 cache reads) instead of 10.
Session Calls Server (input / cached / output) This PR Main (input / output)
ends on its own, 18 calls 18 403,757 / 338,560 / 5,336 427,565 / 338,560 / 5,336 (one 0-input call counted twice by Claude Code) 24,877 / 5,336
ends on its own 5 99,541 / 99,072 / 666 same 341 / 666
three parallel tool calls 2 39,461 / 39,040 / 265 same 37 / 265
two subagents 7 101,969 / 100,864 / 1,013 same 80 / 550
killed during a model call 4, 3 finished 61,599 / 40,320 / 1,902 same (the call in flight is in neither the stream nor a transcript) 61,599 / 0
killed during a tool call 1 19,584 / 19,456 / 133 39,168 / 19,456 / 133 (the 0-input case) 19,584 / 0
killed while subagents ran 5, 4 finished, plus a subagent call in flight 56,317 / 33,536 / 637, plus 12,518 input in flight 68,835 / 33,536 / 637 43,818 / 0
GLM 5.3, no cache fields 12 1,724,053 / 0 / 227,817 (per-message input from the stream; no proxy) same 1,724,053 / 227,817

Rollouts

Live check at 2a528b4:

  • Build: this change merged onto the Terminal-Bench-Science Claude Code profile (feat(terminal-bench-science): run three more harnesses in the task sandbox #5459 at 84ac5a7).
  • Setup: Claude Code ran in the task sandbox on OpenSandbox, with the recording proxy as its Anthropic endpoint and DeepSeek V4.1 Flash behind it.
  • Runs: two capped at 600 s with the official instructions, and two that ended on their own because the task's instruction was replaced by a three-step one.
  • Matching: each rollout's Claude Code session id selects its proxy calls.
Run Calls Server Gym Difference
eeg-erp-recovery, ends on its own 4 79,542 / 60,032 / 420 79,542 / 60,032 / 420 none
mri-harmonization, ends on its own 4 79,201 / 59,776 / 302 79,201 / 59,776 / 302 none
eeg-erp-recovery, capped 41, 40 finished 3,403,953 / 2,863,616 / 75,650 3,403,953 / 2,863,616 / 75,650 none; the call in flight (151,544 input), cut off 4 s after it started, is not counted
mri-harmonization, capped 34, 33 finished 1,803,626 / 1,535,488 / 81,837 1,858,932 / 1,437,568 / 74,631 the three cases below

In the capped mri-harmonization run, Gym equals the server's count:

  • plus the call in flight, cut off 8 s after it started (56,606 input, no output);
  • plus 98,048 for the 0-input case;
  • minus a context compaction (99,348 / 97,920 / 7,206), which is not in the transcripts.

A call counts in flight once Claude Code has recorded it. The capped eeg-erp-recovery run's last call was cut off before that, so neither the transcripts nor the stream name it.

Before this change, all four would have reported cached_tokens 0, and the capped runs 0 output.

Compatibility and benchmark impact

Reported usage for Claude Code runs changes:

Rewards and grading are unchanged. No configuration changes.

A run that ends on its own now takes its usage from the result event's
modelUsage, which covers every call the CLI made (subagents included),
and counts the prompt cache's reads and writes as input, with the reads
as cached_tokens. The result's usage, which NVIDIA-NeMo#5158 used, covers the main
conversation only and splits input the Anthropic way.

A run cut off at the timeout has no result event. It now counts the
calls that finished, from Claude Code's session transcripts (host config
dir or the sandbox's transcript archive), so their output is counted. A
call cut off in flight is not counted.

Tests use eight real Claude Code 2.1.207 sessions with model text
removed, checked against what the server sent for each call.

Signed-off-by: gwarmstrong <gwarmstrong@users.noreply.github.com>
Signed-off-by: gwarmstrong <gwarmstrong@users.noreply.github.com>
@gwarmstrong
gwarmstrong requested a review from lafrancef October 28, 2026 21:28
@copy-pr-bot

copy-pr-bot Bot commented Oct 28, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@gwarmstrong

Copy link
Copy Markdown
Contributor Author

/ok to test 4e49a9b

@svcnvidia-nemo-ci svcnvidia-nemo-ci left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SHIP WITH CARE

One RISK: unreadable host transcripts suppress the stream-usage fallback and report zero tokens.

Reviewed all five changed files and directly affected parser, response, host-cleanup, and sandbox-artifact paths. No other qualifying findings. Static inspection only; tests and rollouts were not executed.

Review generated by codex. This was a full review. Request a full review with /review reset=true.

usage.update(transcript_usage(path.read_text(errors="replace").splitlines()))
except OSError:
LOG.warning("Could not read Claude Code transcript %s for token usage", path, exc_info=True)
return usage if paths else None

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

RISK — Preserve stream fallback when no transcript file could be read

Failure: If a matching transcript exists but read_text() raises OSError, this reader logs the failure and returns {} because paths is nonempty. The total-selection branch treats that mapping as authoritative. Consequently, a local no-result run with one streamed message carrying input_tokens=40 reports (input, cached, output) = (0, 0, 0) rather than the available fallback (40, 0, 0). This is introduced by the new transcript path; the archive reader already returns None on read failure.

Correction: Could transcript_usage_from_dir track successful reads instead of discovered paths? Replace the loop and return after usage initialization with:

readable = False
for path in paths:
    try:
        text = path.read_text(errors="replace")
        usage.update(transcript_usage(text.splitlines()))
        readable = True
    except OSError:
        LOG.warning("Could not read Claude Code transcript %s for token usage", path, exc_info=True)
return usage if readable else None

Preservation: Unavailable transcript evidence must permit stream fallback, while readable evidence remains authoritative—even if it contains no finished calls. Keep the caller's finished is not None check: changing it to truthiness would incorrectly count an in-flight message when a readable transcript returns {}. Also retain successfully read main/subagent counts when another file is unreadable.

Proposed acceptance checks (not executed):

  • Create projects/session.jsonl, make its read_text() raise OSError, and parse a no-result stream containing assistant message m1 with input_tokens=40; expect (40, 0, 0).
  • Make that transcript readable with only m1, stop_reason=null; expect (0, 0, 0).
  • Supply one readable finished message with uncached input 10, cache writes 2, cache reads 3, output 7, plus one unreadable file; expect (15, 3, 7).
  • Retain the killed-session assertions: finished calls match the proxy-derived totals, killed-sub excludes its in-flight call, and archive/directory counts agree.

Suggested code is source-reasoned and requires author verification.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 5f52948. transcript_usage_from_dir now returns None unless at least one transcript was read, so a session whose transcripts can't be read falls back to the stream. A transcript that was read stays authoritative even when it holds no finished call, and the caller keeps its is not None check.

The acceptance checks are built from the captured sessions rather than hand-written numbers:

  • An unreadable projects/session.jsonl (a directory, so read_text() raises IsADirectoryError) gives None, and the killed-mid session then counts what it counts with no transcripts at all, which is non-zero.
  • killed-sub's transcripts, keeping only the entries with no stop reason, count (0, 0, 0).
  • Each killed session's transcripts plus one unreadable file still match the proxy totals.

The first two fail at 4e49a9b. The existing killed-session and archive tests are unchanged and pass.

Signed-off-by: gwarmstrong <gwarmstrong@users.noreply.github.com>
@gwarmstrong

Copy link
Copy Markdown
Contributor Author

/ok to test 5f52948

@svcnvidia-nemo-ci svcnvidia-nemo-ci added area:agent Agent harnesses and Responses API agent behavior bug Something isn't working needs-review PR is ready for code review and waiting on a reviewer labels Oct 28, 2026
parse_stream_json,
artifacts.events,
fold_thinking=not self._token_id_capture_enabled(),
transcript=(lambda: transcript_usage_from_tar(transcripts)) if transcripts else None,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we collect transcript usage during the existing observation pass, or share the parsed events from observability.py:_read_events and the extracted directory in sandbox.py:parse_artifact_observations? On the fallback path, this reopens the archive and parses the JSONL before the observation code extracts and parses it again; the host path similarly rereads projects/**/*.jsonl. Sharing that work would also avoid injecting a transcript-I/O lambda into parse_stream_json.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 0d2ff4a. The transcripts are now read once per run and shared:

  • observability.read_transcripts(config_dir) parses every record of projects/**/*.jsonl and keeps the read gaps and a count of the files it could open.
  • Token usage reads all the records; observation extraction reads the ones with a sessionId. Their usage is matched by message id, which does not need the session id.
  • parse_stream_json takes the records as data (transcript_records, None when no transcript could be read), and the transcript-I/O lambda is gone.
  • On the host path, _run_claude_code reads them once the CLI has stopped. The same object then goes to the observation collector, which now receives the transcripts instead of the directory.
  • On the sandbox path, read_artifact_transcripts extracts the archive once. Its result feeds both parse_stream_json and parse_artifact_observations.

Tests check that the archive is read once per sandbox activation, and that the host path reads before it removes the config dir.

total_output = unidentified_output + sum(counts[1] for counts in message_usage.values())
metadata: dict = {"input_tokens": total_input, "output_tokens": total_output}
finished = transcript() if transcript is not None else None
calls = finished.values() if finished is not None else [unidentified, *message_usage.values()]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we preserve the known input usage for in-flight calls when transcripts are available? transcript_usage drops messages with stop_reason is None, and this branch then replaces all stream usage with the finished-call totals. The same interrupted call therefore contributes input tokens when its transcript is unavailable, but contributes nothing when it is readable (and a readable transcript with no finished calls yields zero). Please reconcile by message ID, preferring final transcript usage while retaining known input for unfinished calls. Alternatively, if the intended contract is finished calls only, apply it consistently and emit ObservationGap(code="partial_model_usage_unavailable") on incomplete/timed-out runs.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 0d2ff4a, reconciling by message id. Without a result event, each call is counted once:

  • a call whose transcript shows it finished counts its final usage;
  • a call cut off in flight counts the input it was sent with (from its transcript entry, else from the stream's message_start usage) and no output.

A readable transcript with no finished call now counts the input of the calls it and the stream name, instead of zero.

The expected values come from the recording proxy, not the parser: finished calls as the server sent them, plus the message_start input of each unfinished call that the stream or a transcript names. In the captured sessions, this changes killed-sub. Its subagent was cut off during a 12,518-token call that its transcript names but the stream does not, so its total goes from 56,317 to 68,835 input. killed-mid's in-flight call appears in neither, so no parser can see it. Its count is unchanged, and the test documents why.

return None
totals: UsageCounts = (0, 0, 0)
for usage in model_usage.values():
if not isinstance(usage, dict):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please require at least one token key (inputTokens or outputTokens) in each model usage record before treating it as valid. For example, modelUsage={"model": {}} currently produces (0, 0, 0), which is truthy and suppresses valid fallback usage. At the call site, use an explicit if counts is None: fallback. In both usage helpers, please validate reported token values with type(v) is int and v >= 0 rather than int(v or 0): coercion accepts booleans, truncates floats, accepts negative counts, and can raise on malformed strings. Missing optional fields can still be handled separately.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 0d2ff4a. Both usage helpers now go through one validator:

  • a record must report inputTokens/outputTokens (or input_tokens/output_tokens);
  • every count it reports must satisfy type(v) is int and v >= 0, otherwise the record is invalid;
  • an absent or null count is treated as unreported.

For modelUsage, one invalid model record makes the whole value invalid. The call site falls back with an explicit if counts is None:.

New tests: {"model": {}}, a cache-only record, True, 1.5, -1, "3", and a mix of a good and a bad model all fall back to usage, where 5f52948 produced (0, 0, 0), (1, 0, 1), (1, 0, 1), (-1, 0, 1) and (3, 0, 1). Invalid usage records count nothing, and null optional counts are unreported.

)

input_tokens = run_metadata.get("input_tokens", 0)
cached_tokens = run_metadata.get("cached_tokens", 0)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we preserve the distinction between unknown usage and a reported zero? Please propagate cached_tokens=None when no usage/cache fields were reported, rather than defaulting to zero here and in the parser. Also set reasoning_tokens=None below: Anthropic does not report a separate reasoning-token count. This matches the convention already used by AnthropicConverter._usage_to_responses_usage.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 0d2ff4a. cached_tokens is now None unless a usage record reported a cache count. The parser carries it as an optional count, and _response_from_output no longer defaults it to 0. reasoning_tokens is None, as in AnthropicConverter._usage_to_responses_usage. NeMoGymResponseUsage.sum_from_list and the observability and ATIF consumers already treat both as optional.

A reported zero stays zero. For example, Claude Code's own result and transcripts report cache counts of 0 against the GLM endpoint, which sends none, and they stay 0. Tests cover a result with no cache fields (None), explicit zeros (0), and the response's reasoning_tokens.

usage=NeMoGymResponseUsage(
input_tokens=input_tokens,
input_tokens_details=NeMoGymResponseInputTokensDetails(cached_tokens=0),
input_tokens_details=NeMoGymResponseInputTokensDetails(cached_tokens=cached_tokens),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If this is expected to report cache usage when claude_code_agent uses a Gym model_server, could we also update nemo_gym/anthropic_converter.py? AnthropicConverter.responses_to_anthropic_response currently emits only input/output totals, so the cache breakdown is lost before Claude Code receives it. The reverse _usage_to_responses_usage already maps cache reads into cached_tokens, but does not add cache reads/writes into total input. Please handle those fields consistently in both directions, preserving total input without double-counting cache tokens, and cover the Gym model-server path.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 0d2ff4a, in both directions:

  • responses_to_anthropic_response now reports Responses' cached_tokens as cache_read_input_tokens, with input_tokens excluding them. The synthesized SSE carries the split in message_start, as Anthropic streams it, and message_delta still sends only the output.
  • _usage_to_responses_usage adds cache reads and writes into input_tokens (so total_tokens includes them), with the reads as cached_tokens.
  • Responses has no count of cache writes, so on the way back they are part of input_tokens. The total is preserved and nothing is counted twice.

New converter tests cover both directions, an Anthropic usage without cache counts, and round trips that keep the input. One existing egress test expected the old behavior (input 10 that left out its 3 cache reads) and now expects 13.

Model-server path test: a served Anthropic usage of 200 uncached plus 800 cache reads and 50 output goes through anthropic_to_responses, responses_to_anthropic_response and the SSE. Claude Code keeps message_start's usage and updates it with message_delta's. Parsing that usage, as the result event or as a finished transcript entry, gives (1000, 800, 50). At 5f52948 the cache split was lost before Claude Code received it.

…counts

Read Claude Code's transcripts once per run and share the records between token usage and observation extraction, instead of a transcript reader injected into parse_stream_json. A run cut off at the timeout now counts each call once by message id: final usage for a call that finished, and the input of a call cut off in flight. Usage records must report input or output tokens as non-negative integers. cached_tokens and reasoning_tokens stay None when unreported. The Anthropic converter carries cache reads both ways: into Responses input and cached_tokens, and back out as cache_read_input_tokens.

Signed-off-by: gwarmstrong <gwarmstrong@users.noreply.github.com>
@gwarmstrong

Copy link
Copy Markdown
Contributor Author

/ok to test 0d2ff4a

@gwarmstrong
gwarmstrong requested a review from lafrancef October 29, 2026 21:38
@gwarmstrong

Copy link
Copy Markdown
Contributor Author

/ok to test 0d2ff4a

…-usage-cache

Signed-off-by: gwarmstrong <gwarmstrong@users.noreply.github.com>

# Conflicts:
#	responses_api_agents/claude_code_agent/parsing.py
@gwarmstrong

Copy link
Copy Markdown
Contributor Author

/ok to test 2a528b4

Signed-off-by: gwarmstrong <gwarmstrong@users.noreply.github.com>
@gwarmstrong

Copy link
Copy Markdown
Contributor Author

/ok to test 134645b

…-usage-cache

Signed-off-by: gwarmstrong <gwarmstrong@users.noreply.github.com>

# Conflicts:
#	responses_api_agents/claude_code_agent/app.py
#	responses_api_agents/claude_code_agent/sandbox.py
#	responses_api_agents/claude_code_agent/tests/test_app.py
#	tests/unit_tests/models/test_anthropic_converter.py
@gwarmstrong

Copy link
Copy Markdown
Contributor Author

/ok to test 9012f48

…-usage-cache

Signed-off-by: gwarmstrong <gwarmstrong@users.noreply.github.com>

# Conflicts:
#	responses_api_agents/claude_code_agent/README.md
#	responses_api_agents/claude_code_agent/app.py
@gwarmstrong

Copy link
Copy Markdown
Contributor Author

/ok to test d8aee9c

This branch was previously deployed

1 inactive deployment
public — d8aee9cf Deployed Nov 4, 2026 by copy-pr-bot[bot] via release / finalize / notify #10001
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:agent Agent harnesses and Responses API agent behavior bug Something isn't working needs-review PR is ready for code review and waiting on a reviewer sla:review-overdue Review response is over the one-business-day SLA

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants