From 23cb897f550f37a4cb59a810f987f763c9ff3d6a Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Sat, 29 Aug 2026 17:09:52 +0800 Subject: [PATCH] fix(providers): recover incomplete Grok searches --- nanobot/providers/xai_grok_provider.py | 215 +++++++++++++++++----- tests/providers/test_xai_grok_provider.py | 157 +++++++++++++++- 2 files changed, 330 insertions(+), 42 deletions(-) diff --git a/nanobot/providers/xai_grok_provider.py b/nanobot/providers/xai_grok_provider.py index 504e9eceb..011fbdf99 100644 --- a/nanobot/providers/xai_grok_provider.py +++ b/nanobot/providers/xai_grok_provider.py @@ -35,6 +35,7 @@ from nanobot.providers.xai_oauth import ( ) DEFAULT_XAI_GROK_URL = "https://cli-chat-proxy.grok.com/v1/responses" +_HOSTED_SEARCH_MAX_TURNS = 5 _MAX_ERROR_BODY_CHARS = 1000 _SENSITIVE_ERROR_KEYS = { "accesstoken", @@ -61,6 +62,10 @@ def _is_named_x_search_tool(value: object) -> bool: class XAIGrokProvider(LLMProvider): """Call xAI's subscription proxy and expose supported hosted tools.""" + # An incomplete hosted-tool stream can already have emitted answer text. Let the + # provider close that stream segment before its one bounded recovery attempt. + supports_stream_recover_callback = True + def __init__( self, default_model: str = DEFAULT_XAI_GROK_MODEL, @@ -100,6 +105,7 @@ class XAIGrokProvider(LLMProvider): on_content_delta: Callable[[str], Awaitable[None]] | None = None, on_thinking_delta: Callable[[str], Awaitable[None]] | None = None, on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None, + on_stream_recover: Callable[[], Awaitable[None]] | None = None, ) -> LLMResponse: wire_model = _strip_model_prefix(model or self.default_model) system_prompt, input_items = convert_messages(messages) @@ -130,6 +136,8 @@ class XAIGrokProvider(LLMProvider): if supports_backend_search: converted_tools.append({"type": "x_search"}) + hosted_search_enabled = supports_backend_search or configured_hosted_search + body: dict[str, Any] = { "model": wire_model, "store": False, @@ -145,6 +153,11 @@ class XAIGrokProvider(LLMProvider): "temperature": temperature, "reasoning": _build_reasoning_options(reasoning_effort), } + if hosted_search_enabled: + # xAI's global default is intentionally unspecified. Five turns is + # their documented balanced setting and prevents a search from + # stopping after a single unsuccessful lookup. + body["max_turns"] = _HOSTED_SEARCH_MAX_TURNS if self._extra_body: body.update({ key: value @@ -156,40 +169,54 @@ class XAIGrokProvider(LLMProvider): headers = _build_headers(token.access, wire_model) stage = "xai_request" - try: - result = await _request_xai( - DEFAULT_XAI_GROK_URL, - headers, - body, - proxy=self.proxy, - on_content_delta=on_content_delta, - on_thinking_delta=on_thinking_delta, - on_tool_call_delta=on_tool_call_delta, - ) - except _XAIHTTPError as exc: - if exc.status_code != 401: - raise - stage = "oauth_refresh" - token = await asyncio.to_thread( - get_xai_oauth_token, - proxy=self.proxy, - force_refresh=True, - ) - self._model_capabilities = None - self._model_capabilities_fetched_at = 0.0 - headers = _build_headers(token.access, wire_model) - stage = "xai_request_retry" - result = await _request_xai( - DEFAULT_XAI_GROK_URL, - headers, - body, - proxy=self.proxy, - on_content_delta=on_content_delta, - on_thinking_delta=on_thinking_delta, - on_tool_call_delta=on_tool_call_delta, - ) + auth_retried = False + hosted_tool_retried = False + retry_usage: LLMUsage | None = None + while True: + try: + result = await _request_xai( + DEFAULT_XAI_GROK_URL, + headers, + body, + proxy=self.proxy, + on_content_delta=on_content_delta, + on_thinking_delta=on_thinking_delta, + on_tool_call_delta=on_tool_call_delta, + ) + break + except _XAIHTTPError as exc: + if exc.status_code != 401 or auth_retried: + raise + auth_retried = True + stage = "oauth_refresh" + token = await asyncio.to_thread( + get_xai_oauth_token, + proxy=self.proxy, + force_refresh=True, + ) + headers = _build_headers(token.access, wire_model) + stage = "xai_request_after_oauth_refresh" + except _XAIIncompleteHostedToolError as exc: + retry_usage = _combine_usage(retry_usage, exc.usage) + cannot_recover_stream = ( + exc.stream_output_emitted + and on_stream_recover is None + ) + if hosted_tool_retried or cannot_recover_stream: + exc.usage = retry_usage + raise + hosted_tool_retried = True + stage = "hosted_tool_recovery" + logger.warning( + "xAI response ended with unfinished hosted tool(s): {}; " + "retrying once", + ", ".join(exc.tool_names), + ) + if on_stream_recover is not None: + await on_stream_recover() content, tool_calls, finish_reason, usage, reasoning_content = result + usage = _combine_usage(retry_usage, usage) return LLMResponse( content=content, tool_calls=tool_calls, @@ -238,6 +265,7 @@ class XAIGrokProvider(LLMProvider): on_content_delta: Callable[[str], Awaitable[None]] | None = None, on_thinking_delta: Callable[[str], Awaitable[None]] | None = None, on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None, + on_stream_recover: Callable[[], Awaitable[None]] | None = None, ) -> LLMResponse: return await self._call_xai( messages, @@ -250,6 +278,7 @@ class XAIGrokProvider(LLMProvider): on_content_delta, on_thinking_delta, on_tool_call_delta, + on_stream_recover, ) def get_default_model(self) -> str: @@ -269,6 +298,14 @@ def _build_reasoning_options(reasoning_effort: str | None) -> dict[str, str]: return options +def _combine_usage(left: LLMUsage | None, right: LLMUsage | None) -> LLMUsage | None: + if left is None: + return right + if right is None: + return left + return left + right + + def _build_headers(token: str, model: str) -> dict[str, str]: conversation_id = str(uuid.uuid4()) return { @@ -310,6 +347,31 @@ class _XAIHTTPError(RuntimeError): self.response_body = response_body +class _XAIIncompleteHostedToolError(RuntimeError): + """A nominally successful xAI stream ended before a hosted tool did.""" + + should_retry = False # _call_xai already performs the one safe recovery attempt. + + def __init__( + self, + active_tools: list[dict[str, Any]], + *, + usage: LLMUsage | None, + stream_output_emitted: bool = False, + ) -> None: + names = [ + str(event.get("name") or "hosted_tool") + for event in active_tools + ] + super().__init__( + "xAI ended the response before its hosted tool completed: " + + ", ".join(names) + ) + self.tool_names = tuple(names) + self.usage = usage + self.stream_output_emitted = stream_output_emitted + + async def _request_xai( url: str, headers: dict[str, str], @@ -320,10 +382,39 @@ async def _request_xai( on_thinking_delta: Callable[[str], Awaitable[None]] | None = None, on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None, ) -> tuple[str, list[ToolCallRequest], str, LLMUsage | None, str | None]: + active_hosted_tools: dict[str, dict[str, Any]] = {} + stream_output_emitted = False + + async def _forward_content_delta(delta: str) -> None: + nonlocal stream_output_emitted + if delta: + stream_output_emitted = True + if on_content_delta is not None: + await on_content_delta(delta) + + async def _forward_thinking_delta(delta: str) -> None: + nonlocal stream_output_emitted + if delta: + stream_output_emitted = True + if on_thinking_delta is not None: + await on_thinking_delta(delta) + + async def _track_and_forward_tool_event(event: dict[str, Any]) -> None: + if event.get("kind") == "hosted_tool": + call_id = event.get("call_id") + if call_id: + call_id = str(call_id) + if event.get("phase") == "start": + active_hosted_tools[call_id] = dict(event) + elif event.get("phase") in {"end", "error"}: + active_hosted_tools.pop(call_id, None) + if on_tool_call_delta is not None: + await on_tool_call_delta(event) + async def _on_response_event(event: dict[str, Any]) -> None: hosted_event = _xai_hosted_tool_event(event) - if hosted_event is not None and on_tool_call_delta is not None: - await on_tool_call_delta(hosted_event) + if hosted_event is not None: + await _track_and_forward_tool_event(hosted_event) client_kwargs: dict[str, Any] = {"timeout": resolve_stream_idle_timeout_s()} if proxy: @@ -334,13 +425,34 @@ async def _request_xai( content = await response.aread() raw = content.decode("utf-8", "ignore") raise _build_xai_http_error(response.status_code, response.headers, raw) - return await consume_sse_with_reasoning( + result = await consume_sse_with_reasoning( response, - on_content_delta=on_content_delta, - on_tool_call_delta=on_tool_call_delta, - on_reasoning_delta=on_thinking_delta, - on_response_event=_on_response_event if on_tool_call_delta else None, + on_content_delta=( + _forward_content_delta if on_content_delta is not None else None + ), + # Always observe tool events so protocol validation also works for + # non-streaming callers that did not request UI progress callbacks. + on_tool_call_delta=_track_and_forward_tool_event, + on_reasoning_delta=( + _forward_thinking_delta if on_thinking_delta is not None else None + ), + on_response_event=_on_response_event, ) + if result[2] != "error" and active_hosted_tools: + active = list(active_hosted_tools.values()) + for event in active: + await _track_and_forward_tool_event({ + **event, + "phase": "error", + "result": None, + "error": "xAI ended the response before this hosted tool completed.", + }) + raise _XAIIncompleteHostedToolError( + active, + usage=result[3], + stream_output_emitted=stream_output_emitted, + ) + return result def _xai_hosted_tool_event(event: dict[str, Any]) -> dict[str, Any] | None: @@ -360,13 +472,31 @@ def _xai_hosted_tool_event(event: dict[str, Any]) -> dict[str, Any] | None: "result": None, } - if event_type != "response.output_item.done": + if event_type not in {"response.output_item.added", "response.output_item.done"}: return None item = event.get("item") if not isinstance(item, dict): return None item = cast(dict[str, Any], item) - if item.get("type") != "custom_tool_call": + item_type = item.get("type") + if item_type == "x_search_call": + call_id = item.get("id") or item.get("call_id") or event.get("item_id") + if not call_id: + return None + phase = "start" if event_type == "response.output_item.added" else "end" + return { + "kind": "hosted_tool", + "phase": phase, + "call_id": str(call_id), + "name": "x_search", + "arguments": _xai_hosted_tool_arguments(item.get("action")), + "result": ( + {"status": str(item.get("status") or "completed")} + if phase == "end" + else None + ), + } + if event_type != "response.output_item.done" or item_type != "custom_tool_call": return None tool_name = item.get("name") if not isinstance(tool_name, str) or not tool_name.startswith("x_"): @@ -490,6 +620,8 @@ def _xai_error_response(exc: Exception) -> LLMResponse: should_retry = True if should_retry is None else should_retry elif isinstance(exc, _XAIHTTPError): error_kind = "http" + elif isinstance(exc, _XAIIncompleteHostedToolError): + error_kind = "provider" if status_code is not None and should_retry is None: should_retry = _should_retry_status( int(status_code), @@ -502,6 +634,7 @@ def _xai_error_response(exc: Exception) -> LLMResponse: return LLMResponse( content=f"Error calling xAI ({type(exc).__name__}): {message}", finish_reason="error", + usage=getattr(exc, "usage", None), retry_after=retry_after, error_status_code=int(status_code) if status_code is not None else None, error_kind=error_kind, diff --git a/tests/providers/test_xai_grok_provider.py b/tests/providers/test_xai_grok_provider.py index abae49320..af89a307a 100644 --- a/tests/providers/test_xai_grok_provider.py +++ b/tests/providers/test_xai_grok_provider.py @@ -23,6 +23,7 @@ from nanobot.providers.xai_grok_provider import ( _request_xai, _xai_error_response, _XAIHTTPError, + _XAIIncompleteHostedToolError, ) @@ -147,6 +148,7 @@ async def test_provider_injects_hosted_x_search_and_required_proxy_headers(monke assert body["stream_tool_calls"] is True assert body["reasoning"] == {"summary": "concise", "effort": "high"} assert body["store"] is False + assert body["max_turns"] == 5 assert headers["Authorization"] == "Bearer subscription-token" assert headers["X-XAI-Token-Auth"] == "xai-grok-cli" assert headers["x-authenticateresponse"] == "authenticate-response" @@ -258,6 +260,7 @@ async def test_explicit_empty_tools_disables_catalog_lookup_and_hosted_tool(monk "description": "Read a file", "parameters": {"type": "object"}, }] + assert "max_turns" not in bodies[0] @pytest.mark.asyncio @@ -296,6 +299,8 @@ async def test_provider_keeps_local_x_search_when_model_does_not_support_hosted_ "parameters": {"type": "object"}, } ] + assert "max_turns" not in bodies[0] + assert bodies[0]["instructions"] == "" @pytest.mark.asyncio @@ -381,7 +386,10 @@ async def test_factory_builds_xai_provider_and_applies_explicit_body_overrides(m "providers": { "xaiGrok": { "proxy": "http://127.0.0.1:7890", - "extraBody": {"parallel_tool_calls": False}, + "extraBody": { + "parallel_tool_calls": False, + "max_turns": 2, + }, } }, } @@ -394,6 +402,7 @@ async def test_factory_builds_xai_provider_and_applies_explicit_body_overrides(m assert provider.proxy == "http://127.0.0.1:7890" assert response.content == "ok" assert bodies[0]["parallel_tool_calls"] is False + assert bodies[0]["max_turns"] == 2 assert {"type": "x_search"} in bodies[0]["tools"] @@ -513,6 +522,152 @@ async def test_raw_response_request_streams_hosted_x_search_lifecycle(monkeypatc assert "large hosted result" not in json.dumps(tool_events) +@pytest.mark.asyncio +async def test_raw_response_request_streams_official_x_search_lifecycle(monkeypatch) -> None: + original_client = httpx.AsyncClient + events = [ + { + "type": "response.output_item.added", + "item": { + "type": "x_search_call", + "id": "x-search-1", + "status": "in_progress", + "action": {"query": "nanobot oauth"}, + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "x_search_call", + "id": "x-search-1", + "status": "completed", + "action": {"query": "nanobot oauth"}, + }, + }, + { + "type": "response.completed", + "response": {"status": "completed", "usage": {}}, + }, + ] + content = "".join(f"data: {json.dumps(event)}\n\n" for event in events) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, content=content, request=request) + + def fake_client(**kwargs) -> httpx.AsyncClient: + return original_client( + transport=httpx.MockTransport(handler), + timeout=kwargs["timeout"], + ) + + monkeypatch.setattr("nanobot.providers.xai_grok_provider.httpx.AsyncClient", fake_client) + tool_events: list[dict[str, Any]] = [] + + await _request_xai( + "https://cli-chat-proxy.grok.com/v1/responses", + _build_headers("secret", "grok-4.6"), + {"model": "grok-4.6", "tools": [{"type": "x_search"}]}, + on_tool_call_delta=lambda event: _append(tool_events, event), + ) + + assert [(event["phase"], event["name"]) for event in tool_events] == [ + ("start", "x_search"), + ("end", "x_search"), + ] + assert tool_events[-1]["result"] == {"status": "completed"} + + +@pytest.mark.asyncio +async def test_raw_response_rejects_unfinished_hosted_tool_and_closes_progress( + monkeypatch, +) -> None: + original_client = httpx.AsyncClient + events = [ + { + "type": "response.custom_tool_call_input.done", + "item_id": "x-search-1", + "input": '{"query":"nanobot oauth"}', + }, + {"type": "response.output_text.delta", "delta": "I will keep searching."}, + { + "type": "response.completed", + "response": { + "status": "completed", + "usage": {"input_tokens": 8, "output_tokens": 4, "total_tokens": 12}, + }, + }, + ] + content = "".join(f"data: {json.dumps(event)}\n\n" for event in events) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, content=content, request=request) + + def fake_client(**kwargs) -> httpx.AsyncClient: + return original_client( + transport=httpx.MockTransport(handler), + timeout=kwargs["timeout"], + ) + + monkeypatch.setattr("nanobot.providers.xai_grok_provider.httpx.AsyncClient", fake_client) + tool_events: list[dict[str, Any]] = [] + + with pytest.raises(_XAIIncompleteHostedToolError) as caught: + await _request_xai( + "https://cli-chat-proxy.grok.com/v1/responses", + _build_headers("secret", "grok-4.6"), + {"model": "grok-4.6", "tools": [{"type": "x_search"}]}, + on_tool_call_delta=lambda event: _append(tool_events, event), + ) + + assert caught.value.usage == LLMUsage.reported(input_tokens=8, output_tokens=4) + assert [event["phase"] for event in tool_events] == ["start", "error"] + assert "before this hosted tool completed" in tool_events[-1]["error"] + + +@pytest.mark.asyncio +async def test_provider_recovers_unfinished_hosted_tool_once_and_preserves_usage( + monkeypatch, +) -> None: + _mock_token(monkeypatch) + _mock_model_capabilities(monkeypatch, supports_backend_search=True) + attempts = 0 + streamed: list[str] = [] + recovered: list[bool] = [] + first_usage = LLMUsage.reported(input_tokens=10, output_tokens=2) + second_usage = LLMUsage.reported(input_tokens=11, output_tokens=4) + + async def fake_request(_url, _headers, body, **kwargs): + nonlocal attempts + attempts += 1 + assert body["max_turns"] == 5 + if attempts == 1: + await kwargs["on_content_delta"]("I will keep searching.") + raise _XAIIncompleteHostedToolError( + [{"name": "x_search", "call_id": "search-1"}], + usage=first_usage, + ) + await kwargs["on_content_delta"]("Final researched answer.") + return "Final researched answer.", [], "stop", second_usage, None + + async def on_recover() -> None: + recovered.append(True) + + monkeypatch.setattr("nanobot.providers.xai_grok_provider._request_xai", fake_request) + provider = XAIGrokProvider() + + response = await provider.chat_stream_with_retry( + [{"role": "user", "content": "Search X"}], + on_content_delta=lambda delta: _append(streamed, delta), + on_stream_recover=on_recover, + ) + + assert attempts == 2 + assert recovered == [True] + assert streamed == ["I will keep searching.", "Final researched answer."] + assert response.content == "Final researched answer." + assert response.usage == first_usage + second_usage + + @pytest.mark.asyncio async def test_raw_response_error_preserves_bounded_redacted_body(monkeypatch) -> None: original_client = httpx.AsyncClient