Skip to content

fix: stamp provider usage at the adapter boundary and close the compact accounting gap - #2138

Open
Q1hangL wants to merge 1 commit into
xorbitsai:mainfrom
Q1hangL:fix/520-compact-usage-stamps
Open

fix: stamp provider usage at the adapter boundary and close the compact accounting gap#2138
Q1hangL wants to merge 1 commit into
xorbitsai:mainfrom
Q1hangL:fix/520-compact-usage-stamps

Conversation

@Q1hangL

@Q1hangL Q1hangL commented Sep 5, 2026

Copy link
Copy Markdown

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)

  • Provider usage stamping at the adapter boundary: OpenAICompatibleLLM snapshots resp.usage once and stamps a top-level usage key 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's cache_write_input_tokens) ride the stamp when non-zero.
  • PatternRuntime usage resolution: _resolve_usage_payload shared by _extract_token_usage and _extract_cached_tokens — top-level first, then one level under raw, 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.
  • Compact-path freshness handling: compaction usage records carry synthetic_purpose and 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.
  • Direct consumer fixes required by the envelope change: shared unwrap_chat_text / classify_chat_response for the three hand-written unwrapping call sites (ContextBuilder._compact_*, optimize_instructions) so tool_call envelopes fail explicitly instead of being repr()ed into compacted context, and the ReAct finalize fallback surfaces the envelope's text rather than the envelope itself.
  • Docs: BaseLLM.chat()/vision_chat() docstrings describe the envelope contract, with raw marked provider-dependent/optional.

Moved to follow-ups (preserved in #1787)

Test plan

  • tests/core/agent/test_compact_llm_usage_contract.py (41 tests): per adapter family, responses built with real SDK types (e.g. ChatCompletion with populated CompletionUsage), only the SDK client patched, real adapter driven through the real compact path — asserting action_end_llm (purpose=context_compaction) token fields, get_total_token_usage(), and end-to-end cached_input_tokens; plus extractor edge shapes (string raw, top-level usage_metadata, all-zero usage), checkpoint round-trip, the truncation-suppression scenario, and a no-double-counting guard between the contextvar ledger and llm_calls.
  • Local: 1707 passed, 0 failed across agent / model.chat / web suites (2 pre-existing environment failures in test_executor.py / test_browser_tools.py, red on clean main too). mypy clean on touched files; ruff check / ruff format clean.
  • Key fixes mutation-verified on this branch: deleting the raw fallback turns the extractor parametrization red; dropping the synthetic-record skip turns the truncation-suppression test red.

Known boundaries

cc @rogercloud

…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).

@gemini-code-assist gemini-code-assist Bot 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.

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.

Comment on lines 677 to 679
add_token_usage(
input_tokens=input_tokens,
output_tokens=output_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.

medium

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.

Suggested change
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
  1. 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.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

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.

Comment on lines 597 to 599
add_token_usage(
input_tokens=input_tokens,
output_tokens=output_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.

medium

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.

Suggested change
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
  1. 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.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix: compact LLM token usage for real provider responses

2 participants