Repository navigation
fix(claude-code): count cached and subagent tokens in usage - #5227
gwarmstrong wants to merge 8 commits into
Conversation
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>
|
/ok to test 4e49a9b |
svcnvidia-nemo-ci
left a comment
There was a problem hiding this comment.
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 |
There was a problem hiding this comment.
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 NonePreservation: 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 itsread_text()raiseOSError, and parse a no-result stream containing assistant messagem1withinput_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-subexcludes its in-flight call, and archive/directory counts agree.
Suggested code is source-reasoned and requires author verification.
There was a problem hiding this comment.
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, soread_text()raisesIsADirectoryError) givesNone, 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>
|
/ok to test 5f52948 |
| parse_stream_json, | ||
| artifacts.events, | ||
| fold_thinking=not self._token_id_capture_enabled(), | ||
| transcript=(lambda: transcript_usage_from_tar(transcripts)) if transcripts else None, |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Done in 0d2ff4a. The transcripts are now read once per run and shared:
observability.read_transcripts(config_dir)parses every record ofprojects/**/*.jsonland 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_jsontakes 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_codereads 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_transcriptsextracts the archive once. Its result feeds bothparse_stream_jsonandparse_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()] |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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): |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Done in 0d2ff4a. Both usage helpers now go through one validator:
- a record must report
inputTokens/outputTokens(orinput_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) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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), |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Done in 0d2ff4a, in both directions:
responses_to_anthropic_responsenow reports Responses'cached_tokensascache_read_input_tokens, withinput_tokensexcluding 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_usageadds cache reads and writes intoinput_tokens(sototal_tokensincludes them), with the reads ascached_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>
|
/ok to test 0d2ff4a |
|
/ok to test 0d2ff4a |
…-usage-cache Signed-off-by: gwarmstrong <gwarmstrong@users.noreply.github.com> # Conflicts: # responses_api_agents/claude_code_agent/parsing.py
|
/ok to test 2a528b4 |
Signed-off-by: gwarmstrong <gwarmstrong@users.noreply.github.com>
|
/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
|
/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
|
/ok to test d8aee9c |
What does this PR do?
Fixes the token usage that
claude_code_agentreports. #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:usage.input_tokens. In Anthropic's usage that is only the uncached part; cache writes and reads are incache_creation_input_tokensandcache_read_input_tokens. A five-call DeepSeek session reported 341 input tokens for the 99,541 the server processed.cached_tokenswas always 0.usagecovers the main conversation only. The subagents' calls are in itsmodelUsage. In a session with two subagents,usageheld 40,016 of 101,969 input tokens.#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:
modelUsage, summed over models. Input isinputTokensplus the cache writes and reads, andcached_tokensis the cache reads. It falls back tousagewhenmodelUsageis absent.message_startusage) 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.CLAUDE_CONFIG_DIRbefore removing it; the sandbox path extracts the transcript archive the runner already collects.modelUsagefalls back tousage.cached_tokensis the cache reads, orNonewhen no cache count was reported;reasoning_tokensisNone, since Anthropic reports none.claude_code_agentbehind a Gym model server reports it too. Responses'cached_tokensbecomescache_read_input_tokens, withinput_tokensexcluding it. Anthropic usage converted to Responses counts cache reads and writes ininput_tokens, with the reads ascached_tokens. Cache writes have no Responses field, so on the way back they are part ofinput_tokens. Any caller that readanthropic_to_responses'sinput_tokensas the uncached part now gets the total.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.
message_start'sinput_tokensand the split only in the finalmessage_delta. Claude Code keeps the start value when the final one is 0, then adds the cache fields on top.message_start.Follow-up to #5158.
Validation
mainmerged 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-commitclean.message_startusage, eachmessage_deltausage, and whethermessage_stoparrived. "Server" below is the sum of the calls that finished (input = the final split's uncached + cache writes + cache reads).message_startinput of each unfinished call that the stream or a transcript names.modelUsagefall back tousage; a usage without cache fields reportscached_tokensNone.cached_tokenskey, and one converter egress test now expects input 13 (10 uncached + 3 cache reads) instead of 10.Rollouts
Live check at 2a528b4:
In the capped mri-harmonization run, Gym equals the server's count:
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_tokens0, 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.