feat(providers): add xAI Grok OAuth with capability-gated X Search (#5035)

2026-07-23 11:55:16 +08:00
parent c22efb5f7a
commit c7393c785e
38 changed files with 3881 additions and 104 deletions
@@ -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()