feat(providers): add xAI Grok OAuth with capability-gated X Search (#5035)
This commit is contained in:
chengyongru
@@ -7,43 +7,61 @@ import sys
|
||||
|
||||
|
||||
def test_importing_providers_package_is_lazy(monkeypatch) -> None:
|
||||
original_package = sys.modules["nanobot.providers"]
|
||||
monkeypatch.delitem(sys.modules, "nanobot.providers", raising=False)
|
||||
monkeypatch.delitem(sys.modules, "nanobot.providers.anthropic_provider", raising=False)
|
||||
monkeypatch.delitem(sys.modules, "nanobot.providers.openai_compat_provider", raising=False)
|
||||
monkeypatch.delitem(sys.modules, "nanobot.providers.openai_codex_provider", raising=False)
|
||||
monkeypatch.delitem(sys.modules, "nanobot.providers.xai_oauth", raising=False)
|
||||
monkeypatch.delitem(sys.modules, "nanobot.providers.xai_grok_provider", raising=False)
|
||||
monkeypatch.delitem(sys.modules, "nanobot.providers.github_copilot_provider", raising=False)
|
||||
monkeypatch.delitem(sys.modules, "nanobot.providers.azure_openai_provider", raising=False)
|
||||
monkeypatch.delitem(sys.modules, "nanobot.providers.bedrock_provider", raising=False)
|
||||
|
||||
providers = importlib.import_module("nanobot.providers")
|
||||
try:
|
||||
providers = importlib.import_module("nanobot.providers")
|
||||
|
||||
assert "nanobot.providers.anthropic_provider" not in sys.modules
|
||||
assert "nanobot.providers.openai_compat_provider" not in sys.modules
|
||||
assert "nanobot.providers.openai_codex_provider" not in sys.modules
|
||||
assert "nanobot.providers.github_copilot_provider" not in sys.modules
|
||||
assert "nanobot.providers.azure_openai_provider" not in sys.modules
|
||||
assert "nanobot.providers.bedrock_provider" not in sys.modules
|
||||
assert providers.__all__ == [
|
||||
"LLMProvider",
|
||||
"LLMResponse",
|
||||
"AnthropicProvider",
|
||||
"OpenAICompatProvider",
|
||||
"OpenAICodexProvider",
|
||||
"GitHubCopilotProvider",
|
||||
"AzureOpenAIProvider",
|
||||
"BedrockProvider",
|
||||
]
|
||||
assert "nanobot.providers.anthropic_provider" not in sys.modules
|
||||
assert "nanobot.providers.openai_compat_provider" not in sys.modules
|
||||
assert "nanobot.providers.openai_codex_provider" not in sys.modules
|
||||
assert "nanobot.providers.xai_oauth" not in sys.modules
|
||||
assert "nanobot.providers.xai_grok_provider" not in sys.modules
|
||||
assert "nanobot.providers.github_copilot_provider" not in sys.modules
|
||||
assert "nanobot.providers.azure_openai_provider" not in sys.modules
|
||||
assert "nanobot.providers.bedrock_provider" not in sys.modules
|
||||
assert providers.__all__ == [
|
||||
"LLMProvider",
|
||||
"LLMResponse",
|
||||
"AnthropicProvider",
|
||||
"OpenAICompatProvider",
|
||||
"OpenAICodexProvider",
|
||||
"XAIGrokProvider",
|
||||
"GitHubCopilotProvider",
|
||||
"AzureOpenAIProvider",
|
||||
"BedrockProvider",
|
||||
]
|
||||
finally:
|
||||
# Importing a replacement subpackage also replaces nanobot.providers on the
|
||||
# parent package. Restore both views so this isolation test cannot pollute
|
||||
# later tests that resolve a module through a dotted monkeypatch target.
|
||||
monkeypatch.undo()
|
||||
setattr(sys.modules["nanobot"], "providers", original_package)
|
||||
|
||||
|
||||
def test_explicit_provider_import_still_works(monkeypatch) -> None:
|
||||
original_package = sys.modules["nanobot.providers"]
|
||||
monkeypatch.delitem(sys.modules, "nanobot.providers", raising=False)
|
||||
monkeypatch.delitem(sys.modules, "nanobot.providers.anthropic_provider", raising=False)
|
||||
|
||||
namespace: dict[str, object] = {}
|
||||
exec("from nanobot.providers import AnthropicProvider", namespace)
|
||||
try:
|
||||
namespace: dict[str, object] = {}
|
||||
exec("from nanobot.providers import AnthropicProvider", namespace)
|
||||
|
||||
assert namespace["AnthropicProvider"].__name__ == "AnthropicProvider"
|
||||
assert "nanobot.providers.anthropic_provider" in sys.modules
|
||||
assert namespace["AnthropicProvider"].__name__ == "AnthropicProvider"
|
||||
assert "nanobot.providers.anthropic_provider" in sys.modules
|
||||
finally:
|
||||
monkeypatch.undo()
|
||||
setattr(sys.modules["nanobot"], "providers", original_package)
|
||||
|
||||
|
||||
def test_openai_codex_supports_progress_deltas() -> None:
|
||||
|
||||
@@ -0,0 +1,526 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from nanobot.config.schema import Config
|
||||
from nanobot.providers.factory import make_provider
|
||||
from nanobot.providers.registry import find_by_name
|
||||
from nanobot.providers.xai_grok_provider import (
|
||||
DEFAULT_XAI_GROK_MODEL,
|
||||
DEFAULT_XAI_GROK_MODELS_URL,
|
||||
XAIGrokProvider,
|
||||
_bounded_error_body,
|
||||
_build_headers,
|
||||
_build_model_headers,
|
||||
_build_reasoning_options,
|
||||
_build_xai_http_error,
|
||||
_fetch_xai_model_capabilities,
|
||||
_parse_xai_model_capabilities,
|
||||
_request_xai,
|
||||
_xai_error_response,
|
||||
_XAIHTTPError,
|
||||
)
|
||||
|
||||
|
||||
def _token(access: str = "subscription-token") -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
access=access,
|
||||
refresh="refresh-token",
|
||||
expires=int(time.time() * 1000) + 3_600_000,
|
||||
account_id="account",
|
||||
)
|
||||
|
||||
|
||||
def _mock_token(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"nanobot.providers.xai_grok_provider.get_xai_oauth_token",
|
||||
lambda **_kwargs: _token(),
|
||||
)
|
||||
|
||||
|
||||
def _mock_model_capabilities(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
supports_backend_search: bool,
|
||||
) -> None:
|
||||
async def fake_fetch(*_args, **_kwargs):
|
||||
return {"grok-4.5": supports_backend_search}
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.providers.xai_grok_provider._fetch_xai_model_capabilities",
|
||||
fake_fetch,
|
||||
)
|
||||
|
||||
|
||||
def test_xai_grok_registry_exposes_curated_x_search_model() -> None:
|
||||
spec = find_by_name("xai_grok")
|
||||
|
||||
assert spec is not None
|
||||
assert spec.is_oauth is True
|
||||
assert spec.backend == "xai_grok"
|
||||
assert spec.builtin_models[0].id == DEFAULT_XAI_GROK_MODEL
|
||||
assert spec.builtin_models[0].context_window == 500000
|
||||
assert "when supported" in spec.builtin_models[0].description
|
||||
|
||||
|
||||
def test_reasoning_options_omit_disabled_effort() -> None:
|
||||
assert _build_reasoning_options("none") == {"summary": "concise"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_injects_hosted_x_search_and_required_proxy_headers(monkeypatch) -> None:
|
||||
_mock_token(monkeypatch)
|
||||
_mock_model_capabilities(monkeypatch, supports_backend_search=True)
|
||||
calls: list[tuple[str, dict[str, str], dict[str, Any]]] = []
|
||||
|
||||
async def fake_request(url, headers, body, **_kwargs):
|
||||
calls.append((url, headers, body))
|
||||
return "answer [[1]](https://x.com/example/status/1)", [], "stop", {}, None
|
||||
|
||||
monkeypatch.setattr("nanobot.providers.xai_grok_provider._request_xai", fake_request)
|
||||
provider = XAIGrokProvider()
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_file",
|
||||
"description": "Read a file",
|
||||
"parameters": {"type": "object"},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "x_search",
|
||||
"description": "A colliding local tool",
|
||||
"parameters": {"type": "object"},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
response = await provider.chat(
|
||||
[{"role": "user", "content": "What is happening on X?"}],
|
||||
tools=tools,
|
||||
max_tokens=1234,
|
||||
temperature=0.2,
|
||||
reasoning_effort="high",
|
||||
)
|
||||
|
||||
assert response.content == "answer [[1]](https://x.com/example/status/1)"
|
||||
url, headers, body = calls[0]
|
||||
assert url == "https://cli-chat-proxy.grok.com/v1/responses"
|
||||
assert body["model"] == "grok-4.5"
|
||||
assert body["tools"] == [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "read_file",
|
||||
"description": "Read a file",
|
||||
"parameters": {"type": "object"},
|
||||
},
|
||||
{"type": "x_search"},
|
||||
]
|
||||
assert body["max_output_tokens"] == 1234
|
||||
assert body["temperature"] == 0.2
|
||||
assert body["stream_tool_calls"] is True
|
||||
assert body["reasoning"] == {"summary": "concise", "effort": "high"}
|
||||
assert body["store"] is False
|
||||
assert headers["Authorization"] == "Bearer subscription-token"
|
||||
assert headers["X-XAI-Token-Auth"] == "xai-grok-cli"
|
||||
assert headers["x-authenticateresponse"] == "authenticate-response"
|
||||
assert headers["x-grok-client-identifier"] == "nanobot"
|
||||
assert headers["x-grok-client-mode"] == "headless"
|
||||
assert headers["x-grok-model-override"] == "grok-4.5"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_keeps_local_x_search_when_model_does_not_support_hosted_search(
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
_mock_token(monkeypatch)
|
||||
_mock_model_capabilities(monkeypatch, supports_backend_search=False)
|
||||
bodies: list[dict[str, Any]] = []
|
||||
|
||||
async def fake_request(_url, _headers, body, **_kwargs):
|
||||
bodies.append(body)
|
||||
return "ok", [], "stop", {}, None
|
||||
|
||||
monkeypatch.setattr("nanobot.providers.xai_grok_provider._request_xai", fake_request)
|
||||
provider = XAIGrokProvider()
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "x_search",
|
||||
"description": "A local search fallback",
|
||||
"parameters": {"type": "object"},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
response = await provider.chat([{"role": "user", "content": "search"}], tools=tools)
|
||||
|
||||
assert response.content == "ok"
|
||||
assert bodies[0]["tools"] == [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "x_search",
|
||||
"description": "A local search fallback",
|
||||
"parameters": {"type": "object"},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_fails_closed_and_caches_model_catalog_failure(monkeypatch) -> None:
|
||||
_mock_token(monkeypatch)
|
||||
fetch_calls = 0
|
||||
bodies: list[dict[str, Any]] = []
|
||||
|
||||
async def failing_fetch(*_args, **_kwargs):
|
||||
nonlocal fetch_calls
|
||||
fetch_calls += 1
|
||||
raise httpx.ConnectError("catalog unavailable")
|
||||
|
||||
async def fake_request(_url, _headers, body, **_kwargs):
|
||||
bodies.append(body)
|
||||
return "ok", [], "stop", {}, None
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.providers.xai_grok_provider._fetch_xai_model_capabilities",
|
||||
failing_fetch,
|
||||
)
|
||||
monkeypatch.setattr("nanobot.providers.xai_grok_provider._request_xai", fake_request)
|
||||
provider = XAIGrokProvider()
|
||||
|
||||
await provider.chat([{"role": "user", "content": "first"}])
|
||||
await provider.chat([{"role": "user", "content": "second"}])
|
||||
|
||||
assert fetch_calls == 1
|
||||
assert all({"type": "x_search"} not in body["tools"] for body in bodies)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_refreshes_and_retries_exactly_once_after_401(monkeypatch) -> None:
|
||||
_mock_model_capabilities(monkeypatch, supports_backend_search=False)
|
||||
token_calls: list[tuple[str | None, bool]] = []
|
||||
|
||||
def fake_token(*, proxy=None, force_refresh=False):
|
||||
token_calls.append((proxy, force_refresh))
|
||||
return _token("fresh-token" if force_refresh else "stale-token")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.providers.xai_grok_provider.get_xai_oauth_token",
|
||||
fake_token,
|
||||
)
|
||||
request_tokens: list[str] = []
|
||||
|
||||
async def fake_request(_url, headers, _body, **_kwargs):
|
||||
request_tokens.append(headers["Authorization"])
|
||||
if len(request_tokens) == 1:
|
||||
raise _XAIHTTPError("unauthorized", status_code=401, should_retry=False)
|
||||
return "ok", [], "stop", {}, None
|
||||
|
||||
monkeypatch.setattr("nanobot.providers.xai_grok_provider._request_xai", fake_request)
|
||||
provider = XAIGrokProvider(proxy="http://127.0.0.1:7890")
|
||||
|
||||
response = await provider.chat([{"role": "user", "content": "hello"}])
|
||||
|
||||
assert response.content == "ok"
|
||||
assert token_calls == [
|
||||
("http://127.0.0.1:7890", False),
|
||||
("http://127.0.0.1:7890", True),
|
||||
]
|
||||
assert request_tokens == ["Bearer stale-token", "Bearer fresh-token"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_second_401_is_non_retryable_and_prompts_reauthentication(monkeypatch) -> None:
|
||||
_mock_token(monkeypatch)
|
||||
_mock_model_capabilities(monkeypatch, supports_backend_search=False)
|
||||
|
||||
async def always_unauthorized(*_args, **_kwargs):
|
||||
raise _XAIHTTPError(
|
||||
"xAI rejected the login. Sign in again with `nanobot provider login xai-grok`.",
|
||||
status_code=401,
|
||||
should_retry=False,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.providers.xai_grok_provider._request_xai",
|
||||
always_unauthorized,
|
||||
)
|
||||
provider = XAIGrokProvider()
|
||||
|
||||
response = await provider.chat([{"role": "user", "content": "hello"}])
|
||||
|
||||
assert response.finish_reason == "error"
|
||||
assert response.error_status_code == 401
|
||||
assert response.error_kind == "http"
|
||||
assert response.error_should_retry is False
|
||||
assert "nanobot provider login xai-grok" in (response.content or "")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_factory_builds_xai_provider_and_applies_explicit_body_overrides(monkeypatch) -> None:
|
||||
_mock_token(monkeypatch)
|
||||
_mock_model_capabilities(monkeypatch, supports_backend_search=True)
|
||||
bodies: list[dict[str, Any]] = []
|
||||
|
||||
async def fake_request(_url, _headers, body, **_kwargs):
|
||||
bodies.append(body)
|
||||
return "ok", [], "stop", {}, None
|
||||
|
||||
monkeypatch.setattr("nanobot.providers.xai_grok_provider._request_xai", fake_request)
|
||||
config = Config.model_validate(
|
||||
{
|
||||
"agents": {
|
||||
"defaults": {
|
||||
"model": "xai-grok/grok-4.5",
|
||||
"provider": "xai_grok",
|
||||
}
|
||||
},
|
||||
"providers": {
|
||||
"xaiGrok": {
|
||||
"proxy": "http://127.0.0.1:7890",
|
||||
"extraBody": {"parallel_tool_calls": False},
|
||||
}
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
provider = make_provider(config)
|
||||
response = await provider.chat([{"role": "user", "content": "hello"}])
|
||||
|
||||
assert isinstance(provider, XAIGrokProvider)
|
||||
assert provider.proxy == "http://127.0.0.1:7890"
|
||||
assert response.content == "ok"
|
||||
assert bodies[0]["parallel_tool_calls"] is False
|
||||
assert {"type": "x_search"} in bodies[0]["tools"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_response_request_streams_text_usage_and_inline_citations(monkeypatch) -> None:
|
||||
original_client = httpx.AsyncClient
|
||||
captured: dict[str, Any] = {}
|
||||
events = [
|
||||
{"type": "response.output_text.delta", "delta": "Live result "},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"delta": "[[1]](https://x.com/example/status/1)",
|
||||
},
|
||||
{
|
||||
"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:
|
||||
captured["request"] = request
|
||||
captured["json"] = json.loads(request.content)
|
||||
return httpx.Response(200, content=content, request=request)
|
||||
|
||||
def fake_client(**kwargs) -> httpx.AsyncClient:
|
||||
captured["kwargs"] = kwargs
|
||||
return original_client(
|
||||
transport=httpx.MockTransport(handler),
|
||||
timeout=kwargs["timeout"],
|
||||
)
|
||||
|
||||
monkeypatch.setattr("nanobot.providers.xai_grok_provider.httpx.AsyncClient", fake_client)
|
||||
deltas: list[str] = []
|
||||
|
||||
result = await _request_xai(
|
||||
"https://cli-chat-proxy.grok.com/v1/responses",
|
||||
_build_headers("secret", "grok-4.5"),
|
||||
{"model": "grok-4.5", "tools": [{"type": "x_search"}]},
|
||||
on_content_delta=lambda delta: _append(deltas, delta),
|
||||
)
|
||||
|
||||
assert result[0] == "Live result [[1]](https://x.com/example/status/1)"
|
||||
assert result[2] == "stop"
|
||||
assert result[3] == {"prompt_tokens": 8, "completion_tokens": 4, "total_tokens": 12}
|
||||
assert deltas == ["Live result ", "[[1]](https://x.com/example/status/1)"]
|
||||
assert captured["json"]["tools"] == [{"type": "x_search"}]
|
||||
|
||||
|
||||
def test_model_capabilities_follow_upstream_aliases_and_default_to_disabled() -> None:
|
||||
capabilities = _parse_xai_model_capabilities(
|
||||
{
|
||||
"data": [
|
||||
{"id": "grok-4.5", "supportsBackendSearch": False},
|
||||
{
|
||||
"model": "grok-search",
|
||||
"supports_backend_search": True,
|
||||
},
|
||||
{
|
||||
"modelId": "grok-meta",
|
||||
"_meta": {"supportsBackendSearch": True},
|
||||
},
|
||||
{"id": "grok-unknown"},
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
assert capabilities == {
|
||||
"grok-4.5": False,
|
||||
"grok-search": True,
|
||||
"grok-meta": True,
|
||||
"grok-unknown": False,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_capability_request_uses_subscription_headers(monkeypatch) -> None:
|
||||
original_client = httpx.AsyncClient
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured["request"] = request
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={"data": [{"id": "grok-search", "supportsBackendSearch": True}]},
|
||||
request=request,
|
||||
)
|
||||
|
||||
def fake_client(**kwargs) -> httpx.AsyncClient:
|
||||
captured["kwargs"] = kwargs
|
||||
return original_client(
|
||||
transport=httpx.MockTransport(handler),
|
||||
timeout=kwargs["timeout"],
|
||||
follow_redirects=kwargs["follow_redirects"],
|
||||
)
|
||||
|
||||
monkeypatch.setattr("nanobot.providers.xai_grok_provider.httpx.AsyncClient", fake_client)
|
||||
payload = base64.urlsafe_b64encode(
|
||||
json.dumps({"sub": "user-42", "email": "user@example.com"}).encode()
|
||||
).decode().rstrip("=")
|
||||
access_token = f"header.{payload}.signature"
|
||||
headers = _build_model_headers(_token(access_token))
|
||||
|
||||
capabilities = await _fetch_xai_model_capabilities(
|
||||
DEFAULT_XAI_GROK_MODELS_URL,
|
||||
headers,
|
||||
)
|
||||
|
||||
request = captured["request"]
|
||||
assert isinstance(request, httpx.Request)
|
||||
assert request.method == "GET"
|
||||
assert str(request.url) == DEFAULT_XAI_GROK_MODELS_URL
|
||||
assert request.headers["Authorization"] == f"Bearer {access_token}"
|
||||
assert request.headers["X-XAI-Token-Auth"] == "xai-grok-cli"
|
||||
assert request.headers["x-userid"] == "user-42"
|
||||
assert request.headers["x-email"] == "user@example.com"
|
||||
assert captured["kwargs"] == {"timeout": 10.0, "follow_redirects": False}
|
||||
assert capabilities == {"grok-search": True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_response_error_preserves_bounded_redacted_body(monkeypatch) -> None:
|
||||
original_client = httpx.AsyncClient
|
||||
raw = json.dumps(
|
||||
{
|
||||
"code": "invalid-argument",
|
||||
"message": "Hosted x_search is not supported by grok-4.5",
|
||||
"access_token": "must-not-leak",
|
||||
}
|
||||
)
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(400, content=raw, 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)
|
||||
|
||||
with pytest.raises(_XAIHTTPError) as caught:
|
||||
await _request_xai(
|
||||
"https://cli-chat-proxy.grok.com/v1/responses",
|
||||
_build_headers("secret", "grok-4.5"),
|
||||
{"model": "grok-4.5"},
|
||||
)
|
||||
|
||||
error = caught.value
|
||||
assert error.status_code == 400
|
||||
assert error.error_code == "invalid-argument"
|
||||
assert error.should_retry is False
|
||||
assert error.response_body == (
|
||||
'{"code":"invalid-argument","message":"Hosted x_search is not supported by '
|
||||
'grok-4.5","access_token":"[REDACTED]"}'
|
||||
)
|
||||
assert f"Response body: {error.response_body}" in str(error)
|
||||
assert "must-not-leak" not in str(error)
|
||||
|
||||
provider_response = _xai_error_response(error)
|
||||
assert provider_response.error_status_code == 400
|
||||
assert provider_response.error_code == "invalid-argument"
|
||||
assert error.response_body in (provider_response.content or "")
|
||||
|
||||
|
||||
def test_plain_error_body_is_single_line_and_bounded() -> None:
|
||||
detail = _bounded_error_body("Bearer secret-token\n" + "x" * 1100)
|
||||
|
||||
assert detail is not None
|
||||
assert detail.startswith("Bearer [REDACTED] ")
|
||||
assert detail.endswith("…")
|
||||
assert len(detail) == 1001
|
||||
|
||||
|
||||
def test_client_version_rejection_explains_update_and_preserves_body() -> None:
|
||||
raw = json.dumps(
|
||||
{
|
||||
"code": "upgrade-required",
|
||||
"message": "Client version 0.2.109 is no longer supported",
|
||||
}
|
||||
)
|
||||
|
||||
error = _build_xai_http_error(426, httpx.Headers(), raw)
|
||||
response = _xai_error_response(error)
|
||||
|
||||
assert error.status_code == 426
|
||||
assert error.should_retry is False
|
||||
assert error.response_body == (
|
||||
'{"code":"upgrade-required","message":"Client version 0.2.109 is no longer supported"}'
|
||||
)
|
||||
assert "xAI requires a newer Grok client version. Update nanobot and try again." in str(error)
|
||||
assert error.response_body in str(error)
|
||||
assert response.error_status_code == 426
|
||||
assert error.response_body in (response.content or "")
|
||||
|
||||
|
||||
def test_large_json_error_body_redacts_camel_case_credentials_before_bounding() -> None:
|
||||
detail = _bounded_error_body(
|
||||
json.dumps(
|
||||
{
|
||||
"accessToken": "access-must-not-leak",
|
||||
"refresh-token": "refresh-must-not-leak",
|
||||
"padding": "x" * 33_000,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
assert detail is not None
|
||||
assert '"accessToken":"[REDACTED]"' in detail
|
||||
assert '"refresh-token":"[REDACTED]"' in detail
|
||||
assert "access-must-not-leak" not in detail
|
||||
assert "refresh-must-not-leak" not in detail
|
||||
assert detail.endswith("…")
|
||||
assert len(detail) == 1001
|
||||
|
||||
|
||||
async def _append(target: list[str], value: str) -> None:
|
||||
target.append(value)
|
||||
@@ -0,0 +1,340 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from urllib.parse import parse_qs, urlencode, urlsplit
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import nanobot.providers.xai_oauth as xai_oauth
|
||||
from nanobot.providers.xai_oauth import (
|
||||
XAI_CLIENT_ID,
|
||||
XAI_OAUTH_SCOPES,
|
||||
XAIOAuthError,
|
||||
XAIToken,
|
||||
_build_authorize_url,
|
||||
_CallbackResult,
|
||||
_Discovery,
|
||||
_generate_pkce,
|
||||
_make_callback_server,
|
||||
_validate_xai_endpoint,
|
||||
complete_xai_oauth_login,
|
||||
get_xai_oauth_login_status,
|
||||
get_xai_oauth_storage_path,
|
||||
get_xai_oauth_token,
|
||||
login_xai_oauth,
|
||||
logout_xai_oauth,
|
||||
start_xai_oauth_login,
|
||||
)
|
||||
|
||||
|
||||
def _use_temp_credentials(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None:
|
||||
monkeypatch.setattr(xai_oauth, "get_data_dir", lambda: tmp_path)
|
||||
|
||||
|
||||
def test_authorize_url_uses_pkce_and_frozen_xai_scope_contract() -> None:
|
||||
verifier, challenge = _generate_pkce()
|
||||
|
||||
url = _build_authorize_url(
|
||||
"https://auth.x.ai/oauth2/authorize",
|
||||
redirect_uri="http://127.0.0.1:54321/callback",
|
||||
challenge=challenge,
|
||||
state="state-value",
|
||||
nonce="nonce-value",
|
||||
)
|
||||
|
||||
params = parse_qs(urlsplit(url).query)
|
||||
assert len(verifier) >= 43
|
||||
assert params == {
|
||||
"response_type": ["code"],
|
||||
"client_id": [XAI_CLIENT_ID],
|
||||
"redirect_uri": ["http://127.0.0.1:54321/callback"],
|
||||
"scope": [" ".join(XAI_OAUTH_SCOPES)],
|
||||
"code_challenge": [challenge],
|
||||
"code_challenge_method": ["S256"],
|
||||
"state": ["state-value"],
|
||||
"nonce": ["nonce-value"],
|
||||
"referrer": ["nanobot"],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"endpoint",
|
||||
[
|
||||
"http://auth.x.ai/oauth2/token",
|
||||
"https://evil.example/oauth2/token",
|
||||
"https://auth.x.ai:444/oauth2/token",
|
||||
"https://user@auth.x.ai/oauth2/token",
|
||||
"https://auth.x.ai:invalid/oauth2/token",
|
||||
],
|
||||
)
|
||||
def test_discovery_rejects_unsafe_endpoints(endpoint: str) -> None:
|
||||
with pytest.raises(XAIOAuthError, match="unsafe token endpoint"):
|
||||
_validate_xai_endpoint(endpoint, "token")
|
||||
|
||||
|
||||
def test_callback_server_accepts_only_matching_state_and_allows_accounts_origin() -> None:
|
||||
results: queue.Queue[_CallbackResult] = queue.Queue(maxsize=1)
|
||||
server = _make_callback_server("expected-state", results)
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
response = httpx.get(
|
||||
f"http://127.0.0.1:{server.server_port}/callback",
|
||||
params={"code": "one-time-code", "state": "expected-state"},
|
||||
headers={"Origin": "https://accounts.x.ai"},
|
||||
)
|
||||
result = results.get(timeout=1)
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join(timeout=2)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["access-control-allow-origin"] == "https://accounts.x.ai"
|
||||
assert response.headers["access-control-allow-private-network"] == "true"
|
||||
assert result == _CallbackResult(code="one-time-code", state="expected-state")
|
||||
assert "one-time-code" not in response.text
|
||||
|
||||
|
||||
def test_login_uses_random_loopback_callback_and_saves_separate_credentials(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
) -> None:
|
||||
_use_temp_credentials(monkeypatch, tmp_path)
|
||||
discovery = _Discovery(
|
||||
authorization_endpoint="https://auth.x.ai/oauth2/authorize",
|
||||
token_endpoint="https://auth.x.ai/oauth2/token",
|
||||
userinfo_endpoint="https://auth.x.ai/oauth2/userinfo",
|
||||
)
|
||||
monkeypatch.setattr(xai_oauth, "_discover", lambda _proxy: discovery)
|
||||
exchanged: dict[str, str] = {}
|
||||
|
||||
def fake_exchange(endpoint: str, **kwargs):
|
||||
exchanged.update(endpoint=endpoint, **kwargs)
|
||||
return {
|
||||
"access_token": "access-secret",
|
||||
"refresh_token": "refresh-secret",
|
||||
"expires_in": 3600,
|
||||
}
|
||||
|
||||
monkeypatch.setattr(xai_oauth, "_exchange_code", fake_exchange)
|
||||
monkeypatch.setattr(
|
||||
xai_oauth,
|
||||
"_fetch_account",
|
||||
lambda endpoint, access, proxy: "user@example.com",
|
||||
)
|
||||
opened_urls: list[str] = []
|
||||
|
||||
def complete_in_browser(authorize_url: str) -> bool:
|
||||
opened_urls.append(authorize_url)
|
||||
params = parse_qs(urlsplit(authorize_url).query)
|
||||
callback_url = params["redirect_uri"][0]
|
||||
callback_query = urlencode({"code": "auth-code", "state": params["state"][0]})
|
||||
response = httpx.get(f"{callback_url}?{callback_query}")
|
||||
assert response.status_code == 200
|
||||
return True
|
||||
|
||||
token = login_xai_oauth(
|
||||
print_fn=lambda _message: None,
|
||||
browser_opener=complete_in_browser,
|
||||
callback_timeout_s=1,
|
||||
)
|
||||
|
||||
assert token.account_id == "user@example.com"
|
||||
assert exchanged["code"] == "auth-code"
|
||||
assert exchanged["redirect_uri"].startswith("http://127.0.0.1:")
|
||||
assert exchanged["verifier"]
|
||||
assert opened_urls
|
||||
assert get_xai_oauth_storage_path() == tmp_path / "auth" / "xai.json"
|
||||
saved = json.loads(get_xai_oauth_storage_path().read_text(encoding="utf-8"))
|
||||
assert saved["access"] == "access-secret"
|
||||
assert saved["refresh"] == "refresh-secret"
|
||||
assert get_xai_oauth_login_status() == token
|
||||
|
||||
|
||||
def test_pending_login_accepts_authorization_code_from_remote_browser(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
) -> None:
|
||||
_use_temp_credentials(monkeypatch, tmp_path)
|
||||
discovery = _Discovery(
|
||||
authorization_endpoint="https://auth.x.ai/oauth2/authorize",
|
||||
token_endpoint="https://auth.x.ai/oauth2/token",
|
||||
userinfo_endpoint=None,
|
||||
)
|
||||
monkeypatch.setattr(xai_oauth, "_discover", lambda _proxy: discovery)
|
||||
exchanged: dict[str, str] = {}
|
||||
|
||||
def fake_exchange(endpoint: str, **kwargs):
|
||||
exchanged.update(endpoint=endpoint, **kwargs)
|
||||
return {"access_token": "remote-access", "expires_in": 3600}
|
||||
|
||||
monkeypatch.setattr(xai_oauth, "_exchange_code", fake_exchange)
|
||||
|
||||
flow = start_xai_oauth_login(timeout_s=5)
|
||||
try:
|
||||
params = parse_qs(urlsplit(flow.authorization_url).query)
|
||||
callback_url = params["redirect_uri"][0]
|
||||
|
||||
token = complete_xai_oauth_login(flow, "remote-code")
|
||||
finally:
|
||||
flow.cancel()
|
||||
|
||||
assert token is not None
|
||||
assert token.access == "remote-access"
|
||||
assert exchanged["code"] == "remote-code"
|
||||
assert exchanged["redirect_uri"] == callback_url
|
||||
assert get_xai_oauth_login_status() == token
|
||||
|
||||
|
||||
def test_expired_token_refreshes_once_and_persists_rotated_refresh_token(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
) -> None:
|
||||
_use_temp_credentials(monkeypatch, tmp_path)
|
||||
expired = XAIToken(
|
||||
access="old-access",
|
||||
refresh="old-refresh",
|
||||
expires=int(time.time() * 1000) - 1,
|
||||
account_id="account",
|
||||
)
|
||||
xai_oauth._write_token(expired)
|
||||
calls: list[tuple[XAIToken, str | None]] = []
|
||||
|
||||
def fake_refresh(token: XAIToken, proxy: str | None) -> XAIToken:
|
||||
calls.append((token, proxy))
|
||||
return XAIToken(
|
||||
access="new-access",
|
||||
refresh="new-refresh",
|
||||
expires=int(time.time() * 1000) + 3_600_000,
|
||||
account_id=token.account_id,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(xai_oauth, "_refresh_token", fake_refresh)
|
||||
|
||||
refreshed = get_xai_oauth_token(proxy="http://127.0.0.1:7890")
|
||||
|
||||
assert refreshed.access == "new-access"
|
||||
assert refreshed.refresh == "new-refresh"
|
||||
assert calls == [(expired, "http://127.0.0.1:7890")]
|
||||
assert get_xai_oauth_login_status() == refreshed
|
||||
|
||||
|
||||
def test_missing_credentials_returns_actionable_login_command(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
) -> None:
|
||||
_use_temp_credentials(monkeypatch, tmp_path)
|
||||
|
||||
with pytest.raises(XAIOAuthError, match="nanobot provider login xai-grok"):
|
||||
get_xai_oauth_token()
|
||||
|
||||
|
||||
def test_logout_waits_for_inflight_refresh_and_removes_rotated_token(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
) -> None:
|
||||
_use_temp_credentials(monkeypatch, tmp_path)
|
||||
expired = XAIToken(
|
||||
access="old-access",
|
||||
refresh="old-refresh",
|
||||
expires=int(time.time() * 1000) - 1,
|
||||
)
|
||||
xai_oauth._write_token(expired)
|
||||
refresh_started = threading.Event()
|
||||
finish_refresh = threading.Event()
|
||||
logout_started = threading.Event()
|
||||
logout_finished = threading.Event()
|
||||
refresh_errors: list[Exception] = []
|
||||
logout_results: list[bool] = []
|
||||
|
||||
def fake_refresh(token: XAIToken, _proxy: str | None) -> XAIToken:
|
||||
refresh_started.set()
|
||||
assert finish_refresh.wait(timeout=2)
|
||||
return XAIToken(
|
||||
access="new-access",
|
||||
refresh=token.refresh,
|
||||
expires=int(time.time() * 1000) + 3_600_000,
|
||||
)
|
||||
|
||||
def refresh() -> None:
|
||||
try:
|
||||
get_xai_oauth_token()
|
||||
except Exception as exc: # pragma: no cover - asserted below
|
||||
refresh_errors.append(exc)
|
||||
|
||||
def logout() -> None:
|
||||
logout_started.set()
|
||||
logout_results.append(logout_xai_oauth())
|
||||
logout_finished.set()
|
||||
|
||||
monkeypatch.setattr(xai_oauth, "_refresh_token", fake_refresh)
|
||||
refresh_thread = threading.Thread(target=refresh)
|
||||
refresh_thread.start()
|
||||
assert refresh_started.wait(timeout=2)
|
||||
|
||||
logout_thread = threading.Thread(target=logout)
|
||||
logout_thread.start()
|
||||
assert logout_started.wait(timeout=2)
|
||||
assert not logout_finished.wait(timeout=0.05)
|
||||
finish_refresh.set()
|
||||
refresh_thread.join(timeout=2)
|
||||
logout_thread.join(timeout=2)
|
||||
|
||||
assert not refresh_thread.is_alive()
|
||||
assert not logout_thread.is_alive()
|
||||
assert refresh_errors == []
|
||||
assert logout_results == [True]
|
||||
assert not get_xai_oauth_storage_path().exists()
|
||||
|
||||
|
||||
def test_refresh_does_not_restore_credentials_removed_after_its_initial_read(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
) -> None:
|
||||
_use_temp_credentials(monkeypatch, tmp_path)
|
||||
expired = XAIToken(
|
||||
access="old-access",
|
||||
refresh="old-refresh",
|
||||
expires=int(time.time() * 1000) - 1,
|
||||
)
|
||||
xai_oauth._write_token(expired)
|
||||
real_load_token = xai_oauth._load_token
|
||||
initial_read_finished = threading.Event()
|
||||
continue_refresh = threading.Event()
|
||||
refresh_errors: list[Exception] = []
|
||||
first_load = True
|
||||
|
||||
def gated_load_token() -> XAIToken | None:
|
||||
nonlocal first_load
|
||||
token = real_load_token()
|
||||
if first_load:
|
||||
first_load = False
|
||||
initial_read_finished.set()
|
||||
assert continue_refresh.wait(timeout=2)
|
||||
return token
|
||||
|
||||
def refresh() -> None:
|
||||
try:
|
||||
get_xai_oauth_token()
|
||||
except Exception as exc:
|
||||
refresh_errors.append(exc)
|
||||
|
||||
monkeypatch.setattr(xai_oauth, "_load_token", gated_load_token)
|
||||
refresh_thread = threading.Thread(target=refresh)
|
||||
refresh_thread.start()
|
||||
assert initial_read_finished.wait(timeout=2)
|
||||
|
||||
assert logout_xai_oauth() is True
|
||||
continue_refresh.set()
|
||||
refresh_thread.join(timeout=2)
|
||||
|
||||
assert not refresh_thread.is_alive()
|
||||
assert len(refresh_errors) == 1
|
||||
assert isinstance(refresh_errors[0], XAIOAuthError)
|
||||
assert "not signed in" in str(refresh_errors[0])
|
||||
assert not get_xai_oauth_storage_path().exists()
|
||||
Reference in New Issue
Block a user