fix: stamp provider usage at the adapter boundary and close the compact accounting gap - #2138
fix: stamp provider usage at the adapter boundary and close the compact accounting gap#2138Q1hangL wants to merge 1 commit into
Conversation
…ct accounting gap Fixes xorbitsai#520. Split out from xorbitsai#1787 as the focused, smallest coherent set: provider usage stamping on chat()/vision_chat() envelopes across the OpenAI family (inherited by DeepSeek/DashScope/OpenRouter/Azure), Zhipu/Gemini/Claude text paths upgraded to the {type, content, usage} envelope the streaming path already emits, Xinference stamped on every branch, plus cache metrics when non-zero. PatternRuntime resolves usage from top-level stamps first and one level under raw as a fallback, with strict counter coercion (no bools, non-finite, negative, or non-integral values shadowing later aliases). Compaction usage records are marked synthetic_purpose and excluded from the context-freshness baseline so a failed LLM compaction can no longer suppress the truncation fallback. Direct consumer fixes required by the envelope change come along: a shared unwrap_chat_text/classify_chat_response path for the three hand-written unwrapping call sites (the xorbitsai#1714 repr-leak class) and the ReAct finalize fallback. Regression coverage at the SDK transport boundary (real SDK types, only the client patched), including the truncation-suppression scenario and double-counting guards; key fixes mutation-verified on this branch (revert -> test goes red).
There was a problem hiding this comment.
Code Review
This pull request introduces structural classification of LLM chat response shapes to prevent raw response envelopes (such as tool calls) from being stringified and leaked into transcripts or compacted contexts. It adds a new classification helper, stamps top-level usage payloads onto result envelopes across various adapters, and introduces a synthetic purpose marker to prevent internal compaction calls from hijacking the context-size freshness baseline. The review feedback correctly identifies opportunities to record cached input tokens in the token ledger for both the Claude and Gemini adapters.
| add_token_usage( | ||
| input_tokens=input_tokens, | ||
| output_tokens=output_tokens, |
There was a problem hiding this comment.
The add_token_usage call currently does not record the extracted cached input tokens (cache_read) to the global/context token ledger. Passing cached_input_tokens=int(cache_read or 0) ensures that the context-local token ledger accurately tracks prompt-cache hits for Claude, while safely coercing missing or null values to 0 to prevent exceptions in token accounting.
| add_token_usage( | |
| input_tokens=input_tokens, | |
| output_tokens=output_tokens, | |
| add_token_usage( | |
| input_tokens=input_tokens, | |
| output_tokens=output_tokens, | |
| cached_input_tokens=int(cache_read or 0), |
References
- For token accounting helpers, prefer fail-open behavior (coercing missing or null values to 0) rather than raising exceptions, to ensure LLM calls do not break due to accounting issues.
There was a problem hiding this comment.
Verified against the current head — this is already handled. The add_token_usage call a few lines below (claude.py:677-685) passes both cached_input_tokens=cache_read and cache_write_input_tokens=cache_write; the two values come from _anthropic_input_usage's normalization of Anthropic's usage. So the ledger does record Claude's prompt-cache hits, and this suggestion would double-pass the same field. Dropping as verified-safe.
| add_token_usage( | ||
| input_tokens=input_tokens, | ||
| output_tokens=output_tokens, |
There was a problem hiding this comment.
The add_token_usage call currently does not record the extracted cached input tokens (cached_tokens) to the global/context token ledger. Passing cached_input_tokens=int(cached_tokens or 0) ensures that the context-local token ledger accurately tracks prompt-cache hits for Gemini, while preferring fail-open behavior by coercing missing or null values to 0.
| add_token_usage( | |
| input_tokens=input_tokens, | |
| output_tokens=output_tokens, | |
| add_token_usage( | |
| input_tokens=input_tokens, | |
| output_tokens=output_tokens, | |
| cached_input_tokens=int(cached_tokens or 0), |
References
- For token accounting helpers, prefer fail-open behavior (coercing missing or null values to 0) rather than raising exceptions, to ensure LLM calls do not break due to accounting issues.
There was a problem hiding this comment.
Verified against the current head — already handled. The add_token_usage call at gemini.py:597-602 passes cached_input_tokens=cached_tokens, where cached_tokens = getattr(usage_metadata, "cached_content_token_count", 0) or 0 (gemini.py:584-586) — i.e. it already fails open to 0 for missing/None values. The stamp above additionally guards the comparison by type so an unexpected non-numeric value can never raise out of chat(). Dropping as verified-safe.
Fixes #520. Split out from #1787 per the review discussion — this is the focused branch from current
main; #1787 stays as the reference implementation and becomes follow-ups (see below).Scope (the smallest coherent set for #520)
OpenAICompatibleLLMsnapshotsresp.usageonce and stamps a top-levelusagekey on all six chat/vision return sites (inherited by DeepSeek/DashScope/OpenRouter/Azure); Zhipu/Gemini/Claude text paths return the same{type, content, usage}envelope the streaming path already emits instead of bare strings; Xinference stamps on every envelope branch. Cache metrics (cached_input_tokens, Claude'scache_write_input_tokens) ride the stamp when non-zero.PatternRuntimeusage resolution:_resolve_usage_payloadshared by_extract_token_usageand_extract_cached_tokens— top-level first, then one level underraw, fail-open on unknown shapes; strict counter coercion (_coerce_usage_int) rejects bools, non-finite, negative, and non-integral values instead of truncating or crashing.synthetic_purposeand are excluded from the context-freshness baseline (_latest_freshness_baseline_call), so a stamped-but-failed LLM compaction can no longer suppress the truncation fallback; the field round-trips through checkpoints with a backward-compatible default.unwrap_chat_text/classify_chat_responsefor the three hand-written unwrapping call sites (ContextBuilder._compact_*,optimize_instructions) so tool_call envelopes fail explicitly instead of beingrepr()ed into compacted context, and the ReAct finalize fallback surfaces the envelope's text rather than the envelope itself.BaseLLM.chat()/vision_chat()docstrings describe the envelope contract, withrawmarked provider-dependent/optional.Moved to follow-ups (preserved in #1787)
usage_attempts/ billed-retry-attempt accounting (RetryWrapper merge, OpenRouter internal retries, exception carriers, multi-attempt cache scope)_normalize_vision_responsedelegation) and the defaultstream_chatboundary/usage-chunk workTest plan
tests/core/agent/test_compact_llm_usage_contract.py(41 tests): per adapter family, responses built with real SDK types (e.g.ChatCompletionwith populatedCompletionUsage), only the SDK client patched, real adapter driven through the real compact path — assertingaction_end_llm(purpose=context_compaction) token fields,get_total_token_usage(), and end-to-endcached_input_tokens; plus extractor edge shapes (stringraw, top-levelusage_metadata, all-zero usage), checkpoint round-trip, the truncation-suppression scenario, and a no-double-counting guard between the contextvar ledger andllm_calls.test_executor.py/test_browser_tools.py, red on clean main too). mypy clean on touched files;ruff check/ruff formatclean.rawfallback turns the extractor parametrization red; dropping the synthetic-record skip turns the truncation-suppression test red.Known boundaries
model_dump()branch, as the real zai-sdk response is pydantic.BaseLLM's defaultstream_chatstill treats any dict as a tool_call envelope — no current adapter triggers it (all overridestream_chat;PatternRuntimebypasses the default); the fix lives in the follow-up.cc @rogercloud