fix(providers): recover incomplete Grok searches

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,

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