fix(xai): cap search results to the caller's max_results

transform_search_response returned every distinct citation regardless
of the caller's max_results, unlike every other search provider. Caps
the returned list while still pricing the per-citation surcharge off
the true distinct citation count, since xAI consulted all of them
internally regardless of how many are returned.

litellm/llms/xai/search/transformation.py

@@ -163,6 +163,19 @@ def _requested_model(response_kwargs: Mapping[str, object]) -> str:
     return _model(optional_params)
 
 
+def _valid_max_results(max_results: object) -> int | None:
+    if isinstance(max_results, bool) or not isinstance(max_results, int):
+        return None
+    return max_results if max_results > 0 else None
+
+
+def _requested_max_results(response_kwargs: Mapping[str, object]) -> int | None:
+    optional_params: Final = response_kwargs.get("optional_params")
+    if not isinstance(optional_params, Mapping):
+        return None
+    return _valid_max_results(optional_params.get("max_results"))
+
+
 def _x_search_tool(optional_params: Mapping[str, object]) -> _XSearchTool:
     return _XSearchTool(
         allowed_x_handles=_typed_str_tuple(optional_params.get("allowed_x_handles")),
@@ -325,19 +338,22 @@ class XAISearchConfig(BaseSearchConfig):
                 else "unknown reason"
             )
             raise self._upstream_error(f"the search was incomplete: {reason}", raw_response)
-        return self._priced(results, parsed.usage, _requested_model(kwargs))
+        max_results: Final = _requested_max_results(kwargs)
+        capped_results: Final = results[:max_results] if max_results is not None else results
+        return self._priced(capped_results, parsed.usage, _requested_model(kwargs), len(results))
 
     def _priced(
         self,
         results: tuple[SearchResult, ...],
         usage: _Usage | None,
         model: str,
+        distinct_citation_count: int,
     ) -> SearchResponse:
         response: Final = SearchResponse(
             results=list(results),  # mutable-ok: SearchResponse.results is list[SearchResult]
             object="search",
         )
-        cost: Final = _resolved_cost(usage, model, len(results))
+        cost: Final = _resolved_cost(usage, model, distinct_citation_count)
         if cost is not None:
             response._hidden_params[  # pyright: ignore[reportPrivateUsage]  # response_cost_calculator's own contract
                 "additional_headers"

tests/test_litellm/llms/xai/search/test_xai_search_transformation.py

@@ -317,6 +317,30 @@ class TestXAISearchConfigTransformResponse:
         resp = _config().transform_search_response(_resp(payload), logging_obj=Mock())
         assert [r.url for r in resp.results] == ["https://example.com"]
 
+    def test_caps_results_to_max_results(self):
+        annotations = [_citation(f"https://example.com/{i}", f"T{i}") for i in range(5)]
+        resp = _config().transform_search_response(
+            _resp(_message_response("claim", annotations)), logging_obj=Mock(), optional_params={"max_results": 2}
+        )
+        assert [r.url for r in resp.results] == ["https://example.com/0", "https://example.com/1"]
+
+    def test_without_max_results_returns_all_citations(self):
+        annotations = [_citation(f"https://example.com/{i}", f"T{i}") for i in range(4)]
+        resp = _config().transform_search_response(
+            _resp(_message_response("claim", annotations)), logging_obj=Mock()
+        )
+        assert len(resp.results) == 4
+
+    @pytest.mark.parametrize("max_results", [True, False, 0, -1, "5"])
+    def test_ignores_invalid_max_results_and_returns_all_citations(self, max_results: object):
+        annotations = [_citation(f"https://example.com/{i}", f"T{i}") for i in range(3)]
+        resp = _config().transform_search_response(
+            _resp(_message_response("claim", annotations)),
+            logging_obj=Mock(),
+            optional_params={"max_results": max_results},
+        )
+        assert len(resp.results) == 3
+
 
 class TestXAISearchConfigCost:
     def test_uses_xai_reported_ticks_when_present(self):
@@ -417,6 +441,27 @@ class TestXAISearchConfigCost:
             "llm_provider-x-litellm-response-cost"
         ] == pytest.approx(transformation._PER_CITATION_SURCHARGE_USD * 2)
 
+    def test_surcharge_uses_full_citation_count_even_when_results_are_capped(self, monkeypatch: pytest.MonkeyPatch):
+        def fake_get_model_info(model: str, custom_llm_provider: str):
+            if model == "xai/grok-4-fast-non-reasoning":
+                return {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0}
+            raise Exception("not mapped")
+
+        monkeypatch.setattr(transformation, "get_model_info", fake_get_model_info)
+        annotations = [_citation(f"https://example.com/{i}", f"T{i}") for i in range(5)]
+        payload = _message_response(
+            "claim",
+            annotations,
+            usage={"input_tokens": 0, "output_tokens": 0, "output_tokens_details": {"reasoning_tokens": 0}},
+        )
+        resp = _config().transform_search_response(
+            _resp(payload), logging_obj=Mock(), optional_params={"max_results": 2}
+        )
+        assert len(resp.results) == 2
+        assert resp._hidden_params["additional_headers"][
+            "llm_provider-x-litellm-response-cost"
+        ] == pytest.approx(transformation._PER_CITATION_SURCHARGE_USD * 5)
+
 
 class TestXAISearchConfigGetErrorClass:
     def test_attributes_the_provider(self):