fix(xai): include reasoning tokens in completion usage

ag2/config/xai/mappers.py

@@ -313,9 +313,15 @@ def normalize_usage(usage: usage_pb2.SamplingUsage | None) -> Usage:
     if usage is None:
         return Usage()
 
+    # xAI reports reasoning beside the completion (total = prompt + completion +
+    # reasoning); completion includes it, with thinking_tokens as the breakdown.
+    completion = float(usage.completion_tokens) if usage.completion_tokens else None
+    if usage.reasoning_tokens:
+        completion = (completion or 0.0) + usage.reasoning_tokens
+
     return Usage(
         prompt_tokens=float(usage.prompt_tokens) if usage.prompt_tokens else None,
-        completion_tokens=float(usage.completion_tokens) if usage.completion_tokens else None,
+        completion_tokens=completion,
         total_tokens=float(usage.total_tokens) if usage.total_tokens else None,
         cache_read_input_tokens=float(usage.cached_prompt_text_tokens) if usage.cached_prompt_text_tokens else None,
         thinking_tokens=float(usage.reasoning_tokens) if usage.reasoning_tokens else None,

test/ag_ui/test_usage.py

@@ -629,8 +629,8 @@ async def test_bedrock_cache_is_added_to_its_input(self) -> None:
 
     @pytest.mark.parametrize("provider", ["openai", "xai"])
     async def test_other_providers_are_reported_as_they_count(self, provider: str) -> None:
-        """OpenAI counts cache and reasoning inside its totals already. xAI's reasoning
-        is not known to sit outside them, so it is not added in.
+        """OpenAI and xAI count reasoning inside the normalized completion already,
+        so it is not added in again.
         """
         usage = Usage(
             prompt_tokens=100,

test/config/xai/test_xai_usage.py

@@ -32,13 +32,29 @@ def test_cached_prompt_tokens_map_to_cache_read() -> None:
 
 
 def test_reasoning_tokens_map_to_thinking_tokens() -> None:
-    raw = usage_pb2.SamplingUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15, reasoning_tokens=3)
+    raw = usage_pb2.SamplingUsage(prompt_tokens=10, completion_tokens=5, total_tokens=18, reasoning_tokens=3)
 
     result = normalize_usage(raw)
 
     assert result.thinking_tokens == 3.0
 
 
+def test_completion_includes_reasoning_tokens() -> None:
+    # xAI reports reasoning beside the completion: total = prompt + completion + reasoning.
+    raw = usage_pb2.SamplingUsage(prompt_tokens=204, completion_tokens=32, total_tokens=536, reasoning_tokens=300)
+
+    result = normalize_usage(raw)
+
+    assert result == Usage(prompt_tokens=204.0, completion_tokens=332.0, total_tokens=536.0, thinking_tokens=300.0)
+    assert result.prompt_tokens + result.completion_tokens == result.total_tokens
+
+
+def test_reasoning_without_completion_tokens() -> None:
+    raw = usage_pb2.SamplingUsage(prompt_tokens=10, total_tokens=50, reasoning_tokens=40)
+
+    assert normalize_usage(raw).completion_tokens == 40.0
+
+
 def test_zero_or_missing_fields_become_none() -> None:
     raw = usage_pb2.SamplingUsage()