Skip to content

fix(xai): include reasoning tokens in completion usage - #3493

Merged
Lancetnik merged 6 commits into
mainfrom
fix/xai-reasoning-completion-usage
Oct 26, 2026
Merged

Lancetnik merged 6 commits into
mainfrom
fix/xai-reasoning-completion-usage

Conversation

@vvlrff

@vvlrff vvlrff commented Oct 24, 2026

Copy link
Copy Markdown
Collaborator

Why are these changes needed?

xAI reports reasoning tokens beside the completion, not inside it, so Usage.completion_tokens for Grok reasoning models leaves them out and prompt + completion falls short of total. Fold reasoning into completion_tokens and keep thinking_tokens as its breakdown, as the OpenAI, Anthropic and Gemini mappers already do. Follow-up to #3446, which fixes the same gap for Gemini.

 normalize_usage(usage)               # ag2/config/xai/mappers.py
-  completion = usage.completion_tokens
+  completion = usage.completion_tokens
+  if usage.reasoning_tokens:
+    completion += usage.reasoning_tokens
   thinking_tokens = usage.reasoning_tokens   # unchanged, the breakdown

Evidence, from the live xAI API (grok-4-fast-reasoning):

  • Before: raw SamplingUsage is prompt=204, completion=32, reasoning=300, total=536. The old mapper passed completion=32 through, so prompt + completion = 236 != 536. prompt + completion + reasoning == total does hold.
  • After: Agent.ask() on the patched mapper gives p+c == total in both modes: plain 201 + 283 = 484, streaming 201 + 313 = 514, each with thinking_tokens kept as the breakdown (282 / 312).
  • New tests: test_completion_includes_reasoning_tokens (the numbers above) and test_reasoning_without_completion_tokens. The existing reasoning test's total_tokens is corrected to a consistent value.

Merge danger:

  • Door: two-way. It is a one-mapper change that reverts cleanly.
  • Blast radius: xAI consumers. completion_tokens (and AG-UI output_tokens) now grow by the reasoning count. Pricing and cross-provider comparisons become correct, but anything that budgets on completion_tokens will reach its limit sooner for reasoning models.
  • Other mappers were audited: OpenAI, Z.ai and Anthropic already include reasoning in completion, and Gemini Live's response_token_count already includes thoughts. ACP's thought_tokens was not verified against a live agent and is left unchanged.

Related issue number

Follows up #3446 (fixes #3443 for Gemini).

Checks

AI assistance

  • I understand the changes in this PR and can explain them in my own words.
  • I have verified that the PR description accurately reflects the actual diff.
  • If AI assistance was used, I reviewed, tested, and validated the generated code/text before submitting.

original4422 and others added 3 commits October 18, 2026 01:43
xAI reports reasoning_tokens beside completion_tokens (total = prompt +
completion + reasoning). Fold them into Usage.completion_tokens and keep
thinking_tokens as the breakdown, matching the OpenAI, Anthropic and
Gemini mappers.
@github-actions github-actions Bot added area:ag-ui AG-UI integration (ag2/ag_ui) area:providers LLM provider clients and mappers (ag2/config) labels Oct 24, 2026
@Lancetnik
Lancetnik merged commit ce17bf2 into main Oct 26, 2026
15 of 22 checks passed
@Lancetnik
Lancetnik deleted the fix/xai-reasoning-completion-usage branch October 26, 2026 04:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:ag-ui AG-UI integration (ag2/ag_ui) area:providers LLM provider clients and mappers (ag2/config)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: Gemini usage leaves thinking tokens out of completion_tokens

3 participants