fix(xai): bill web_search from server_side_tool_usage_details

litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py

@@ -2,6 +2,7 @@
 Helper utilities for tracking the cost of built-in tools.
 """
 
+from collections.abc import Mapping
 from typing import Any, Final, Literal
 
 import litellm
@@ -23,6 +24,14 @@
 )
 
 
+def _usage_reports_server_side_web_search_calls(usage: Usage) -> bool:
+    details: Final = getattr(usage, "server_side_tool_usage_details", None)
+    if not isinstance(details, Mapping):
+        return False
+    calls: Final = details.get("web_search_calls")
+    return isinstance(calls, int) and calls > 0
+
+
 class StandardBuiltInToolCostTracking:
     """
     Helper class for tracking the cost of built-in tools
@@ -351,6 +360,10 @@ def response_object_includes_web_search_call(response_object: Any, usage: Usage
                 # and _handle_web_search_cost() is never called.
                 if hasattr(usage, "server_tool_use") and _get_web_search_requests(usage.server_tool_use) is not None:
                     return True
+                # xAI reports usage.server_side_tool_usage_details.web_search_calls; a searched
+                # answer with no url_citation annotations has no other chat-path signal
+                if _usage_reports_server_side_web_search_calls(usage):
+                    return True
             return False
         elif isinstance(response_object, ResponsesAPIResponse):
             # response api explicitly includes web_search_call in the output
@@ -370,6 +383,8 @@ def response_object_includes_web_search_call(response_object: Any, usage: Usage
                 )
             ):
                 return True
+            if _usage_reports_server_side_web_search_calls(usage):
+                return True
 
         return False
 
@@ -432,7 +447,9 @@ def response_includes_output_type(
         """
         output: Final = response_object.output
         for output_item in output:
-            _output_type: str | None = getattr(output_item, "type", None)
+            _output_type: str | None = (
+                output_item.get("type") if isinstance(output_item, dict) else getattr(output_item, "type", None)
+            )
             if _output_type == output_type:
                 return True
         return False

litellm/llms/xai/chat/transformation.py

@@ -1,4 +1,4 @@
-from collections.abc import AsyncIterator, Iterator
+from collections.abc import AsyncIterator, Iterator, Mapping
 from typing import Any, Final
 
 import httpx
@@ -12,13 +12,15 @@
     strip_name_from_messages,
 )
 from litellm.llms.xai.common_utils import XAIModelInfo
+from litellm.llms.xai.cost_calculator import (
+    apply_server_side_tool_usage_details_to_usage,
+)
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.utils import (
     Choices,
     ModelResponse,
     ModelResponseStream,
-    PromptTokensDetailsWrapper,
     Usage,
 )
 
@@ -248,7 +250,7 @@ def transform_response(
         XAI API returns empty string for finish_reason when using tools,
         so we need to fix this after the standard OpenAI transformation.
 
-        Also handles X.AI web search usage tracking by extracting num_sources_used.
+        Also handles X.AI web search usage tracking.
         """
 
         # First, let the parent class handle the standard transformation
@@ -351,25 +353,20 @@ def _fold_reasoning_tokens_into_completion(
 
     def _enhance_usage_with_xai_web_search_fields(self, model_response: ModelResponse, raw_response_json: dict) -> None:
         """
-        Extract num_sources_used from X.AI response and map it to web_search_requests.
+        Copy usage.server_side_tool_usage_details from the provider usage block
+        onto model_response.usage for tool cost calculation.
         """
         if not hasattr(model_response, "usage") or model_response.usage is None:
             return
 
         usage: Final[Usage] = model_response.usage
-        num_sources_used = None
-        response_usage: Final = raw_response_json.get("usage", {})
-        if isinstance(response_usage, dict) and "num_sources_used" in response_usage:
-            num_sources_used = response_usage.get("num_sources_used")
-
-        # Map num_sources_used to web_search_requests for cost detection
-        if num_sources_used is not None and num_sources_used > 0:
-            if usage.prompt_tokens_details is None:
-                usage.prompt_tokens_details = PromptTokensDetailsWrapper()
-
-            usage.prompt_tokens_details.web_search_requests = int(num_sources_used)
-            setattr(usage, "num_sources_used", int(num_sources_used))
-            verbose_logger.debug("X.AI web search sources used: %s", num_sources_used)
+        response_usage: Final = raw_response_json.get("usage")
+        if not isinstance(response_usage, dict):
+            return
+        details: Final = response_usage.get("server_side_tool_usage_details")
+        if isinstance(details, Mapping):
+            apply_server_side_tool_usage_details_to_usage(usage, details)
+            verbose_logger.debug("X.AI server_side_tool_usage_details: %s", details)
 
     @staticmethod
     def _normalize_openai_compatible_usage_totals(

litellm/llms/xai/cost_calculator.py

@@ -4,14 +4,37 @@
 - Handles XAI-specific reasoning token billing (billed as part of completion tokens)
 """
 
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Final
 
 from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
-from litellm.types.utils import Usage
+from litellm.types.utils import PromptTokensDetailsWrapper, Usage
 
 if TYPE_CHECKING:
     from litellm.types.utils import ModelInfo
 
+# https://docs.x.ai/developers/pricing#tools-pricing — default when unset in model map
+_DEFAULT_WEB_SEARCH_COST_PER_CALL: Final = 5.0 / 1000.0
+
+
+def apply_server_side_tool_usage_details_to_usage(usage: Usage, details: Mapping[str, object] | None) -> None:
+    """
+    Attach server_side_tool_usage_details and mirror web_search_calls onto
+    prompt_tokens_details.web_search_requests for built-in tool cost gating.
+    """
+    if details is None:
+        return
+    usage.server_side_tool_usage_details = details  # pyright: ignore[reportAttributeAccessIssue]  # extra  # rebind-ok: extras
+    try:
+        web_search_calls: Final = int(details.get("web_search_calls") or 0)
+    except (TypeError, ValueError):
+        return
+    if web_search_calls <= 0:
+        return
+    prompt_tokens_details: Final = usage.prompt_tokens_details or PromptTokensDetailsWrapper()
+    prompt_tokens_details.web_search_requests = web_search_calls
+    usage.prompt_tokens_details = prompt_tokens_details  # rebind-ok: write details onto caller usage
+
 
 def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
     """
@@ -32,9 +55,11 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
     prompt_tokens: Final = int(getattr(usage, "prompt_tokens", 0) or 0)
     completion_tokens: Final = int(getattr(usage, "completion_tokens", 0) or 0)
     total_tokens: Final = int(getattr(usage, "total_tokens", 0) or 0)
-    reasoning_tokens = 0
-    if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details:
-        reasoning_tokens = int(getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0)
+    reasoning_tokens: Final = (
+        int(getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0)
+        if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details
+        else 0
+    )
 
     already_normalised: Final = total_tokens == prompt_tokens + completion_tokens
     total_completion_tokens: Final = completion_tokens if already_normalised else completion_tokens + reasoning_tokens
@@ -52,33 +77,48 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
     return prompt_cost, completion_cost
 
 
-def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> float:
+def _web_search_cost_per_call_from_model_info(model_info: "ModelInfo") -> float:
     """
-    Calculate the cost of web search requests for X.AI models.
+    Per-invocation web_search price from model_info when configured.
 
-    X.AI Live Search costs $25 per 1,000 sources used.
-    Each source costs $0.025.
-
-    The number of sources is stored in prompt_tokens_details.web_search_requests
-    by the transformation layer to be compatible with the existing detection system.
+    Prefer ``search_context_cost_per_query`` (same shape as Gemini/Anthropic web
+    search pricing in the model cost map). Fall back to current xAI list pricing.
     """
-    # Cost per source used: $25 per 1,000 sources = $0.025 per source
-    cost_per_source: Final = 25.0 / 1000.0  # $0.025
-
-    num_sources_used = 0
-
-    if (
-        hasattr(usage, "prompt_tokens_details")
-        and usage.prompt_tokens_details is not None
-        and hasattr(usage.prompt_tokens_details, "web_search_requests")
-        and usage.prompt_tokens_details.web_search_requests is not None
+    search_costs: Final = model_info.get("search_context_cost_per_query")
+    if not isinstance(search_costs, Mapping):
+        return _DEFAULT_WEB_SEARCH_COST_PER_CALL
+    for key in (
+        "search_context_size_medium",
+        "search_context_size_low",
+        "search_context_size_high",
     ):
-        num_sources_used = int(usage.prompt_tokens_details.web_search_requests)
+        value = search_costs.get(key)
+        if value is None:
+            continue
+        try:
+            cost = float(value)
+        except (TypeError, ValueError):
+            continue
+        if cost > 0:
+            return cost
+    return _DEFAULT_WEB_SEARCH_COST_PER_CALL
 
-    # Fallback: try to get from num_sources_used if set directly
-    elif hasattr(usage, "num_sources_used") and usage.num_sources_used is not None:
-        num_sources_used = int(usage.num_sources_used)
 
-    total_cost: Final = cost_per_source * num_sources_used
+def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> float:
+    """
+    Calculate the cost of web search requests for X.AI models.
 
-    return total_cost
+    Counts invocations from usage.server_side_tool_usage_details.web_search_calls.
+    Per-call rate comes from model_info.search_context_cost_per_query when set,
+    otherwise the default xAI tools rate ($5 / 1k calls).
+    """
+    details: Final = getattr(usage, "server_side_tool_usage_details", None)
+    if not isinstance(details, Mapping):
+        return 0.0
+    try:
+        web_search_calls: Final = int(details.get("web_search_calls") or 0)
+    except (TypeError, ValueError):
+        return 0.0
+    if web_search_calls <= 0:
+        return 0.0
+    return _web_search_cost_per_call_from_model_info(model_info) * web_search_calls

litellm/llms/xai/responses/transformation.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Final
+from typing import Any, Final
 
 import litellm
 from litellm._logging import verbose_logger
@@ -12,13 +12,6 @@
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import LlmProviders
 
-if TYPE_CHECKING:
-    from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
-
-    LiteLLMLoggingObj = _LiteLLMLoggingObj
-else:
-    LiteLLMLoggingObj = Any
-
 
 class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
     """

litellm/responses/utils.py

@@ -1033,20 +1033,29 @@ def _is_response_api_usage(usage: dict | ResponseAPIUsage) -> bool:
 
     @staticmethod
     def _transform_response_api_usage_to_chat_usage(
-        usage_input: dict | ResponseAPIUsage | None,
+        usage_input: Mapping[str, object] | ResponseAPIUsage | Usage | None,
     ) -> Usage:
         """
         Transforms ResponseAPIUsage or ImageUsage to a Usage object.
 
         Both have the same spec with input_tokens, output_tokens, and
         input_tokens_details (text_tokens, image_tokens).
+
+        Usage inputs are returned as-is so re-running this helper never drops
+        fields. Non-standard provider fields (e.g. xAI's
+        server_side_tool_usage_details) are carried onto the returned Usage so
+        provider cost calculators can read them after normalization.
         """
         if usage_input is None:
             return Usage(
                 prompt_tokens=0,
                 completion_tokens=0,
                 total_tokens=0,
             )
+        if isinstance(usage_input, Usage):
+            return usage_input
+        if isinstance(usage_input, dict) and not ResponseAPILoggingUtils._is_response_api_usage(usage_input):
+            return Usage(**usage_input)
         response_api_usage: ResponseAPIUsage
         if isinstance(usage_input, dict):
             usage_input = dict(usage_input)  # shallow copy; avoid mutating caller
@@ -1055,13 +1064,11 @@ def _transform_response_api_usage_to_chat_usage(
                 usage_input["input_tokens_details"] = usage_input["input_token_details"]
             if usage_input.get("output_tokens_details") is None and "output_token_details" in usage_input:
                 usage_input["output_tokens_details"] = usage_input["output_token_details"]
-            total_tokens = usage_input.get("total_tokens")
-            if total_tokens is None:
+            if usage_input.get("total_tokens") is None:
                 input_tokens: Final = usage_input.get("input_tokens")
                 output_tokens: Final = usage_input.get("output_tokens")
-                if input_tokens is not None and output_tokens is not None:
-                    total_tokens = input_tokens + output_tokens
-                    usage_input["total_tokens"] = total_tokens
+                if isinstance(input_tokens, int) and isinstance(output_tokens, int):
+                    usage_input["total_tokens"] = input_tokens + output_tokens
             response_api_usage = ResponseAPIUsage(**usage_input)
         else:
             response_api_usage = usage_input
@@ -1089,12 +1096,27 @@ def _transform_response_api_usage_to_chat_usage(
                 audio_tokens=getattr(output_tokens_details, "audio_tokens", None),
             )
 
+        extra_usage_fields: Final = {
+            key: value
+            for key, value in (response_api_usage.model_extra or {}).items()
+            if key
+            not in (
+                "input_token_details",
+                "output_token_details",
+                "prompt_tokens",
+                "completion_tokens",
+                "total_tokens",
+                "prompt_tokens_details",
+                "completion_tokens_details",
+            )
+        }
         chat_usage: Final = Usage(
             prompt_tokens=prompt_tokens,
             completion_tokens=completion_tokens,
             total_tokens=prompt_tokens + completion_tokens,
             prompt_tokens_details=prompt_tokens_details,
             completion_tokens_details=completion_tokens_details,
+            **extra_usage_fields,
         )
 
         # Preserve cost attribute if it exists on ResponseAPIUsage