From 5922c4ebea06d3f8cd609cd70f34becf3000e919 Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Wed, 20 May 2026 23:07:58 +0800 Subject: [PATCH] feat(providers): add xAI Grok OAuth support --- nanobot/cli/commands.py | 166 ++++- nanobot/config/schema.py | 26 + nanobot/providers/__init__.py | 28 +- nanobot/providers/factory.py | 4 + nanobot/providers/registry.py | 14 +- nanobot/providers/xai_oauth_provider.py | 768 +++++++++++++++++++++ pyproject.toml | 1 + tests/cli/test_commands.py | 200 +++++- tests/providers/test_providers_init.py | 9 + tests/providers/test_xai_oauth_auth.py | 146 ++++ tests/providers/test_xai_oauth_provider.py | 141 ++++ 11 files changed, 1494 insertions(+), 9 deletions(-) create mode 100644 nanobot/providers/xai_oauth_provider.py create mode 100644 tests/providers/test_xai_oauth_auth.py create mode 100644 tests/providers/test_xai_oauth_provider.py diff --git a/nanobot/cli/commands.py b/nanobot/cli/commands.py index f7bf043a4..a291d1e17 100644 --- a/nanobot/cli/commands.py +++ b/nanobot/cli/commands.py @@ -1,12 +1,14 @@ """CLI commands for nanobot.""" import asyncio +import json import os import select import signal import sys from collections.abc import Callable from contextlib import nullcontext, suppress +from inspect import signature from pathlib import Path from typing import Any @@ -1527,6 +1529,106 @@ def status(): console.print(f"{spec.label}: {'[green]✓[/green]' if has_key else '[dim]not set[/dim]'}") +# ============================================================================ +# Config Commands +# ============================================================================ + +config_app = typer.Typer(help="Manage configuration") +app.add_typer(config_app, name="config") + + +@config_app.command("set") +def config_set( + path: str = typer.Argument(..., help="Dot path, e.g. agents.defaults.model"), + value: str = typer.Argument(..., help="Value. Use null/true/false or JSON for structured values."), + config_path: str | None = typer.Option(None, "--config", "-c", help="Path to config file"), +): + """Set one config value by dot path.""" + from pydantic import ValidationError + + from nanobot.config.loader import get_config_path, load_config, save_config, set_config_path + from nanobot.config.schema import Config + + resolved_path = Path(config_path).expanduser().resolve() if config_path else get_config_path() + if config_path: + set_config_path(resolved_path) + + config = load_config(resolved_path) + parsed = _parse_config_cli_value(value) + try: + _set_config_cli_value(config, path, parsed) + validated = Config.model_validate(config.model_dump(mode="json", by_alias=True)) + except (AttributeError, KeyError, TypeError, ValueError, ValidationError) as exc: + console.print(f"[red]Could not set config value:[/red] {exc}") + raise typer.Exit(1) + + save_config(validated, resolved_path) + console.print(f"[green]✓[/green] Set [cyan]{path}[/cyan] = [bold]{value}[/bold]") + console.print(f"[dim]Config: {resolved_path}[/dim]") + if path in {"agents.defaults.provider", "agents.defaults.model"} and validated.agents.defaults.model_preset: + console.print( + "[yellow]! agents.defaults.model_preset is set and may override this. " + "Clear it with: nanobot config set agents.defaults.model_preset null[/yellow]" + ) + + +def _parse_config_cli_value(raw: str) -> Any: + lowered = raw.strip().lower() + if lowered == "null": + return None + if lowered == "true": + return True + if lowered == "false": + return False + with suppress(Exception): + return json.loads(raw) + return raw + + +def _resolve_config_field(obj: Any, key: str) -> str: + from pydantic import BaseModel + from pydantic.alias_generators import to_camel, to_snake + + if not isinstance(obj, BaseModel): + return key + fields = type(obj).model_fields + if key in fields: + return key + normalized = to_snake(key.replace("-", "_")) + if normalized in fields: + return normalized + for name, field in fields.items(): + aliases = { + to_camel(name), + str(field.alias) if field.alias else "", + str(field.serialization_alias) if field.serialization_alias else "", + } + if key in aliases: + return name + raise AttributeError(f"Unknown config path segment {key!r}") + + +def _set_config_cli_value(config: Any, path: str, value: Any) -> None: + parts = [part for part in path.split(".") if part] + if not parts: + raise ValueError("Config path cannot be empty.") + + current = config + for raw_part in parts[:-1]: + if isinstance(current, dict): + current = current.setdefault(raw_part, {}) + continue + part = _resolve_config_field(current, raw_part) + current = getattr(current, part) + + leaf = parts[-1] + if isinstance(current, dict): + current[leaf] = value + return + leaf = _resolve_config_field(current, leaf) + setattr(current, leaf, value) + + # ============================================================================ # OAuth Login # ============================================================================ @@ -1541,6 +1643,7 @@ _LOGOUT_HANDLERS: dict[str, Callable[[], None]] = {} _PROVIDER_DISPLAY: dict[str, str] = { "openai_codex": "OpenAI Codex", "github_copilot": "GitHub Copilot", + "xai_oauth": "xAI Grok OAuth", } @@ -1576,7 +1679,9 @@ def _resolve_oauth_provider(provider: str): @provider_app.command("login") def provider_login( - provider: str = typer.Argument(..., help="OAuth provider (e.g. 'openai-codex', 'github-copilot')"), + provider: str = typer.Argument(..., help="OAuth provider (e.g. 'openai-codex', 'github-copilot', 'xai-oauth')"), + no_browser: bool = typer.Option(False, "--no-browser", help="Print the auth URL instead of opening a browser when supported."), + manual_paste: bool = typer.Option(False, "--manual-paste", help="Prompt for a callback URL or fallback code when supported."), ): """Authenticate with an OAuth provider.""" spec = _resolve_oauth_provider(provider) @@ -1587,12 +1692,18 @@ def provider_login( raise typer.Exit(1) console.print(f"{__logo__} OAuth Login - {spec.label}\n") - handler() + params = signature(handler).parameters + kwargs: dict[str, bool] = {} + if "no_browser" in params: + kwargs["no_browser"] = no_browser + if "manual_paste" in params: + kwargs["manual_paste"] = manual_paste + handler(**kwargs) @provider_app.command("logout") def provider_logout( - provider: str = typer.Argument(..., help="OAuth provider (e.g. 'openai-codex', 'github-copilot')"), + provider: str = typer.Argument(..., help="OAuth provider (e.g. 'openai-codex', 'github-copilot', 'xai-oauth')"), ): """Log out from an OAuth provider.""" spec = _resolve_oauth_provider(provider) @@ -1656,6 +1767,24 @@ def _logout_github_copilot() -> None: _delete_oauth_files(storage.get_token_path(), _PROVIDER_DISPLAY["github_copilot"]) +@_register_logout("xai_oauth") +def _logout_xai_oauth() -> None: + """Clear local OAuth credentials for xAI Grok OAuth.""" + try: + from nanobot.providers.xai_oauth_provider import delete_xai_oauth_credentials + except ImportError: + console.print("[red]xAI Grok OAuth provider unavailable.[/red]") + raise typer.Exit(1) + + removed_paths = delete_xai_oauth_credentials() + if not removed_paths: + console.print(f"[yellow]! No local OAuth credentials found for {_PROVIDER_DISPLAY['xai_oauth']}[/yellow]") + return + console.print(f"[green]✓ Logged out from {_PROVIDER_DISPLAY['xai_oauth']}[/green]") + for path in removed_paths: + console.print(f"[dim]Removed: {path}[/dim]") + + def _delete_oauth_files(token_path: Path, provider_label: str) -> None: """Delete OAuth token and lock files, reporting the result.""" removed_paths: list[Path] = [] @@ -1699,5 +1828,36 @@ def _login_github_copilot() -> None: raise typer.Exit(1) +@_register_login("xai_oauth") +def _login_xai_oauth( + *, + no_browser: bool = False, + manual_paste: bool = False, +) -> None: + try: + from nanobot.providers.xai_oauth_provider import login_xai_oauth_interactive + from nanobot.providers.xai_oauth_provider import DEFAULT_XAI_MODEL + + console.print("[cyan]Starting xAI Grok OAuth login...[/cyan]\n") + credential = login_xai_oauth_interactive( + print_fn=lambda s: console.print(s), + prompt_fn=lambda s: typer.prompt(s), + open_browser=not no_browser, + manual_paste=manual_paste, + ) + account = credential.account_id or "xAI" + storage = "OS keychain" if credential.storage == "keyring" else "private file" + console.print(f"[green]✓ Authenticated with xAI Grok OAuth[/green] [dim]{account} · {storage}[/dim]") + console.print("[dim]To use it for chat:[/dim]") + console.print("[dim] nanobot config set agents.defaults.model_preset null[/dim]") + console.print("[dim] nanobot config set agents.defaults.provider xai-oauth[/dim]") + console.print(f"[dim] nanobot config set agents.defaults.model {DEFAULT_XAI_MODEL}[/dim]") + console.print("[dim]Hosted X Search is enabled by default for xAI OAuth.[/dim]") + console.print("[dim]To disable it: nanobot config set providers.xai_oauth.x_search.enable false[/dim]") + except Exception as e: + console.print(f"[red]Authentication error: {e}[/red]") + raise typer.Exit(1) + + if __name__ == "__main__": app() diff --git a/nanobot/config/schema.py b/nanobot/config/schema.py index c0ad7e758..89b6a8d9d 100644 --- a/nanobot/config/schema.py +++ b/nanobot/config/schema.py @@ -180,6 +180,28 @@ class BedrockProviderConfig(ProviderConfig): profile: str | None = None # Optional AWS shared config profile +class XaiOAuthXSearchConfig(Base): + """xAI hosted X Search configuration.""" + + enable: bool = True + allowed_x_handles: list[str] | None = None + excluded_x_handles: list[str] | None = None + from_date: str | None = None + to_date: str | None = None + enable_image_understanding: bool = False + enable_video_understanding: bool = False + + +class XaiOAuthProviderConfig(ProviderConfig): + """xAI OAuth provider configuration.""" + + x_search: XaiOAuthXSearchConfig = Field(default_factory=XaiOAuthXSearchConfig) + + +def _is_default_xai_oauth_config(value: Any) -> bool: + return isinstance(value, XaiOAuthProviderConfig) and value == XaiOAuthProviderConfig() + + class ProvidersConfig(Base): """Configuration for LLM providers.""" @@ -217,6 +239,10 @@ class ProvidersConfig(Base): byteplus_coding_plan: ProviderConfig = Field(default_factory=ProviderConfig) # BytePlus Coding Plan openai_codex: ProviderConfig = Field(default_factory=ProviderConfig, exclude=True) # OpenAI Codex (OAuth) github_copilot: ProviderConfig = Field(default_factory=ProviderConfig, exclude=True) # Github Copilot (OAuth) + xai_oauth: XaiOAuthProviderConfig = Field( + default_factory=XaiOAuthProviderConfig, + exclude_if=_is_default_xai_oauth_config, + ) # xAI Grok OAuth qianfan: ProviderConfig = Field(default_factory=ProviderConfig) # Qianfan (百度千帆) nvidia: ProviderConfig = Field(default_factory=ProviderConfig) # NVIDIA NIM (nvapi- keys) diff --git a/nanobot/providers/__init__.py b/nanobot/providers/__init__.py index 6150788a7..fcbf6e30b 100644 --- a/nanobot/providers/__init__.py +++ b/nanobot/providers/__init__.py @@ -14,6 +14,7 @@ __all__ = [ "OpenAICompatProvider", "OpenAICodexProvider", "GitHubCopilotProvider", + "XaiOAuthProvider", "AzureOpenAIProvider", "BedrockProvider", ] @@ -23,10 +24,23 @@ _LAZY_IMPORTS = { "OpenAICompatProvider": ".openai_compat_provider", "OpenAICodexProvider": ".openai_codex_provider", "GitHubCopilotProvider": ".github_copilot_provider", + "XaiOAuthProvider": ".xai_oauth_provider", "AzureOpenAIProvider": ".azure_openai_provider", "BedrockProvider": ".bedrock_provider", } +_LAZY_SUBMODULES = { + "anthropic_provider": ".anthropic_provider", + "openai_compat_provider": ".openai_compat_provider", + "openai_codex_provider": ".openai_codex_provider", + "github_copilot_provider": ".github_copilot_provider", + "xai_oauth_provider": ".xai_oauth_provider", + "azure_openai_provider": ".azure_openai_provider", + "bedrock_provider": ".bedrock_provider", + "factory": ".factory", + "registry": ".registry", +} + if TYPE_CHECKING: from nanobot.providers.anthropic_provider import AnthropicProvider from nanobot.providers.azure_openai_provider import AzureOpenAIProvider @@ -34,12 +48,18 @@ if TYPE_CHECKING: from nanobot.providers.github_copilot_provider import GitHubCopilotProvider from nanobot.providers.openai_compat_provider import OpenAICompatProvider from nanobot.providers.openai_codex_provider import OpenAICodexProvider + from nanobot.providers.xai_oauth_provider import XaiOAuthProvider def __getattr__(name: str): """Lazily expose provider implementations without importing all backends up front.""" module_name = _LAZY_IMPORTS.get(name) - if module_name is None: - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - module = import_module(module_name, __name__) - return getattr(module, name) + if module_name is not None: + module = import_module(module_name, __name__) + return getattr(module, name) + module_name = _LAZY_SUBMODULES.get(name) + if module_name is not None: + module = import_module(module_name, __name__) + globals()[name] = module + return module + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/nanobot/providers/factory.py b/nanobot/providers/factory.py index 288611392..213ef15df 100644 --- a/nanobot/providers/factory.py +++ b/nanobot/providers/factory.py @@ -68,6 +68,10 @@ def _make_provider_core( from nanobot.providers.github_copilot_provider import GitHubCopilotProvider provider = GitHubCopilotProvider(default_model=model) + elif backend == "xai_oauth": + from nanobot.providers.xai_oauth_provider import XaiOAuthProvider + + provider = XaiOAuthProvider(default_model=model, config=p) elif backend == "anthropic": from nanobot.providers.anthropic_provider import AnthropicProvider diff --git a/nanobot/providers/registry.py b/nanobot/providers/registry.py index 7c8edd271..39cd751be 100644 --- a/nanobot/providers/registry.py +++ b/nanobot/providers/registry.py @@ -34,7 +34,7 @@ class ProviderSpec: display_name: str = "" # shown in `nanobot status` # which provider implementation to use - # "openai_compat" | "anthropic" | "azure_openai" | "openai_codex" | "github_copilot" | "bedrock" + # "openai_compat" | "anthropic" | "azure_openai" | "openai_codex" | "github_copilot" | "xai_oauth" | "bedrock" backend: str = "openai_compat" # extra env vars, e.g. (("ZHIPUAI_API_KEY", "{api_key}"),) @@ -291,6 +291,18 @@ PROVIDERS: tuple[ProviderSpec, ...] = ( is_oauth=True, supports_max_completion_tokens=True, ), + # xAI Grok OAuth: SuperGrok subscription-backed Responses API provider + ProviderSpec( + name="xai_oauth", + keywords=("xai-oauth", "grok-oauth", "x-ai-oauth", "xai-grok-oauth"), + env_key="", + display_name="xAI Grok OAuth", + backend="xai_oauth", + default_api_base="https://api.x.ai/v1", + strip_model_prefix=True, + is_oauth=True, + supports_max_completion_tokens=True, + ), # DeepSeek: OpenAI-compatible at api.deepseek.com ProviderSpec( name="deepseek", diff --git a/nanobot/providers/xai_oauth_provider.py b/nanobot/providers/xai_oauth_provider.py new file mode 100644 index 000000000..1b0943802 --- /dev/null +++ b/nanobot/providers/xai_oauth_provider.py @@ -0,0 +1,768 @@ +"""xAI Grok OAuth credential flow and Responses provider.""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import os +import secrets +import time +import webbrowser +from collections.abc import Awaitable, Callable +from contextlib import suppress +from dataclasses import dataclass +from hashlib import sha256 +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from threading import Event, Thread +from typing import Any +from urllib.parse import parse_qs, urlencode, urlparse + +import httpx +from filelock import FileLock + +from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest +from nanobot.providers.openai_responses import consume_sse, convert_messages, convert_tools + +DEFAULT_XAI_API_BASE = "https://api.x.ai/v1" +DEFAULT_XAI_AUTH_ISSUER = "https://auth.x.ai" +DEFAULT_XAI_DISCOVERY_URL = f"{DEFAULT_XAI_AUTH_ISSUER}/.well-known/openid-configuration" +DEFAULT_XAI_REDIRECT_URI = "http://127.0.0.1:56121/callback" +DEFAULT_XAI_CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828" +DEFAULT_XAI_SCOPE = "openid profile email offline_access grok-cli:access api:access" + +_SERVICE_NAME = "nanobot.xai_oauth" +_SECRET_USERNAME = "default" +_TOKEN_SKEW_SECONDS = 60 +_LOGIN_TIMEOUT_SECONDS = 300 + + +@dataclass(frozen=True) +class XaiOAuthEndpoints: + authorization_endpoint: str + token_endpoint: str + + +@dataclass(frozen=True) +class XaiOAuthCredential: + access_token: str + refresh_token: str = "" + expires_at: float | None = None + account_id: str | None = None + token_type: str = "Bearer" + api_base: str = DEFAULT_XAI_API_BASE + storage: str = "unknown" + + @property + def is_expiring(self) -> bool: + return self.expires_at is not None and self.expires_at <= time.time() + _TOKEN_SKEW_SECONDS + + +def _nanobot_home() -> Path: + override = os.environ.get("NANOBOT_HOME") + if override: + return Path(override).expanduser() + from nanobot.config.loader import get_config_path + + return get_config_path().parent + + +def _auth_dir() -> Path: + return _nanobot_home() / "auth" + + +def get_xai_oauth_metadata_path() -> Path: + """Return the non-secret xAI OAuth metadata path.""" + return _auth_dir() / "xai-oauth.json" + + +def _lock_path() -> Path: + return get_xai_oauth_metadata_path().with_suffix(".lock") + + +def _write_private_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with suppress(OSError): + path.parent.chmod(0o700) + tmp = path.with_suffix(path.suffix + ".tmp") + tmp.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8") + with suppress(OSError): + tmp.chmod(0o600) + tmp.replace(path) + with suppress(OSError): + path.chmod(0o600) + + +def _read_json(path: Path) -> dict[str, Any]: + return json.loads(path.read_text(encoding="utf-8")) + + +def _keyring_set(tokens: dict[str, Any]) -> bool: + try: + import keyring # type: ignore[import-not-found] + + keyring.set_password(_SERVICE_NAME, _SECRET_USERNAME, json.dumps(tokens)) + return True + except Exception: + return False + + +def _keyring_get() -> dict[str, Any] | None: + try: + import keyring # type: ignore[import-not-found] + + raw = keyring.get_password(_SERVICE_NAME, _SECRET_USERNAME) + except Exception: + return None + if not raw: + return None + try: + payload = json.loads(raw) + except json.JSONDecodeError: + return None + return payload if isinstance(payload, dict) else None + + +def _keyring_delete() -> None: + try: + import keyring # type: ignore[import-not-found] + + keyring.delete_password(_SERVICE_NAME, _SECRET_USERNAME) + except Exception: + pass + + +def _token_payload(credential: XaiOAuthCredential) -> dict[str, Any]: + return { + "access_token": credential.access_token, + "refresh_token": credential.refresh_token, + "expires_at": credential.expires_at, + "token_type": credential.token_type, + } + + +def save_xai_oauth_credential(credential: XaiOAuthCredential) -> XaiOAuthCredential: + """Persist xAI OAuth tokens, preferring OS keychain storage.""" + with FileLock(str(_lock_path())): + tokens = _token_payload(credential) + metadata: dict[str, Any] = { + "provider": "xai_oauth", + "api_base": credential.api_base, + "account_id": credential.account_id, + "expires_at": credential.expires_at, + "updated_at": int(time.time()), + } + if _keyring_set(tokens): + metadata["storage"] = "keyring" + else: + metadata["storage"] = "file" + metadata["tokens"] = tokens + _write_private_json(get_xai_oauth_metadata_path(), metadata) + return XaiOAuthCredential( + access_token=credential.access_token, + refresh_token=credential.refresh_token, + expires_at=credential.expires_at, + account_id=credential.account_id, + token_type=credential.token_type, + api_base=credential.api_base, + storage=str(metadata["storage"]), + ) + + +def load_xai_oauth_credential() -> XaiOAuthCredential | None: + """Load xAI OAuth credentials from keyring or the private file fallback.""" + path = get_xai_oauth_metadata_path() + if not path.exists(): + return None + with FileLock(str(_lock_path())): + try: + metadata = _read_json(path) + except (OSError, json.JSONDecodeError): + return None + + storage = str(metadata.get("storage") or "file") + tokens = _keyring_get() if storage == "keyring" else metadata.get("tokens") + if not isinstance(tokens, dict): + return None + access_token = str(tokens.get("access_token") or "") + if not access_token: + return None + + return XaiOAuthCredential( + access_token=access_token, + refresh_token=str(tokens.get("refresh_token") or ""), + expires_at=_as_float(tokens.get("expires_at") or metadata.get("expires_at")), + account_id=_as_str(metadata.get("account_id")), + token_type=str(tokens.get("token_type") or "Bearer"), + api_base=str(metadata.get("api_base") or DEFAULT_XAI_API_BASE), + storage=storage, + ) + + +def delete_xai_oauth_credentials() -> list[Path]: + """Delete persisted xAI OAuth credentials and return removed local paths.""" + removed: list[Path] = [] + path = get_xai_oauth_metadata_path() + lock_path = _lock_path() + with FileLock(str(lock_path)): + _keyring_delete() + try: + path.unlink() + removed.append(path) + except FileNotFoundError: + pass + try: + lock_path.unlink() + except FileNotFoundError: + pass + return removed + + +def get_xai_oauth_login_status() -> XaiOAuthCredential | None: + return load_xai_oauth_credential() + + +def pkce_challenge(verifier: str) -> str: + digest = sha256(verifier.encode("ascii")).digest() + return base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=") + + +def _new_pkce_verifier() -> str: + return base64.urlsafe_b64encode(secrets.token_bytes(48)).decode("ascii").rstrip("=") + + +def build_xai_authorization_url( + endpoints: XaiOAuthEndpoints, + *, + verifier: str, + state: str, + nonce: str | None = None, + redirect_uri: str = DEFAULT_XAI_REDIRECT_URI, +) -> str: + params = { + "response_type": "code", + "client_id": DEFAULT_XAI_CLIENT_ID, + "redirect_uri": redirect_uri, + "scope": DEFAULT_XAI_SCOPE, + "code_challenge": pkce_challenge(verifier), + "code_challenge_method": "S256", + "state": state, + "nonce": nonce or secrets.token_urlsafe(16), + "plan": "generic", + "referrer": "nanobot", + } + return f"{endpoints.authorization_endpoint}?{urlencode(params)}" + + +def discover_xai_oauth_endpoints() -> XaiOAuthEndpoints: + try: + with httpx.Client(timeout=20.0, follow_redirects=True, trust_env=True) as client: + response = client.get(DEFAULT_XAI_DISCOVERY_URL) + response.raise_for_status() + payload = response.json() + except Exception: + payload = {} + + endpoints = XaiOAuthEndpoints( + authorization_endpoint=str( + payload.get("authorization_endpoint") + or f"{DEFAULT_XAI_AUTH_ISSUER}/authorize" + ), + token_endpoint=str( + payload.get("token_endpoint") + or f"{DEFAULT_XAI_AUTH_ISSUER}/oauth/token" + ), + ) + _validate_xai_endpoint(endpoints.authorization_endpoint, "authorization_endpoint") + _validate_xai_endpoint(endpoints.token_endpoint, "token_endpoint") + return endpoints + + +def _validate_xai_endpoint(url: str, label: str) -> None: + parsed = urlparse(url) + host = parsed.hostname or "" + if parsed.scheme != "https" or not (host == "x.ai" or host.endswith(".x.ai")): + raise RuntimeError(f"Refusing non-xAI OAuth {label}: {url}") + + +def _parse_callback_value(raw: str) -> tuple[str, str | None]: + raw = raw.strip() + parsed = urlparse(raw) + if parsed.scheme and parsed.netloc: + params = parse_qs(parsed.query) + code = (params.get("code") or [""])[0] + state = (params.get("state") or [None])[0] + if not code: + raise RuntimeError("OAuth callback URL did not contain a code.") + return code, state + if raw.startswith("?") or "=" in raw: + params = parse_qs(raw.lstrip("?")) + code = (params.get("code") or [""])[0] + state = (params.get("state") or [None])[0] + if not code: + raise RuntimeError("OAuth callback query did not contain a code.") + return code, state + if raw: + return raw, None + raise RuntimeError("No OAuth code provided.") + + +def _decode_jwt_payload(token: str) -> dict[str, Any]: + parts = token.split(".") + if len(parts) < 2: + return {} + data = parts[1] + "=" * (-len(parts[1]) % 4) + try: + decoded = base64.urlsafe_b64decode(data.encode("ascii")) + payload = json.loads(decoded) + except Exception: + return {} + return payload if isinstance(payload, dict) else {} + + +def _credential_from_token_response(payload: dict[str, Any], previous: XaiOAuthCredential | None = None) -> XaiOAuthCredential: + access_token = str(payload.get("access_token") or "") + if not access_token: + raise RuntimeError("xAI token response did not include an access token.") + + claims = _decode_jwt_payload(access_token) + id_claims = _decode_jwt_payload(str(payload.get("id_token") or "")) + expires_at = _as_float(payload.get("expires_at")) + if expires_at is None: + expires_in = _as_float(payload.get("expires_in")) + expires_at = time.time() + expires_in if expires_in else _as_float(claims.get("exp")) + + account_id = ( + _as_str(id_claims.get("email")) + or _as_str(id_claims.get("preferred_username")) + or _as_str(id_claims.get("sub")) + or _as_str(claims.get("sub")) + or (previous.account_id if previous else None) + ) + refresh_token = str(payload.get("refresh_token") or (previous.refresh_token if previous else "")) + + return XaiOAuthCredential( + access_token=access_token, + refresh_token=refresh_token, + expires_at=expires_at, + account_id=account_id, + token_type=str(payload.get("token_type") or (previous.token_type if previous else "Bearer")), + api_base=previous.api_base if previous else DEFAULT_XAI_API_BASE, + ) + + +def exchange_xai_oauth_code( + code: str, + *, + verifier: str, + endpoints: XaiOAuthEndpoints | None = None, + redirect_uri: str = DEFAULT_XAI_REDIRECT_URI, +) -> XaiOAuthCredential: + endpoints = endpoints or discover_xai_oauth_endpoints() + challenge = pkce_challenge(verifier) + with httpx.Client(timeout=30.0, follow_redirects=True, trust_env=True) as client: + response = client.post( + endpoints.token_endpoint, + headers={"Accept": "application/json"}, + data={ + "grant_type": "authorization_code", + "client_id": DEFAULT_XAI_CLIENT_ID, + "code": code, + "redirect_uri": redirect_uri, + "code_verifier": verifier, + "code_challenge": challenge, + "code_challenge_method": "S256", + }, + ) + if response.status_code >= 400: + raise RuntimeError(f"xAI token exchange failed: HTTP {response.status_code}: {response.text[:500]}") + return _credential_from_token_response(response.json()) + + +def refresh_xai_oauth_credential(credential: XaiOAuthCredential | None = None) -> XaiOAuthCredential: + credential = credential or load_xai_oauth_credential() + if not credential or not credential.refresh_token: + raise RuntimeError("xAI Grok OAuth is not logged in. Run: nanobot provider login xai-oauth") + + endpoints = discover_xai_oauth_endpoints() + with httpx.Client(timeout=30.0, follow_redirects=True, trust_env=True) as client: + response = client.post( + endpoints.token_endpoint, + headers={"Accept": "application/json"}, + data={ + "grant_type": "refresh_token", + "client_id": DEFAULT_XAI_CLIENT_ID, + "refresh_token": credential.refresh_token, + }, + ) + if response.status_code >= 400: + raise RuntimeError(f"xAI token refresh failed: HTTP {response.status_code}: {response.text[:500]}") + return save_xai_oauth_credential(_credential_from_token_response(response.json(), previous=credential)) + + +def resolve_xai_oauth_credential(*, force_refresh: bool = False) -> XaiOAuthCredential: + credential = load_xai_oauth_credential() + if not credential: + raise RuntimeError("xAI Grok OAuth is not logged in. Run: nanobot provider login xai-oauth") + if force_refresh or credential.is_expiring: + credential = refresh_xai_oauth_credential(credential) + return credential + + +def login_xai_oauth_interactive( + print_fn: Callable[[str], None] | None = None, + prompt_fn: Callable[[str], str] | None = None, + open_browser: bool = True, + manual_paste: bool = False, + timeout_seconds: int = _LOGIN_TIMEOUT_SECONDS, +) -> XaiOAuthCredential: + """Run browser PKCE login and persist xAI OAuth credentials.""" + printer = print_fn or print + prompt = prompt_fn or input + endpoints = discover_xai_oauth_endpoints() + verifier = _new_pkce_verifier() + state = secrets.token_urlsafe(24) + nonce = secrets.token_urlsafe(24) + authorize_url = build_xai_authorization_url( + endpoints, + verifier=verifier, + state=state, + nonce=nonce, + ) + + callback = _LoopbackCallback() + server_started = False if manual_paste else callback.start() + printer(f"Open: {authorize_url}") + if open_browser: + with suppress(Exception): + webbrowser.open(authorize_url) + + result: dict[str, str] | None = None + if manual_paste: + printer("Paste the callback URL or xAI fallback code after authorization.") + elif server_started: + try: + result = callback.wait(timeout_seconds) + finally: + callback.stop() + else: + printer("Loopback port 56121 is unavailable; paste the callback URL or xAI fallback code.") + + if result: + code = result.get("code") or "" + returned_state = result.get("state") + else: + pasted = prompt("Paste callback URL or fallback code") + code, returned_state = _parse_callback_value(pasted) + + if not code: + raise RuntimeError("OAuth login did not return a code.") + if returned_state and returned_state != state: + raise RuntimeError("OAuth state mismatch. Please retry login.") + + credential = exchange_xai_oauth_code(code, verifier=verifier, endpoints=endpoints) + return save_xai_oauth_credential(credential) + + +class _LoopbackCallback: + def __init__(self) -> None: + self._event = Event() + self._result: dict[str, str] = {} + self._server: ThreadingHTTPServer | None = None + self._thread: Thread | None = None + + def start(self) -> bool: + owner = self + + class Handler(BaseHTTPRequestHandler): + def do_GET(self) -> None: # noqa: N802 - stdlib callback name + parsed = urlparse(self.path) + params = parse_qs(parsed.query) + code = (params.get("code") or [""])[0] + state = (params.get("state") or [""])[0] + if parsed.path != "/callback" or not code: + self.send_response(404) + self.end_headers() + return + owner._result = {"code": code, "state": state} + owner._event.set() + self.send_response(200) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.end_headers() + self.wfile.write(b"nanobot xAI OAuth complete. You may close this tab.") + + def log_message(self, format: str, *args: Any) -> None: # noqa: A002 + return + + class Server(ThreadingHTTPServer): + allow_reuse_address = True + daemon_threads = True + + try: + self._server = Server(("127.0.0.1", 56121), Handler) + except OSError: + return False + self._thread = Thread(target=self._server.serve_forever, daemon=True) + self._thread.start() + return True + + def wait(self, timeout_seconds: int) -> dict[str, str] | None: + if self._event.wait(timeout_seconds): + return dict(self._result) + return None + + def stop(self) -> None: + if self._server: + self._server.shutdown() + self._server.server_close() + if self._thread: + self._thread.join(timeout=1) + + +def _as_float(value: Any) -> float | None: + try: + return float(value) + except (TypeError, ValueError): + return None + + +def _as_str(value: Any) -> str | None: + return value if isinstance(value, str) and value else None + + +DEFAULT_XAI_MODEL = "xai-oauth/grok-4.3" + + +class XaiOAuthProvider(LLMProvider): + """Use a SuperGrok OAuth session to call xAI's Responses API.""" + + supports_progress_deltas = True + + def __init__(self, default_model: str = DEFAULT_XAI_MODEL, config: Any | None = None): + super().__init__(api_key=None, api_base=DEFAULT_XAI_API_BASE) + self.default_model = default_model + self.config = config + + async def _call_xai( + self, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None, + model: str | None, + max_tokens: int, + temperature: float, + reasoning_effort: str | None, + tool_choice: str | dict[str, Any] | None, + on_content_delta: Callable[[str], Awaitable[None]] | None = None, + on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None, + ) -> LLMResponse: + body = _build_xai_responses_body( + messages=messages, + tools=tools, + model=model or self.default_model, + max_tokens=max_tokens, + temperature=temperature, + reasoning_effort=reasoning_effort, + tool_choice=tool_choice, + hosted_x_search=getattr(self.config, "x_search", None), + ) + try: + credential = await asyncio.to_thread(resolve_xai_oauth_credential) + try: + content, tool_calls, finish_reason = await _request_xai( + credential, + body, + on_content_delta=on_content_delta, + on_tool_call_delta=on_tool_call_delta, + ) + except _XaiHTTPError as exc: + if exc.status_code != 401: + raise + credential = await asyncio.to_thread(resolve_xai_oauth_credential, force_refresh=True) + content, tool_calls, finish_reason = await _request_xai( + credential, + body, + on_content_delta=on_content_delta, + on_tool_call_delta=on_tool_call_delta, + ) + return LLMResponse(content=content, tool_calls=tool_calls, finish_reason=finish_reason) + except Exception as exc: + msg = f"Error calling xAI Grok OAuth: {exc}" + retry_after = getattr(exc, "retry_after", None) or self._extract_retry_after(msg) + return LLMResponse(content=msg, finish_reason="error", retry_after=retry_after) + + async def chat( + self, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None = None, + model: str | None = None, + max_tokens: int = 4096, + temperature: float = 0.7, + reasoning_effort: str | None = None, + tool_choice: str | dict[str, Any] | None = None, + ) -> LLMResponse: + return await self._call_xai( + messages, + tools, + model, + max_tokens, + temperature, + reasoning_effort, + tool_choice, + ) + + async def chat_stream( + self, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None = None, + model: str | None = None, + max_tokens: int = 4096, + temperature: float = 0.7, + reasoning_effort: str | None = None, + tool_choice: str | dict[str, Any] | None = None, + 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, + ) -> LLMResponse: + _ = on_thinking_delta + return await self._call_xai( + messages, + tools, + model, + max_tokens, + temperature, + reasoning_effort, + tool_choice, + on_content_delta, + on_tool_call_delta, + ) + + def get_default_model(self) -> str: + return self.default_model + + +def _strip_model_prefix(model: str) -> str: + for prefix in ("xai-oauth/", "xai_oauth/", "grok-oauth/", "grok_oauth/"): + if model.startswith(prefix): + return model.split("/", 1)[1] + return model + + +def _build_xai_responses_body( + *, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None, + model: str, + max_tokens: int, + temperature: float, + reasoning_effort: str | None, + tool_choice: str | dict[str, Any] | None, + hosted_x_search: Any | None = None, +) -> dict[str, Any]: + system_prompt, input_items = convert_messages(LLMProvider._sanitize_empty_content(messages)) + if system_prompt: + input_items = [ + {"role": "system", "content": [{"type": "input_text", "text": system_prompt}]}, + *input_items, + ] + + body: dict[str, Any] = { + "model": _strip_model_prefix(model), + "store": False, + "stream": True, + "input": input_items, + "tool_choice": tool_choice or "auto", + "parallel_tool_calls": True, + } + if max_tokens: + body["max_output_tokens"] = max_tokens + if temperature is not None: + body["temperature"] = temperature + if reasoning_effort and reasoning_effort.lower() != "none": + body["reasoning"] = {"effort": reasoning_effort} + converted_tools = convert_tools(tools) if tools else [] + hosted_tool = _build_xai_hosted_x_search_tool(hosted_x_search) + if hosted_tool: + converted_tools.append(hosted_tool) + if converted_tools: + body["tools"] = converted_tools + return body + + +def _clean_x_handles(handles: list[str] | None) -> list[str] | None: + if not handles: + return None + cleaned = [str(handle).strip().lstrip("@") for handle in handles if str(handle).strip()] + return cleaned[:10] or None + + +def _build_xai_hosted_x_search_tool(config: Any | None) -> dict[str, Any] | None: + if not config or not getattr(config, "enable", False): + return None + + allowed = _clean_x_handles(getattr(config, "allowed_x_handles", None)) + excluded = _clean_x_handles(getattr(config, "excluded_x_handles", None)) + if allowed and excluded: + raise ValueError("providers.xai_oauth.x_search cannot set both allowed_x_handles and excluded_x_handles") + + tool: dict[str, Any] = {"type": "x_search"} + if allowed: + tool["allowed_x_handles"] = allowed + if excluded: + tool["excluded_x_handles"] = excluded + if getattr(config, "from_date", None): + tool["from_date"] = config.from_date + if getattr(config, "to_date", None): + tool["to_date"] = config.to_date + if getattr(config, "enable_image_understanding", False): + tool["enable_image_understanding"] = True + if getattr(config, "enable_video_understanding", False): + tool["enable_video_understanding"] = True + return tool + + +class _XaiHTTPError(RuntimeError): + def __init__(self, message: str, *, status_code: int, retry_after: float | None = None): + super().__init__(message) + self.status_code = status_code + self.retry_after = retry_after + + +async def _request_xai( + credential: XaiOAuthCredential, + body: dict[str, Any], + on_content_delta: Callable[[str], Awaitable[None]] | None = None, + on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None, +) -> tuple[str, list[ToolCallRequest], str]: + url = credential.api_base.rstrip("/") + "/responses" + headers = { + "Authorization": f"Bearer {credential.access_token}", + "Accept": "text/event-stream", + "Content-Type": "application/json", + "User-Agent": "nanobot (python)", + } + timeout = httpx.Timeout(120.0, connect=20.0) + async with httpx.AsyncClient(timeout=timeout, follow_redirects=True, trust_env=True) as client: + async with client.stream("POST", url, headers=headers, json=body) as response: + if response.status_code != 200: + raw = await response.aread() + retry_after = LLMProvider._extract_retry_after_from_headers(response.headers) + raise _XaiHTTPError( + _friendly_error(response.status_code, raw.decode("utf-8", "ignore")), + status_code=response.status_code, + retry_after=retry_after, + ) + return await consume_sse(response, on_content_delta, on_tool_call_delta) + + +def _friendly_error(status_code: int, raw: str) -> str: + if status_code == 401: + return "xAI OAuth session expired or was revoked. Run: nanobot provider login xai-oauth" + if status_code == 403: + return ( + "xAI accepted the OAuth token, but this account is not entitled for the requested " + "Grok API capability yet. Check the active Grok subscription and selected model." + ) + if status_code == 429: + return "xAI Grok subscription quota or rate limit was reached. Please try again later." + return f"HTTP {status_code}: {raw[:500]}" diff --git a/pyproject.toml b/pyproject.toml index eaf57a2ad..e8dd51e12 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,7 @@ dependencies = [ "openpyxl>=3.1.0,<4.0.0", "python-pptx>=1.0.0,<2.0.0", "filelock>=3.25.2", + "keyring>=25.0.0,<26.0.0", "boto3>=1.43.0", ] diff --git a/tests/cli/test_commands.py b/tests/cli/test_commands.py index 1dede1e13..646860e73 100644 --- a/tests/cli/test_commands.py +++ b/tests/cli/test_commands.py @@ -11,7 +11,7 @@ from typer.testing import CliRunner from nanobot.bus.events import OutboundMessage from nanobot.cli.commands import app from nanobot.providers.factory import make_provider -from nanobot.config.schema import Config +from nanobot.config.schema import Config, ModelPresetConfig from nanobot.cron.types import CronJob, CronPayload from nanobot.providers.factory import ProviderSnapshot from nanobot.providers.openai_codex_provider import _strip_model_prefix @@ -226,6 +226,16 @@ def test_config_dump_excludes_oauth_provider_blocks(): assert "openaiCodex" not in providers assert "githubCopilot" not in providers + assert "xaiOauth" not in providers + + +def test_config_dump_includes_xai_oauth_when_hosted_search_is_disabled(): + config = Config() + config.providers.xai_oauth.x_search.enable = False + + providers = config.model_dump(by_alias=True)["providers"] + + assert providers["xaiOauth"]["xSearch"]["enable"] is False def test_provider_logout_openai_codex_removes_local_oauth_files(tmp_path, monkeypatch): @@ -280,6 +290,175 @@ def test_provider_logout_github_copilot_succeeds_when_no_local_oauth_file(monkey assert "No local OAuth credentials found for GitHub Copilot" in result.stdout +def test_provider_logout_xai_oauth_removes_local_oauth_files(tmp_path, monkeypatch): + token_path = tmp_path / "auth" / "xai-oauth.json" + lock_path = token_path.with_suffix(".lock") + token_path.parent.mkdir(parents=True, exist_ok=True) + token_path.write_text("{}", encoding="utf-8") + lock_path.write_text("", encoding="utf-8") + monkeypatch.setenv("NANOBOT_HOME", str(tmp_path)) + monkeypatch.setattr("nanobot.providers.xai_oauth_provider._keyring_delete", lambda: None) + + result = runner.invoke(app, ["provider", "logout", "xai-oauth"]) + + assert result.exit_code == 0 + assert not token_path.exists() + assert not lock_path.exists() + assert "Logged out from xAI Grok OAuth" in result.stdout + + +def test_provider_logout_xai_oauth_succeeds_when_no_local_oauth_file(monkeypatch, tmp_path): + monkeypatch.setenv("NANOBOT_HOME", str(tmp_path)) + monkeypatch.setattr("nanobot.providers.xai_oauth_provider._keyring_delete", lambda: None) + + result = runner.invoke(app, ["provider", "logout", "xai-oauth"]) + + assert result.exit_code == 0 + assert "No local OAuth credentials found for xAI Grok OAuth" in result.stdout + + +def test_provider_login_xai_oauth_forwards_manual_options(monkeypatch): + from nanobot.providers.xai_oauth_provider import XaiOAuthCredential + + captured: dict[str, object] = {} + + def fake_login_xai_oauth_interactive(**kwargs): + captured.update(kwargs) + return XaiOAuthCredential(access_token="access", account_id="acct", storage="keyring") + + monkeypatch.setattr( + "nanobot.providers.xai_oauth_provider.login_xai_oauth_interactive", + fake_login_xai_oauth_interactive, + ) + + result = runner.invoke(app, ["provider", "login", "xai-oauth", "--no-browser", "--manual-paste"]) + + assert result.exit_code == 0 + assert captured["open_browser"] is False + assert captured["manual_paste"] is True + assert "Authenticated with xAI Grok OAuth" in result.stdout + assert "nanobot config set agents.defaults.provider xai-oauth" in result.stdout + + +def test_config_set_updates_default_model_selection(tmp_path): + config_path = tmp_path / "config.json" + + result = runner.invoke(app, [ + "config", + "set", + "--config", + str(config_path), + "agents.defaults.model_preset", + "null", + ]) + assert result.exit_code == 0 + + result = runner.invoke(app, [ + "config", + "set", + "--config", + str(config_path), + "agents.defaults.provider", + "xai-oauth", + ]) + assert result.exit_code == 0 + + result = runner.invoke(app, [ + "config", + "set", + "--config", + str(config_path), + "agents.defaults.model", + "xai-oauth/grok-4.3", + ]) + assert result.exit_code == 0 + + data = json.loads(config_path.read_text(encoding="utf-8")) + config = Config.model_validate(data) + assert config.agents.defaults.model_preset is None + assert config.agents.defaults.provider == "xai-oauth" + assert config.agents.defaults.model == "xai-oauth/grok-4.3" + + +def test_config_set_warns_when_model_preset_would_override_selection(tmp_path): + config = Config() + config.agents.defaults.model_preset = "fast" + config.model_presets["fast"] = ModelPresetConfig( + provider="openrouter", + model="openrouter/openai/gpt-4o-mini", + ) + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps(config.model_dump(mode="json", by_alias=True)), encoding="utf-8") + + result = runner.invoke(app, [ + "config", + "set", + "--config", + str(config_path), + "agents.defaults.provider", + "xai-oauth", + ]) + + assert result.exit_code == 0 + assert "model_preset is set and may override this" in result.stdout + + +def test_config_set_disables_xai_oauth_hosted_search(tmp_path): + config_path = tmp_path / "config.json" + + result = runner.invoke(app, [ + "config", + "set", + "--config", + str(config_path), + "providers.xai_oauth.x_search.enable", + "false", + ]) + + assert result.exit_code == 0 + data = json.loads(config_path.read_text(encoding="utf-8")) + assert data["providers"]["xaiOauth"]["xSearch"]["enable"] is False + assert Config.model_validate(data).providers.xai_oauth.x_search.enable is False + + +def test_config_set_rejects_unknown_path(tmp_path): + result = runner.invoke(app, [ + "config", + "set", + "--config", + str(tmp_path / "config.json"), + "agents.defaults.not_a_field", + "value", + ]) + + assert result.exit_code == 1 + assert "Could not set config value" in result.stdout + + +def test_provider_login_xai_oauth_does_not_update_config(monkeypatch, tmp_path): + from nanobot.providers.xai_oauth_provider import XaiOAuthCredential + + config = Config() + config.agents.defaults.provider = "auto" + config.agents.defaults.model = "anthropic/claude-opus-4-5" + config_path = tmp_path / "config.json" + + monkeypatch.setattr( + "nanobot.providers.xai_oauth_provider.login_xai_oauth_interactive", + lambda **_kwargs: XaiOAuthCredential(access_token="access", account_id="acct", storage="keyring"), + ) + monkeypatch.setattr("nanobot.config.loader.get_config_path", lambda: config_path) + monkeypatch.setattr("nanobot.config.loader.load_config", lambda _path=None: config) + save_config = MagicMock() + monkeypatch.setattr("nanobot.config.loader.save_config", save_config) + + result = runner.invoke(app, ["provider", "login", "xai-oauth"]) + + assert result.exit_code == 0 + save_config.assert_not_called() + assert "nanobot config set agents.defaults.model xai-oauth/grok-4.3" in result.stdout + + def test_provider_logout_rejects_unknown_provider(): result = runner.invoke(app, ["provider", "logout", "not-a-real-provider"]) @@ -398,6 +577,8 @@ def test_find_by_name_accepts_camel_case_and_hyphen_aliases(): assert find_by_name("volcengineCodingPlan").name == "volcengine_coding_plan" assert find_by_name("github-copilot") is not None assert find_by_name("github-copilot").name == "github_copilot" + assert find_by_name("xai-oauth") is not None + assert find_by_name("xai-oauth").name == "xai_oauth" assert find_by_name("longcat") is not None assert find_by_name("longcat").name == "longcat" assert find_by_name("atomic-chat") is not None @@ -540,6 +721,23 @@ def test_make_provider_uses_github_copilot_backend(): assert provider.__class__.__name__ == "GitHubCopilotProvider" +def test_make_provider_uses_xai_oauth_backend(): + config = Config.model_validate( + { + "agents": { + "defaults": { + "provider": "xai-oauth", + "model": "xai-oauth/grok-4.3", + } + } + } + ) + + provider = make_provider(config) + + assert provider.__class__.__name__ == "XaiOAuthProvider" + + def test_github_copilot_provider_strips_prefixed_model_name(): from nanobot.providers.github_copilot_provider import GitHubCopilotProvider diff --git a/tests/providers/test_providers_init.py b/tests/providers/test_providers_init.py index 707a3f683..9b26a1f76 100644 --- a/tests/providers/test_providers_init.py +++ b/tests/providers/test_providers_init.py @@ -12,6 +12,7 @@ def test_importing_providers_package_is_lazy(monkeypatch) -> None: 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.github_copilot_provider", raising=False) + monkeypatch.delitem(sys.modules, "nanobot.providers.xai_oauth_provider", raising=False) monkeypatch.delitem(sys.modules, "nanobot.providers.azure_openai_provider", raising=False) monkeypatch.delitem(sys.modules, "nanobot.providers.bedrock_provider", raising=False) @@ -21,6 +22,7 @@ def test_importing_providers_package_is_lazy(monkeypatch) -> None: 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.xai_oauth_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__ == [ @@ -30,6 +32,7 @@ def test_importing_providers_package_is_lazy(monkeypatch) -> None: "OpenAICompatProvider", "OpenAICodexProvider", "GitHubCopilotProvider", + "XaiOAuthProvider", "AzureOpenAIProvider", "BedrockProvider", ] @@ -50,3 +53,9 @@ def test_openai_codex_supports_progress_deltas() -> None: from nanobot.providers.openai_codex_provider import OpenAICodexProvider assert OpenAICodexProvider.supports_progress_deltas is True + + +def test_xai_oauth_supports_progress_deltas() -> None: + from nanobot.providers.xai_oauth_provider import XaiOAuthProvider + + assert XaiOAuthProvider.supports_progress_deltas is True diff --git a/tests/providers/test_xai_oauth_auth.py b/tests/providers/test_xai_oauth_auth.py new file mode 100644 index 000000000..3fd4cdd2e --- /dev/null +++ b/tests/providers/test_xai_oauth_auth.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +import json +import os +import stat +from urllib.parse import parse_qs, urlparse + +import pytest + +import nanobot.providers.xai_oauth_provider as auth + + +def test_build_xai_authorization_url_includes_pkce_and_grok_scope() -> None: + endpoints = auth.XaiOAuthEndpoints( + authorization_endpoint="https://auth.x.ai/authorize", + token_endpoint="https://auth.x.ai/oauth/token", + ) + + url = auth.build_xai_authorization_url( + endpoints, + verifier="verifier", + state="state", + nonce="nonce", + ) + + parsed = urlparse(url) + params = parse_qs(parsed.query) + assert parsed.scheme == "https" + assert parsed.hostname == "auth.x.ai" + assert params["client_id"] == [auth.DEFAULT_XAI_CLIENT_ID] + assert params["code_challenge"] == [auth.pkce_challenge("verifier")] + assert params["code_challenge_method"] == ["S256"] + assert params["scope"] == [auth.DEFAULT_XAI_SCOPE] + assert params["nonce"] == ["nonce"] + assert params["plan"] == ["generic"] + assert params["referrer"] == ["nanobot"] + + +def test_parse_callback_value_accepts_fallback_shapes() -> None: + assert auth._parse_callback_value("https://localhost/callback?code=abc&state=state") == ("abc", "state") + assert auth._parse_callback_value("?code=abc&state=state") == ("abc", "state") + assert auth._parse_callback_value("code=abc&state=state") == ("abc", "state") + assert auth._parse_callback_value("fallback-code") == ("fallback-code", None) + + +def test_file_storage_fallback_is_private_and_round_trips(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("NANOBOT_HOME", str(tmp_path)) + monkeypatch.setattr(auth, "_keyring_set", lambda _tokens: False) + monkeypatch.setattr(auth, "_keyring_get", lambda: None) + + saved = auth.save_xai_oauth_credential( + auth.XaiOAuthCredential( + access_token="access", + refresh_token="refresh", + expires_at=123.0, + account_id="acct", + ) + ) + + path = auth.get_xai_oauth_metadata_path() + payload = json.loads(path.read_text(encoding="utf-8")) + assert saved.storage == "file" + assert payload["storage"] == "file" + assert payload["tokens"]["access_token"] == "access" + if os.name != "nt": + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + + loaded = auth.load_xai_oauth_credential() + assert loaded is not None + assert loaded.access_token == "access" + assert loaded.refresh_token == "refresh" + assert loaded.account_id == "acct" + assert loaded.storage == "file" + + +def test_keyring_storage_keeps_tokens_out_of_metadata(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("NANOBOT_HOME", str(tmp_path)) + secret: dict[str, object] = {} + + def fake_set(tokens: dict[str, object]) -> bool: + secret.update(tokens) + return True + + monkeypatch.setattr(auth, "_keyring_set", fake_set) + monkeypatch.setattr(auth, "_keyring_get", lambda: dict(secret)) + + auth.save_xai_oauth_credential( + auth.XaiOAuthCredential( + access_token="access", + refresh_token="refresh", + expires_at=123.0, + account_id="acct", + ) + ) + + payload = json.loads(auth.get_xai_oauth_metadata_path().read_text(encoding="utf-8")) + assert payload["storage"] == "keyring" + assert "tokens" not in payload + assert auth.load_xai_oauth_credential().access_token == "access" + + +def test_exchange_xai_oauth_code_sends_required_code_challenge(monkeypatch) -> None: + captured: dict[str, object] = {} + + class FakeResponse: + status_code = 200 + text = "" + + def json(self) -> dict[str, object]: + return {"access_token": "access", "refresh_token": "refresh", "expires_in": 3600} + + class FakeClient: + def __init__(self, *args, **kwargs) -> None: + pass + + def __enter__(self): + return self + + def __exit__(self, *args) -> None: + pass + + def post(self, url: str, headers: dict[str, str], data: dict[str, str]) -> FakeResponse: + captured["url"] = url + captured["headers"] = headers + captured["data"] = data + return FakeResponse() + + monkeypatch.setattr(auth.httpx, "Client", FakeClient) + endpoints = auth.XaiOAuthEndpoints( + authorization_endpoint="https://auth.x.ai/authorize", + token_endpoint="https://auth.x.ai/oauth/token", + ) + + credential = auth.exchange_xai_oauth_code("code", verifier="verifier", endpoints=endpoints) + + assert credential.access_token == "access" + assert captured["url"] == "https://auth.x.ai/oauth/token" + data = captured["data"] + assert data["code_verifier"] == "verifier" + assert data["code_challenge"] == auth.pkce_challenge("verifier") + assert data["code_challenge_method"] == "S256" + + +def test_rejects_non_xai_discovery_endpoints() -> None: + with pytest.raises(RuntimeError): + auth._validate_xai_endpoint("https://example.com/oauth/token", "token_endpoint") diff --git a/tests/providers/test_xai_oauth_provider.py b/tests/providers/test_xai_oauth_provider.py new file mode 100644 index 000000000..201407cdb --- /dev/null +++ b/tests/providers/test_xai_oauth_provider.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +import asyncio + +from nanobot.config.schema import XaiOAuthXSearchConfig +import nanobot.providers.xai_oauth_provider as xai_oauth_provider +from nanobot.providers.xai_oauth_provider import ( + XaiOAuthCredential, + XaiOAuthProvider, + _build_xai_responses_body, + _strip_model_prefix, +) + + +def test_xai_oauth_strip_prefix_supports_aliases() -> None: + assert _strip_model_prefix("xai-oauth/grok-4.3") == "grok-4.3" + assert _strip_model_prefix("xai_oauth/grok-4.3") == "grok-4.3" + assert _strip_model_prefix("grok-oauth/grok-4.3") == "grok-4.3" + assert _strip_model_prefix("grok-4.3") == "grok-4.3" + + +def test_build_xai_responses_body_keeps_system_prompt_in_input() -> None: + body = _build_xai_responses_body( + messages=[ + {"role": "system", "content": "You are nanobot."}, + {"role": "user", "content": "hi"}, + ], + tools=[ + { + "type": "function", + "function": { + "name": "ping", + "description": "Ping", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + model="xai-oauth/grok-4.3", + max_tokens=32, + temperature=0.2, + reasoning_effort="high", + tool_choice=None, + ) + + assert body["model"] == "grok-4.3" + assert "instructions" not in body + assert body["input"][0] == { + "role": "system", + "content": [{"type": "input_text", "text": "You are nanobot."}], + } + assert body["input"][1]["role"] == "user" + assert body["max_output_tokens"] == 32 + assert body["temperature"] == 0.2 + assert body["reasoning"] == {"effort": "high"} + assert body["tools"][0]["name"] == "ping" + + +def test_build_xai_responses_body_attaches_hosted_x_search_by_default() -> None: + body = _build_xai_responses_body( + messages=[{"role": "user", "content": "what is happening on X?"}], + tools=None, + model="xai-oauth/grok-4.3", + max_tokens=32, + temperature=0.2, + reasoning_effort=None, + tool_choice=None, + hosted_x_search=XaiOAuthXSearchConfig(), + ) + + assert body["tools"] == [{"type": "x_search"}] + + +def test_build_xai_responses_body_can_customize_hosted_x_search() -> None: + body = _build_xai_responses_body( + messages=[{"role": "user", "content": "what is happening on X?"}], + tools=None, + model="xai-oauth/grok-4.3", + max_tokens=32, + temperature=0.2, + reasoning_effort=None, + tool_choice=None, + hosted_x_search=XaiOAuthXSearchConfig( + allowed_x_handles=["@xai", " nanobot "], + enable_image_understanding=True, + ), + ) + + assert body["tools"] == [ + { + "type": "x_search", + "allowed_x_handles": ["xai", "nanobot"], + "enable_image_understanding": True, + } + ] + + +def test_build_xai_responses_body_omits_disabled_hosted_x_search() -> None: + body = _build_xai_responses_body( + messages=[{"role": "user", "content": "hi"}], + tools=None, + model="xai-oauth/grok-4.3", + max_tokens=32, + temperature=0.2, + reasoning_effort=None, + tool_choice=None, + hosted_x_search=XaiOAuthXSearchConfig(enable=False), + ) + + assert "tools" not in body + + +def test_xai_oauth_provider_refreshes_once_on_401(monkeypatch) -> None: + async def run() -> None: + response = await provider.chat([{"role": "user", "content": "hi"}]) + + assert response.content == "ok" + assert response.finish_reason == "stop" + assert calls == [("resolve", False), ("resolve", True)] + + provider = XaiOAuthProvider(default_model="xai-oauth/grok-4.3") + credentials = [ + XaiOAuthCredential(access_token="expired"), + XaiOAuthCredential(access_token="fresh"), + ] + calls: list[tuple[str, bool]] = [] + + def fake_resolve(*, force_refresh: bool = False) -> XaiOAuthCredential: + calls.append(("resolve", force_refresh)) + return credentials.pop(0) + + async def fake_request(credential, body, on_content_delta=None, on_tool_call_delta=None): + from nanobot.providers.xai_oauth_provider import _XaiHTTPError + + if credential.access_token == "expired": + raise _XaiHTTPError("expired", status_code=401) + return "ok", [], "stop" + + monkeypatch.setattr(xai_oauth_provider, "resolve_xai_oauth_credential", fake_resolve) + monkeypatch.setattr(xai_oauth_provider, "_request_xai", fake_request) + + asyncio.run(run())