feat(providers): add xAI Grok OAuth with capability-gated X Search (#5035)
This commit is contained in:
chengyongru
@@ -13,6 +13,7 @@ __all__ = [
|
||||
"AnthropicProvider",
|
||||
"OpenAICompatProvider",
|
||||
"OpenAICodexProvider",
|
||||
"XAIGrokProvider",
|
||||
"GitHubCopilotProvider",
|
||||
"AzureOpenAIProvider",
|
||||
"BedrockProvider",
|
||||
@@ -22,6 +23,7 @@ _LAZY_IMPORTS = {
|
||||
"AnthropicProvider": ".anthropic_provider",
|
||||
"OpenAICompatProvider": ".openai_compat_provider",
|
||||
"OpenAICodexProvider": ".openai_codex_provider",
|
||||
"XAIGrokProvider": ".xai_grok_provider",
|
||||
"GitHubCopilotProvider": ".github_copilot_provider",
|
||||
"AzureOpenAIProvider": ".azure_openai_provider",
|
||||
"BedrockProvider": ".bedrock_provider",
|
||||
@@ -34,6 +36,7 @@ if TYPE_CHECKING:
|
||||
from nanobot.providers.github_copilot_provider import GitHubCopilotProvider
|
||||
from nanobot.providers.openai_codex_provider import OpenAICodexProvider
|
||||
from nanobot.providers.openai_compat_provider import OpenAICompatProvider
|
||||
from nanobot.providers.xai_grok_provider import XAIGrokProvider
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
|
||||
@@ -60,10 +60,10 @@ def _make_provider_core(
|
||||
if spec and spec.is_transcription_only:
|
||||
raise ValueError(f"Provider '{provider_name}' only supports transcription.")
|
||||
backend = spec.backend if spec else "openai_compat"
|
||||
if p and p.proxy and backend not in {"openai_compat", "openai_codex"}:
|
||||
if p and p.proxy and backend not in {"openai_compat", "openai_codex", "xai_grok"}:
|
||||
raise ValueError(
|
||||
f"providers.{provider_name}.proxy is only supported for "
|
||||
"OpenAI-compatible providers and OpenAI Codex."
|
||||
"OpenAI-compatible providers, OpenAI Codex, and xAI Grok."
|
||||
)
|
||||
|
||||
if backend == "azure_openai":
|
||||
@@ -91,6 +91,14 @@ def _make_provider_core(
|
||||
proxy=getattr(p, "proxy", None) if p else None,
|
||||
extra_body=p.extra_body if p else None,
|
||||
)
|
||||
elif backend == "xai_grok":
|
||||
from nanobot.providers.xai_grok_provider import XAIGrokProvider
|
||||
|
||||
provider = XAIGrokProvider(
|
||||
default_model=model,
|
||||
proxy=getattr(p, "proxy", None) if p else None,
|
||||
extra_body=p.extra_body if p else None,
|
||||
)
|
||||
elif backend == "azure_openai":
|
||||
from nanobot.providers.azure_openai_provider import AzureOpenAIProvider
|
||||
|
||||
|
||||
@@ -47,7 +47,8 @@ class ProviderSpec:
|
||||
settings_alias_for: str = "" # compatibility alias grouped under this provider in Settings
|
||||
|
||||
# which provider implementation to use
|
||||
# "openai_compat" | "anthropic" | "azure_openai" | "openai_codex" | "github_copilot" | "bedrock"
|
||||
# "openai_compat" | "anthropic" | "azure_openai" | "openai_codex" | "xai_grok"
|
||||
# | "github_copilot" | "bedrock"
|
||||
backend: str = "openai_compat"
|
||||
|
||||
# extra env vars / request headers supplied by the provider integration.
|
||||
@@ -420,6 +421,25 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
|
||||
default_api_base="https://chatgpt.com/backend-api",
|
||||
is_oauth=True,
|
||||
),
|
||||
# xAI subscription: OAuth-based, with capability-gated server-hosted X Search.
|
||||
ProviderSpec(
|
||||
name="xai_grok",
|
||||
keywords=("xai-grok", "xai_grok"),
|
||||
env_key="",
|
||||
display_name="xAI Grok",
|
||||
model_catalog="builtin",
|
||||
builtin_models=(
|
||||
ProviderModelSpec(
|
||||
id="xai-grok/grok-4.5",
|
||||
label="Grok 4.5",
|
||||
description="Grok via xAI subscription; X Search is enabled when supported.",
|
||||
context_window=500000,
|
||||
),
|
||||
),
|
||||
backend="xai_grok",
|
||||
default_api_base="https://cli-chat-proxy.grok.com/v1",
|
||||
is_oauth=True,
|
||||
),
|
||||
# GitHub Copilot: OAuth-based
|
||||
ProviderSpec(
|
||||
name="github_copilot",
|
||||
|
||||
@@ -0,0 +1,545 @@
|
||||
"""xAI subscription provider with capability-gated hosted X Search."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from loguru import logger
|
||||
|
||||
from nanobot import __version__
|
||||
from nanobot.providers.base import (
|
||||
LLMProvider,
|
||||
LLMResponse,
|
||||
ToolCallRequest,
|
||||
resolve_stream_idle_timeout_s,
|
||||
)
|
||||
from nanobot.providers.openai_responses import (
|
||||
consume_sse_with_reasoning,
|
||||
convert_messages,
|
||||
convert_tools,
|
||||
)
|
||||
from nanobot.providers.xai_oauth import (
|
||||
XAI_CLIENT_VERSION,
|
||||
XAIToken,
|
||||
get_xai_oauth_token,
|
||||
)
|
||||
|
||||
DEFAULT_XAI_GROK_URL = "https://cli-chat-proxy.grok.com/v1/responses"
|
||||
DEFAULT_XAI_GROK_MODELS_URL = "https://cli-chat-proxy.grok.com/v1/models"
|
||||
DEFAULT_XAI_GROK_MODEL = "xai-grok/grok-4.5"
|
||||
_MODEL_CAPABILITIES_TTL_S = 5 * 60
|
||||
_MAX_ERROR_BODY_CHARS = 1000
|
||||
_SENSITIVE_ERROR_KEYS = {
|
||||
"accesstoken",
|
||||
"apikey",
|
||||
"authorization",
|
||||
"idtoken",
|
||||
"refreshtoken",
|
||||
}
|
||||
|
||||
|
||||
class XAIGrokProvider(LLMProvider):
|
||||
"""Call xAI's subscription proxy and expose supported hosted tools."""
|
||||
|
||||
supports_progress_deltas = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
default_model: str = DEFAULT_XAI_GROK_MODEL,
|
||||
proxy: str | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
):
|
||||
super().__init__(api_key=None, api_base=None)
|
||||
self.default_model = default_model
|
||||
self.proxy = proxy or None
|
||||
self._extra_body = dict(extra_body or {})
|
||||
self._model_capabilities: dict[str, bool] | None = None
|
||||
self._model_capabilities_fetched_at = 0.0
|
||||
|
||||
async def _supports_backend_search(self, token: XAIToken, model: str) -> bool:
|
||||
now = time.monotonic()
|
||||
capabilities = self._model_capabilities
|
||||
if (
|
||||
capabilities is None
|
||||
or now - self._model_capabilities_fetched_at >= _MODEL_CAPABILITIES_TTL_S
|
||||
):
|
||||
try:
|
||||
capabilities = await _fetch_xai_model_capabilities(
|
||||
DEFAULT_XAI_GROK_MODELS_URL,
|
||||
_build_model_headers(token),
|
||||
proxy=self.proxy,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"xAI model capability lookup failed; hosted X Search disabled for model {}: "
|
||||
"type={} error={}",
|
||||
model,
|
||||
type(exc).__name__,
|
||||
str(exc).strip() or "unexpected error",
|
||||
)
|
||||
capabilities = {}
|
||||
self._model_capabilities = capabilities
|
||||
self._model_capabilities_fetched_at = now
|
||||
else:
|
||||
self._model_capabilities = capabilities
|
||||
self._model_capabilities_fetched_at = now
|
||||
return capabilities.get(model, False)
|
||||
|
||||
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_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
|
||||
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
|
||||
) -> LLMResponse:
|
||||
wire_model = _strip_model_prefix(model or self.default_model)
|
||||
system_prompt, input_items = convert_messages(messages)
|
||||
|
||||
stage = "oauth_token"
|
||||
try:
|
||||
token = await asyncio.to_thread(get_xai_oauth_token, proxy=self.proxy)
|
||||
stage = "model_capabilities"
|
||||
supports_backend_search = await self._supports_backend_search(token, wire_model)
|
||||
converted_tools = convert_tools(tools or [])
|
||||
if supports_backend_search:
|
||||
converted_tools = [
|
||||
tool for tool in converted_tools if tool.get("name") != "x_search"
|
||||
]
|
||||
converted_tools.append({"type": "x_search"})
|
||||
|
||||
body: dict[str, Any] = {
|
||||
"model": wire_model,
|
||||
"store": False,
|
||||
"stream": True,
|
||||
"instructions": system_prompt,
|
||||
"input": input_items,
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"tools": converted_tools,
|
||||
"tool_choice": tool_choice or "auto",
|
||||
"parallel_tool_calls": True,
|
||||
"stream_tool_calls": True,
|
||||
"max_output_tokens": max_tokens,
|
||||
"temperature": temperature,
|
||||
"reasoning": _build_reasoning_options(reasoning_effort),
|
||||
}
|
||||
if self._extra_body:
|
||||
body.update(self._extra_body)
|
||||
|
||||
headers = _build_headers(token.access, wire_model)
|
||||
stage = "xai_request"
|
||||
try:
|
||||
result = await _request_xai(
|
||||
DEFAULT_XAI_GROK_URL,
|
||||
headers,
|
||||
body,
|
||||
proxy=self.proxy,
|
||||
on_content_delta=on_content_delta,
|
||||
on_thinking_delta=on_thinking_delta,
|
||||
on_tool_call_delta=on_tool_call_delta,
|
||||
)
|
||||
except _XAIHTTPError as exc:
|
||||
if exc.status_code != 401:
|
||||
raise
|
||||
stage = "oauth_refresh"
|
||||
token = await asyncio.to_thread(
|
||||
get_xai_oauth_token,
|
||||
proxy=self.proxy,
|
||||
force_refresh=True,
|
||||
)
|
||||
self._model_capabilities = None
|
||||
self._model_capabilities_fetched_at = 0.0
|
||||
headers = _build_headers(token.access, wire_model)
|
||||
stage = "xai_request_retry"
|
||||
result = await _request_xai(
|
||||
DEFAULT_XAI_GROK_URL,
|
||||
headers,
|
||||
body,
|
||||
proxy=self.proxy,
|
||||
on_content_delta=on_content_delta,
|
||||
on_thinking_delta=on_thinking_delta,
|
||||
on_tool_call_delta=on_tool_call_delta,
|
||||
)
|
||||
|
||||
content, tool_calls, finish_reason, usage, reasoning_content = result
|
||||
return LLMResponse(
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
finish_reason=finish_reason,
|
||||
usage=usage,
|
||||
reasoning_content=reasoning_content,
|
||||
)
|
||||
except Exception as exc:
|
||||
response = _xai_error_response(exc)
|
||||
logger.warning(
|
||||
"xAI subscription request failed: stage={} type={} retryable={} status={} "
|
||||
"error_type={} error_code={} response_body={}",
|
||||
stage,
|
||||
type(exc).__name__,
|
||||
response.error_should_retry,
|
||||
response.error_status_code,
|
||||
response.error_type,
|
||||
response.error_code,
|
||||
getattr(exc, "response_body", None),
|
||||
)
|
||||
return response
|
||||
|
||||
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:
|
||||
return await self._call_xai(
|
||||
messages,
|
||||
tools,
|
||||
model,
|
||||
max_tokens,
|
||||
temperature,
|
||||
reasoning_effort,
|
||||
tool_choice,
|
||||
on_content_delta,
|
||||
on_thinking_delta,
|
||||
on_tool_call_delta,
|
||||
)
|
||||
|
||||
def get_default_model(self) -> str:
|
||||
return self.default_model
|
||||
|
||||
|
||||
def _strip_model_prefix(model: str) -> str:
|
||||
if model.startswith("xai-grok/") or model.startswith("xai_grok/"):
|
||||
return model.split("/", 1)[1]
|
||||
return model
|
||||
|
||||
|
||||
def _build_reasoning_options(reasoning_effort: str | None) -> dict[str, str]:
|
||||
options = {"summary": "concise"}
|
||||
if reasoning_effort and reasoning_effort.lower() != "none":
|
||||
options["effort"] = reasoning_effort
|
||||
return options
|
||||
|
||||
|
||||
def _build_headers(token: str, model: str) -> dict[str, str]:
|
||||
conversation_id = str(uuid.uuid4())
|
||||
return {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"X-XAI-Token-Auth": "xai-grok-cli",
|
||||
"x-authenticateresponse": "authenticate-response",
|
||||
"x-grok-client-version": XAI_CLIENT_VERSION,
|
||||
"x-grok-client-identifier": "nanobot",
|
||||
"x-grok-client-mode": "headless",
|
||||
"x-grok-conv-id": conversation_id,
|
||||
"x-grok-req-id": str(uuid.uuid4()),
|
||||
"x-grok-model-override": model,
|
||||
"x-grok-session-id": conversation_id,
|
||||
"x-grok-agent-id": str(uuid.uuid4()),
|
||||
"User-Agent": f"nanobot/{__version__} (python)",
|
||||
"accept": "text/event-stream",
|
||||
"content-type": "application/json",
|
||||
}
|
||||
|
||||
|
||||
def _build_model_headers(token: XAIToken) -> dict[str, str]:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token.access}",
|
||||
"X-XAI-Token-Auth": "xai-grok-cli",
|
||||
"x-grok-client-version": XAI_CLIENT_VERSION,
|
||||
"x-grok-client-identifier": "nanobot",
|
||||
"x-grok-client-mode": "headless",
|
||||
"User-Agent": f"nanobot/{__version__} (python)",
|
||||
"accept": "application/json",
|
||||
}
|
||||
claims = _decode_access_token_claims(token.access)
|
||||
user_id = claims.get("sub")
|
||||
if claims.get("principal_type") == "Team":
|
||||
user_id = claims.get("principal_id") or user_id
|
||||
if isinstance(user_id, str) and user_id:
|
||||
headers["x-userid"] = user_id
|
||||
email = claims.get("email")
|
||||
if not isinstance(email, str) or "@" not in email:
|
||||
email = token.account_id if token.account_id and "@" in token.account_id else None
|
||||
if email:
|
||||
headers["x-email"] = email
|
||||
return headers
|
||||
|
||||
|
||||
def _decode_access_token_claims(token: str) -> dict[str, Any]:
|
||||
"""Read identity hints from the signed token; the server still authenticates it."""
|
||||
parts = token.split(".")
|
||||
if len(parts) < 2 or not parts[1]:
|
||||
return {}
|
||||
payload = parts[1]
|
||||
try:
|
||||
decoded = base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4))
|
||||
claims = json.loads(decoded)
|
||||
except (ValueError, TypeError):
|
||||
return {}
|
||||
return claims if isinstance(claims, dict) else {}
|
||||
|
||||
|
||||
class _XAIHTTPError(RuntimeError):
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
*,
|
||||
status_code: int,
|
||||
retry_after: float | None = None,
|
||||
error_type: str | None = None,
|
||||
error_code: str | None = None,
|
||||
should_retry: bool | None = None,
|
||||
response_body: str | None = None,
|
||||
):
|
||||
super().__init__(message)
|
||||
self.status_code = status_code
|
||||
self.retry_after = retry_after
|
||||
self.error_type = error_type
|
||||
self.error_code = error_code
|
||||
self.should_retry = should_retry
|
||||
self.response_body = response_body
|
||||
|
||||
|
||||
async def _fetch_xai_model_capabilities(
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
*,
|
||||
proxy: str | None = None,
|
||||
) -> dict[str, bool]:
|
||||
client_kwargs: dict[str, Any] = {"timeout": 10.0, "follow_redirects": False}
|
||||
if proxy:
|
||||
client_kwargs.update(proxy=proxy, trust_env=False)
|
||||
async with httpx.AsyncClient(**client_kwargs) as client:
|
||||
response = await client.get(url, headers=headers)
|
||||
if response.status_code != 200:
|
||||
raw = response.content.decode("utf-8", "ignore")
|
||||
raise _build_xai_http_error(response.status_code, response.headers, raw)
|
||||
try:
|
||||
payload = response.json()
|
||||
except ValueError as exc:
|
||||
raise RuntimeError("xAI model catalog returned invalid JSON.") from exc
|
||||
return _parse_xai_model_capabilities(payload)
|
||||
|
||||
|
||||
def _parse_xai_model_capabilities(payload: Any) -> dict[str, bool]:
|
||||
if isinstance(payload, dict):
|
||||
rows = payload.get("data")
|
||||
if not isinstance(rows, list):
|
||||
rows = payload.get("models")
|
||||
else:
|
||||
rows = payload
|
||||
if not isinstance(rows, list):
|
||||
return {}
|
||||
|
||||
capabilities: dict[str, bool] = {}
|
||||
for row in rows:
|
||||
if not isinstance(row, dict):
|
||||
continue
|
||||
meta = row.get("_meta") if isinstance(row.get("_meta"), dict) else {}
|
||||
support_value = row.get("supportsBackendSearch")
|
||||
if not isinstance(support_value, bool):
|
||||
support_value = row.get("supports_backend_search")
|
||||
if not isinstance(support_value, bool):
|
||||
support_value = meta.get("supportsBackendSearch")
|
||||
if not isinstance(support_value, bool):
|
||||
support_value = meta.get("supports_backend_search")
|
||||
supports_backend_search = support_value if isinstance(support_value, bool) else False
|
||||
|
||||
identifiers = (
|
||||
row.get("model"),
|
||||
row.get("modelId"),
|
||||
row.get("id"),
|
||||
meta.get("model"),
|
||||
meta.get("modelId"),
|
||||
)
|
||||
for identifier in identifiers:
|
||||
if isinstance(identifier, str) and identifier.strip():
|
||||
capabilities[_strip_model_prefix(identifier.strip())] = supports_backend_search
|
||||
return capabilities
|
||||
|
||||
|
||||
async def _request_xai(
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
body: dict[str, Any],
|
||||
*,
|
||||
proxy: str | 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,
|
||||
) -> tuple[str, list[ToolCallRequest], str, dict[str, int], str | None]:
|
||||
client_kwargs: dict[str, Any] = {"timeout": resolve_stream_idle_timeout_s()}
|
||||
if proxy:
|
||||
client_kwargs.update(proxy=proxy, trust_env=False)
|
||||
async with httpx.AsyncClient(**client_kwargs) as client:
|
||||
async with client.stream("POST", url, headers=headers, json=body) as response:
|
||||
if response.status_code != 200:
|
||||
content = await response.aread()
|
||||
raw = content.decode("utf-8", "ignore")
|
||||
raise _build_xai_http_error(response.status_code, response.headers, raw)
|
||||
return await consume_sse_with_reasoning(
|
||||
response,
|
||||
on_content_delta=on_content_delta,
|
||||
on_tool_call_delta=on_tool_call_delta,
|
||||
on_reasoning_delta=on_thinking_delta,
|
||||
)
|
||||
|
||||
|
||||
def _build_xai_http_error(
|
||||
status_code: int,
|
||||
headers: httpx.Headers,
|
||||
raw: str,
|
||||
) -> _XAIHTTPError:
|
||||
retry_after = LLMProvider._extract_retry_after_from_headers(headers)
|
||||
error_type, error_code = LLMProvider._extract_error_type_code(raw)
|
||||
response_body = _bounded_error_body(raw)
|
||||
return _XAIHTTPError(
|
||||
_friendly_error(status_code, response_body),
|
||||
status_code=status_code,
|
||||
retry_after=retry_after,
|
||||
error_type=error_type,
|
||||
error_code=error_code,
|
||||
should_retry=_should_retry_status(status_code, error_type, error_code, raw),
|
||||
response_body=response_body,
|
||||
)
|
||||
|
||||
|
||||
def _bounded_error_body(raw: str) -> str | None:
|
||||
text = raw.strip()
|
||||
if not text:
|
||||
return None
|
||||
|
||||
try:
|
||||
payload = json.loads(text)
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
else:
|
||||
text = json.dumps(
|
||||
_redact_error_payload(payload),
|
||||
ensure_ascii=False,
|
||||
separators=(",", ":"),
|
||||
)
|
||||
|
||||
text = re.sub(r"(?i)(bearer\s+)[a-z0-9._~+/=-]+", r"\1[REDACTED]", text)
|
||||
text = " ".join(text.split())
|
||||
if len(text) > _MAX_ERROR_BODY_CHARS:
|
||||
return f"{text[:_MAX_ERROR_BODY_CHARS]}…"
|
||||
return text
|
||||
|
||||
|
||||
def _redact_error_payload(payload: Any) -> Any:
|
||||
if isinstance(payload, dict):
|
||||
return {
|
||||
key: "[REDACTED]" if _is_sensitive_error_key(key) else _redact_error_payload(value)
|
||||
for key, value in payload.items()
|
||||
}
|
||||
if isinstance(payload, list):
|
||||
return [_redact_error_payload(value) for value in payload]
|
||||
return payload
|
||||
|
||||
|
||||
def _is_sensitive_error_key(key: str) -> bool:
|
||||
normalized = re.sub(r"[^a-z0-9]", "", key.casefold())
|
||||
return normalized in _SENSITIVE_ERROR_KEYS
|
||||
|
||||
|
||||
def _friendly_error(status_code: int, response_body: str | None = None) -> str:
|
||||
if status_code == 401:
|
||||
message = "xAI rejected the login. Sign in again with `nanobot provider login xai-grok`."
|
||||
elif status_code == 403:
|
||||
message = "This xAI account or subscription cannot access the Grok subscription endpoint."
|
||||
elif status_code == 426:
|
||||
message = "xAI requires a newer Grok client version. Update nanobot and try again."
|
||||
elif status_code == 429:
|
||||
message = "xAI usage quota or rate limit reached. Please try again later."
|
||||
else:
|
||||
message = f"xAI subscription endpoint returned HTTP {status_code}."
|
||||
if response_body:
|
||||
return f"{message} Response body: {response_body}"
|
||||
return message
|
||||
|
||||
|
||||
def _xai_error_response(exc: Exception) -> LLMResponse:
|
||||
status_code = getattr(exc, "status_code", None)
|
||||
should_retry = getattr(exc, "should_retry", None)
|
||||
error_kind: str | None = None
|
||||
if isinstance(exc, (httpx.TimeoutException, asyncio.TimeoutError)):
|
||||
error_kind = "timeout"
|
||||
should_retry = True if should_retry is None else should_retry
|
||||
elif isinstance(exc, (httpx.NetworkError, httpx.TransportError)):
|
||||
error_kind = "connection"
|
||||
should_retry = True if should_retry is None else should_retry
|
||||
elif isinstance(exc, _XAIHTTPError):
|
||||
error_kind = "http"
|
||||
if status_code is not None and should_retry is None:
|
||||
should_retry = _should_retry_status(
|
||||
int(status_code),
|
||||
getattr(exc, "error_type", None),
|
||||
getattr(exc, "error_code", None),
|
||||
None,
|
||||
)
|
||||
message = str(exc).strip() or "unexpected error"
|
||||
retry_after = getattr(exc, "retry_after", None)
|
||||
return LLMResponse(
|
||||
content=f"Error calling xAI ({type(exc).__name__}): {message}",
|
||||
finish_reason="error",
|
||||
retry_after=retry_after,
|
||||
error_status_code=int(status_code) if status_code is not None else None,
|
||||
error_kind=error_kind,
|
||||
error_type=getattr(exc, "error_type", None),
|
||||
error_code=getattr(exc, "error_code", None),
|
||||
error_retry_after_s=retry_after,
|
||||
error_should_retry=should_retry,
|
||||
)
|
||||
|
||||
|
||||
def _should_retry_status(
|
||||
status_code: int,
|
||||
error_type: str | None,
|
||||
error_code: str | None,
|
||||
content: str | None,
|
||||
) -> bool:
|
||||
if status_code == 429:
|
||||
return LLMProvider._is_retryable_429_response(
|
||||
LLMResponse(
|
||||
content=content or "",
|
||||
finish_reason="error",
|
||||
error_status_code=status_code,
|
||||
error_type=error_type,
|
||||
error_code=error_code,
|
||||
)
|
||||
)
|
||||
return status_code in LLMProvider._RETRYABLE_STATUS_CODES or status_code >= 500
|
||||
@@ -0,0 +1,745 @@
|
||||
"""xAI subscription OAuth (Authorization Code + PKCE) support.
|
||||
|
||||
This integration follows the public OAuth client contract used by Grok Build.
|
||||
Credentials are stored separately from Grok Build so rotating refresh tokens are
|
||||
never shared between the two applications.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import os
|
||||
import queue
|
||||
import secrets
|
||||
import threading
|
||||
import time
|
||||
import webbrowser
|
||||
from collections.abc import Callable
|
||||
from contextlib import suppress
|
||||
from dataclasses import asdict, dataclass
|
||||
from http import HTTPStatus
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qs, urlencode, urlsplit
|
||||
|
||||
import httpx
|
||||
from filelock import FileLock
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.config.paths import get_data_dir
|
||||
from nanobot.utils.helpers import _write_text_atomic
|
||||
|
||||
XAI_OAUTH_ISSUER = "https://auth.x.ai"
|
||||
XAI_CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828"
|
||||
XAI_CLIENT_VERSION = "0.2.109"
|
||||
XAI_ALLOWED_CALLBACK_ORIGIN = "https://accounts.x.ai"
|
||||
XAI_OAUTH_SCOPES = (
|
||||
"openid",
|
||||
"profile",
|
||||
"email",
|
||||
"offline_access",
|
||||
"grok-cli:access",
|
||||
"api:access",
|
||||
"conversations:read",
|
||||
"conversations:write",
|
||||
"workspaces:read",
|
||||
"workspaces:write",
|
||||
)
|
||||
|
||||
_DISCOVERY_URL = f"{XAI_OAUTH_ISSUER}/.well-known/openid-configuration"
|
||||
_TOKEN_REFRESH_MARGIN_MS = 5 * 60 * 1000
|
||||
_DEFAULT_TOKEN_TTL_S = 60 * 60
|
||||
_HTTP_TIMEOUT_S = 15.0
|
||||
|
||||
|
||||
class XAIOAuthError(RuntimeError):
|
||||
"""An actionable xAI OAuth failure with no credential material."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class XAIToken:
|
||||
"""Persisted xAI OAuth token material."""
|
||||
|
||||
access: str
|
||||
refresh: str | None
|
||||
expires: int
|
||||
account_id: str | None = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, value: Any) -> XAIToken | None:
|
||||
if not isinstance(value, dict):
|
||||
return None
|
||||
access = value.get("access")
|
||||
if not isinstance(access, str) or not access:
|
||||
return None
|
||||
refresh = value.get("refresh")
|
||||
if not isinstance(refresh, str) or not refresh:
|
||||
refresh = None
|
||||
try:
|
||||
expires = int(value.get("expires") or 0)
|
||||
except (TypeError, ValueError):
|
||||
expires = 0
|
||||
account_id = value.get("account_id")
|
||||
if not isinstance(account_id, str) or not account_id:
|
||||
account_id = None
|
||||
return cls(access=access, refresh=refresh, expires=expires, account_id=account_id)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Discovery:
|
||||
authorization_endpoint: str
|
||||
token_endpoint: str
|
||||
userinfo_endpoint: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _CallbackResult:
|
||||
code: str | None = None
|
||||
state: str | None = None
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class XAIOAuthLoginFlow:
|
||||
"""Pending xAI OAuth login that can finish through loopback or an authorization code."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
authorization_url: str,
|
||||
redirect_uri: str,
|
||||
verifier: str,
|
||||
state: str,
|
||||
discovery: _Discovery,
|
||||
proxy: str | None,
|
||||
result_queue: queue.Queue[_CallbackResult],
|
||||
server: ThreadingHTTPServer,
|
||||
timeout_s: float,
|
||||
) -> None:
|
||||
self.authorization_url = authorization_url
|
||||
self.redirect_uri = redirect_uri
|
||||
self._verifier = verifier
|
||||
self._state = state
|
||||
self._discovery = discovery
|
||||
self._proxy = proxy
|
||||
self._result_queue = result_queue
|
||||
self._server = server
|
||||
self._expires_at = time.monotonic() + timeout_s
|
||||
self._lock = threading.Lock()
|
||||
self._stop_event = threading.Event()
|
||||
self._token: XAIToken | None = None
|
||||
self._error: Exception | None = None
|
||||
self._closed = False
|
||||
self._server_thread = threading.Thread(
|
||||
target=_serve_callback_server,
|
||||
args=(server, self._stop_event),
|
||||
name="nanobot-xai-grok-oauth-callback",
|
||||
daemon=True,
|
||||
)
|
||||
self._server_thread.start()
|
||||
self._timeout_timer = threading.Timer(timeout_s, self._expire)
|
||||
self._timeout_timer.daemon = True
|
||||
self._timeout_timer.start()
|
||||
|
||||
@property
|
||||
def expired(self) -> bool:
|
||||
return time.monotonic() >= self._expires_at
|
||||
|
||||
@property
|
||||
def remaining_seconds(self) -> int:
|
||||
return max(0, int(self._expires_at - time.monotonic()))
|
||||
|
||||
def complete(self, authorization_code: str | None = None) -> XAIToken | None:
|
||||
"""Complete this flow, or return ``None`` while loopback is still pending."""
|
||||
with self._lock:
|
||||
if self._token is not None:
|
||||
return self._token
|
||||
self._raise_if_finished()
|
||||
|
||||
callback: _CallbackResult | None
|
||||
if authorization_code is not None:
|
||||
callback = _CallbackResult(code=authorization_code.strip())
|
||||
else:
|
||||
try:
|
||||
callback = self._result_queue.get_nowait()
|
||||
except queue.Empty:
|
||||
callback = None
|
||||
if callback is None:
|
||||
with self._lock:
|
||||
if self._token is not None:
|
||||
return self._token
|
||||
self._raise_if_finished()
|
||||
if self.expired:
|
||||
self._expire_locked()
|
||||
self._raise_if_finished()
|
||||
return None
|
||||
return self._finish(callback)
|
||||
|
||||
def wait(self, timeout_s: float) -> XAIToken:
|
||||
"""Wait for the loopback callback and complete this flow."""
|
||||
with self._lock:
|
||||
if self._token is not None:
|
||||
return self._token
|
||||
self._raise_if_finished()
|
||||
try:
|
||||
callback = self._result_queue.get(timeout=timeout_s)
|
||||
except queue.Empty as exc:
|
||||
with self._lock:
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
raise XAIOAuthError(
|
||||
"Timed out waiting for xAI sign-in. Run "
|
||||
"`nanobot provider login xai-grok` to try again."
|
||||
) from exc
|
||||
return self._finish(callback)
|
||||
|
||||
def cancel(self) -> None:
|
||||
"""Stop the callback listener for an abandoned flow."""
|
||||
with self._lock:
|
||||
if self._token is None and self._error is None:
|
||||
self._error = XAIOAuthError("xAI sign-in was cancelled.")
|
||||
self._close_locked()
|
||||
|
||||
def _finish(self, callback: _CallbackResult) -> XAIToken:
|
||||
with self._lock:
|
||||
if self._token is not None:
|
||||
return self._token
|
||||
self._raise_if_finished()
|
||||
self._close_locked()
|
||||
try:
|
||||
token = _exchange_callback(
|
||||
callback,
|
||||
expected_state=self._state,
|
||||
discovery=self._discovery,
|
||||
verifier=self._verifier,
|
||||
redirect_uri=self.redirect_uri,
|
||||
proxy=self._proxy,
|
||||
)
|
||||
with _token_lock():
|
||||
_write_token(token)
|
||||
except Exception as exc:
|
||||
self._error = exc
|
||||
raise
|
||||
self._token = token
|
||||
return token
|
||||
|
||||
def _expire(self) -> None:
|
||||
with self._lock:
|
||||
if self._token is not None or self._error is not None:
|
||||
return
|
||||
self._expire_locked()
|
||||
|
||||
def _expire_locked(self) -> None:
|
||||
self._error = XAIOAuthError("xAI sign-in expired. Start a new sign-in flow.")
|
||||
self._close_locked()
|
||||
|
||||
def _raise_if_finished(self) -> None:
|
||||
if self._token is not None:
|
||||
return
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
|
||||
def _close_locked(self) -> None:
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
self._timeout_timer.cancel()
|
||||
self._stop_event.set()
|
||||
if threading.current_thread() is not self._server_thread:
|
||||
self._server_thread.join(timeout=2)
|
||||
|
||||
|
||||
def get_xai_oauth_storage_path() -> Path:
|
||||
"""Return the instance-scoped xAI OAuth credential path."""
|
||||
return get_data_dir() / "auth" / "xai.json"
|
||||
|
||||
|
||||
def get_xai_oauth_login_status() -> XAIToken | None:
|
||||
"""Return locally stored login state without making a network request."""
|
||||
return _load_token()
|
||||
|
||||
|
||||
def logout_xai_oauth() -> bool:
|
||||
"""Remove this instance's credentials while excluding token refreshes."""
|
||||
path = get_xai_oauth_storage_path()
|
||||
with _token_lock():
|
||||
try:
|
||||
path.unlink()
|
||||
except FileNotFoundError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def login_xai_oauth(
|
||||
*,
|
||||
print_fn: Callable[[str], None] = print,
|
||||
prompt_fn: Callable[[str], str] | None = None,
|
||||
proxy: str | None = None,
|
||||
callback_timeout_s: float = 600,
|
||||
browser_opener: Callable[[str], bool] = webbrowser.open,
|
||||
) -> XAIToken:
|
||||
"""Run xAI's browser-based OAuth flow and persist the resulting token.
|
||||
|
||||
``prompt_fn`` is used only when no local browser could be opened. It lets a
|
||||
headless user paste either the final callback URL or its authorization code.
|
||||
"""
|
||||
flow = start_xai_oauth_login(proxy=proxy, timeout_s=callback_timeout_s)
|
||||
|
||||
try:
|
||||
print_fn("Opening xAI sign-in in your browser...")
|
||||
print_fn(f"If it does not open automatically, visit:\n{flow.authorization_url}")
|
||||
opened = False
|
||||
with suppress(Exception):
|
||||
opened = bool(browser_opener(flow.authorization_url))
|
||||
|
||||
if not opened and prompt_fn is not None:
|
||||
pasted = prompt_fn("Paste the final callback URL (or authorization code)")
|
||||
token = flow.complete(pasted)
|
||||
if token is None: # pragma: no cover - pasted input always resolves a callback
|
||||
raise XAIOAuthError("xAI sign-in returned no authorization code.")
|
||||
return token
|
||||
return flow.wait(callback_timeout_s)
|
||||
finally:
|
||||
flow.cancel()
|
||||
|
||||
|
||||
def start_xai_oauth_login(
|
||||
*,
|
||||
proxy: str | None = None,
|
||||
timeout_s: float = 600,
|
||||
) -> XAIOAuthLoginFlow:
|
||||
"""Create a non-blocking OAuth flow for browser or pasted-callback completion."""
|
||||
discovery = _discover(proxy)
|
||||
verifier, challenge = _generate_pkce()
|
||||
state = secrets.token_urlsafe(32)
|
||||
nonce = secrets.token_urlsafe(32)
|
||||
result_queue: queue.Queue[_CallbackResult] = queue.Queue(maxsize=1)
|
||||
server = _make_callback_server(state, result_queue)
|
||||
redirect_uri = f"http://127.0.0.1:{server.server_port}/callback"
|
||||
authorization_url = _build_authorize_url(
|
||||
discovery.authorization_endpoint,
|
||||
redirect_uri=redirect_uri,
|
||||
challenge=challenge,
|
||||
state=state,
|
||||
nonce=nonce,
|
||||
)
|
||||
return XAIOAuthLoginFlow(
|
||||
authorization_url=authorization_url,
|
||||
redirect_uri=redirect_uri,
|
||||
verifier=verifier,
|
||||
state=state,
|
||||
discovery=discovery,
|
||||
proxy=proxy,
|
||||
result_queue=result_queue,
|
||||
server=server,
|
||||
timeout_s=timeout_s,
|
||||
)
|
||||
|
||||
|
||||
def complete_xai_oauth_login(
|
||||
flow: XAIOAuthLoginFlow,
|
||||
authorization_code: str | None = None,
|
||||
) -> XAIToken | None:
|
||||
"""Complete a pending login from loopback state or a pasted authorization code."""
|
||||
return flow.complete(authorization_code)
|
||||
|
||||
|
||||
def get_xai_oauth_token(
|
||||
*,
|
||||
proxy: str | None = None,
|
||||
min_ttl_ms: int = _TOKEN_REFRESH_MARGIN_MS,
|
||||
force_refresh: bool = False,
|
||||
) -> XAIToken:
|
||||
"""Load a usable token, refreshing it under an inter-process lock when needed."""
|
||||
token = _load_token()
|
||||
if token is None:
|
||||
raise XAIOAuthError(
|
||||
"xAI is not signed in. Run `nanobot provider login xai-grok` first."
|
||||
)
|
||||
if not force_refresh and _token_is_fresh(token, min_ttl_ms):
|
||||
return token
|
||||
if not token.refresh:
|
||||
if not force_refresh and token.expires > _now_ms():
|
||||
return token
|
||||
raise XAIOAuthError(
|
||||
"The xAI login has expired and cannot be refreshed. "
|
||||
"Run `nanobot provider login xai-grok` again."
|
||||
)
|
||||
|
||||
with _token_lock():
|
||||
latest = _load_token()
|
||||
if latest is None:
|
||||
raise XAIOAuthError(
|
||||
"xAI is not signed in. Run `nanobot provider login xai-grok` first."
|
||||
)
|
||||
if not force_refresh and _token_is_fresh(latest, min_ttl_ms):
|
||||
return latest
|
||||
if not latest.refresh:
|
||||
raise XAIOAuthError(
|
||||
"The xAI login has expired and cannot be refreshed. "
|
||||
"Run `nanobot provider login xai-grok` again."
|
||||
)
|
||||
refreshed = _refresh_token(latest, proxy)
|
||||
_write_token(refreshed)
|
||||
return refreshed
|
||||
|
||||
|
||||
def _discover(proxy: str | None) -> _Discovery:
|
||||
try:
|
||||
with _http_client(proxy) as client:
|
||||
response = client.get(_DISCOVERY_URL)
|
||||
except httpx.HTTPError as exc:
|
||||
raise XAIOAuthError(f"Could not reach xAI sign-in: {type(exc).__name__}.") from exc
|
||||
if response.status_code != HTTPStatus.OK:
|
||||
raise _oauth_http_error(response, "discovery")
|
||||
try:
|
||||
payload = response.json()
|
||||
except ValueError as exc:
|
||||
raise XAIOAuthError("xAI sign-in discovery returned invalid JSON.") from exc
|
||||
if payload.get("issuer", "").rstrip("/") != XAI_OAUTH_ISSUER:
|
||||
raise XAIOAuthError("xAI sign-in discovery returned an unexpected issuer.")
|
||||
authorization_endpoint = _validate_xai_endpoint(
|
||||
payload.get("authorization_endpoint"), "authorization"
|
||||
)
|
||||
token_endpoint = _validate_xai_endpoint(payload.get("token_endpoint"), "token")
|
||||
userinfo_raw = payload.get("userinfo_endpoint")
|
||||
userinfo_endpoint = (
|
||||
_validate_xai_endpoint(userinfo_raw, "userinfo") if userinfo_raw else None
|
||||
)
|
||||
return _Discovery(authorization_endpoint, token_endpoint, userinfo_endpoint)
|
||||
|
||||
|
||||
def _validate_xai_endpoint(value: Any, label: str) -> str:
|
||||
if not isinstance(value, str):
|
||||
raise XAIOAuthError(f"xAI sign-in discovery omitted the {label} endpoint.")
|
||||
parsed = urlsplit(value)
|
||||
try:
|
||||
port = parsed.port
|
||||
except ValueError as exc:
|
||||
raise XAIOAuthError(
|
||||
f"xAI sign-in discovery returned an unsafe {label} endpoint."
|
||||
) from exc
|
||||
if (
|
||||
parsed.scheme != "https"
|
||||
or parsed.hostname != "auth.x.ai"
|
||||
or parsed.username is not None
|
||||
or parsed.password is not None
|
||||
or port not in (None, 443)
|
||||
):
|
||||
raise XAIOAuthError(f"xAI sign-in discovery returned an unsafe {label} endpoint.")
|
||||
return value
|
||||
|
||||
|
||||
def _generate_pkce() -> tuple[str, str]:
|
||||
verifier = _base64url(secrets.token_bytes(32))
|
||||
challenge = _base64url(hashlib.sha256(verifier.encode("ascii")).digest())
|
||||
return verifier, challenge
|
||||
|
||||
|
||||
def _base64url(value: bytes) -> str:
|
||||
return base64.urlsafe_b64encode(value).rstrip(b"=").decode("ascii")
|
||||
|
||||
|
||||
def _build_authorize_url(
|
||||
endpoint: str,
|
||||
*,
|
||||
redirect_uri: str,
|
||||
challenge: str,
|
||||
state: str,
|
||||
nonce: str,
|
||||
) -> str:
|
||||
params = {
|
||||
"response_type": "code",
|
||||
"client_id": XAI_CLIENT_ID,
|
||||
"redirect_uri": redirect_uri,
|
||||
"scope": " ".join(XAI_OAUTH_SCOPES),
|
||||
"code_challenge": challenge,
|
||||
"code_challenge_method": "S256",
|
||||
"state": state,
|
||||
"nonce": nonce,
|
||||
"referrer": "nanobot",
|
||||
}
|
||||
return f"{endpoint}?{urlencode(params)}"
|
||||
|
||||
|
||||
def _make_callback_server(
|
||||
expected_state: str,
|
||||
result_queue: queue.Queue[_CallbackResult],
|
||||
) -> ThreadingHTTPServer:
|
||||
class CallbackHandler(BaseHTTPRequestHandler):
|
||||
def do_OPTIONS(self) -> None: # noqa: N802
|
||||
if self.path.split("?", 1)[0] != "/callback":
|
||||
self.send_error(HTTPStatus.NOT_FOUND)
|
||||
return
|
||||
self.send_response(HTTPStatus.NO_CONTENT)
|
||||
self._send_cors_headers()
|
||||
self.send_header("Access-Control-Allow-Methods", "GET, OPTIONS")
|
||||
self.send_header("Access-Control-Allow-Headers", "Content-Type")
|
||||
self.send_header("Access-Control-Allow-Private-Network", "true")
|
||||
self.end_headers()
|
||||
|
||||
def do_GET(self) -> None: # noqa: N802
|
||||
parsed = urlsplit(self.path)
|
||||
if parsed.path != "/callback":
|
||||
self.send_error(HTTPStatus.NOT_FOUND)
|
||||
return
|
||||
params = parse_qs(parsed.query)
|
||||
code = _first(params, "code")
|
||||
received_state = _first(params, "state")
|
||||
error = _first(params, "error_description") or _first(params, "error")
|
||||
if code and received_state and hmac.compare_digest(received_state, expected_state):
|
||||
result = _CallbackResult(code=code, state=received_state)
|
||||
title = "Signed in to xAI"
|
||||
message = "You can close this tab and return to nanobot."
|
||||
elif code:
|
||||
result = _CallbackResult(error="OAuth state mismatch")
|
||||
title = "Sign-in failed"
|
||||
message = "The sign-in response could not be verified. Return to nanobot and retry."
|
||||
else:
|
||||
result = _CallbackResult(error=error or "access denied")
|
||||
title = "Access denied"
|
||||
message = "Return to nanobot and try signing in again."
|
||||
with suppress(queue.Full):
|
||||
result_queue.put_nowait(result)
|
||||
body = _callback_page(title, message)
|
||||
encoded = body.encode("utf-8")
|
||||
self.send_response(HTTPStatus.OK)
|
||||
self._send_cors_headers()
|
||||
self.send_header("Content-Type", "text/html; charset=utf-8")
|
||||
self.send_header("Content-Length", str(len(encoded)))
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.end_headers()
|
||||
self.wfile.write(encoded)
|
||||
|
||||
def _send_cors_headers(self) -> None:
|
||||
origin = self.headers.get("Origin")
|
||||
if origin == XAI_ALLOWED_CALLBACK_ORIGIN:
|
||||
self.send_header("Access-Control-Allow-Origin", origin)
|
||||
self.send_header("Vary", "Origin")
|
||||
self.send_header("Access-Control-Allow-Private-Network", "true")
|
||||
|
||||
def log_message(self, _format: str, *_args: Any) -> None:
|
||||
# Callback query strings contain an authorization code.
|
||||
return
|
||||
|
||||
return ThreadingHTTPServer(("127.0.0.1", 0), CallbackHandler)
|
||||
|
||||
|
||||
def _serve_callback_server(
|
||||
server: ThreadingHTTPServer,
|
||||
stop_event: threading.Event,
|
||||
) -> None:
|
||||
server.timeout = 0.2
|
||||
try:
|
||||
while not stop_event.is_set():
|
||||
server.handle_request()
|
||||
finally:
|
||||
server.server_close()
|
||||
|
||||
|
||||
def _first(params: dict[str, list[str]], key: str) -> str | None:
|
||||
values = params.get(key)
|
||||
return values[0] if values else None
|
||||
|
||||
|
||||
def _callback_page(title: str, message: str) -> str:
|
||||
return f"""<!doctype html>
|
||||
<html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width">
|
||||
<title>{title}</title><style>
|
||||
body{{margin:0;min-height:100vh;display:grid;place-items:center;background:#f5f7fb;color:#172033;
|
||||
font:16px/1.5 system-ui,sans-serif}}main{{max-width:30rem;margin:1.5rem;padding:2rem;border:1px solid #dfe4ee;
|
||||
border-radius:18px;background:white;box-shadow:0 16px 50px #17203318}}h1{{margin:0 0 .65rem;font-size:1.5rem}}
|
||||
p{{margin:0;color:#526078}}</style></head><body><main><h1>{title}</h1><p>{message}</p></main></body></html>"""
|
||||
|
||||
|
||||
def _exchange_callback(
|
||||
callback: _CallbackResult,
|
||||
*,
|
||||
expected_state: str,
|
||||
discovery: _Discovery,
|
||||
verifier: str,
|
||||
redirect_uri: str,
|
||||
proxy: str | None,
|
||||
) -> XAIToken:
|
||||
if callback.error:
|
||||
raise XAIOAuthError(f"xAI sign-in was not completed: {callback.error}")
|
||||
if not callback.code:
|
||||
raise XAIOAuthError("xAI sign-in returned no authorization code.")
|
||||
if callback.state and not hmac.compare_digest(callback.state, expected_state):
|
||||
raise XAIOAuthError("xAI sign-in failed because the OAuth state did not match.")
|
||||
|
||||
payload = _exchange_code(
|
||||
discovery.token_endpoint,
|
||||
code=callback.code,
|
||||
verifier=verifier,
|
||||
redirect_uri=redirect_uri,
|
||||
proxy=proxy,
|
||||
)
|
||||
account_id = _fetch_account(discovery.userinfo_endpoint, payload["access_token"], proxy)
|
||||
return _token_from_response(payload, account_id=account_id)
|
||||
|
||||
|
||||
def _exchange_code(
|
||||
token_endpoint: str,
|
||||
*,
|
||||
code: str,
|
||||
verifier: str,
|
||||
redirect_uri: str,
|
||||
proxy: str | None,
|
||||
) -> dict[str, Any]:
|
||||
data = {
|
||||
"grant_type": "authorization_code",
|
||||
"code": code,
|
||||
"redirect_uri": redirect_uri,
|
||||
"client_id": XAI_CLIENT_ID,
|
||||
"code_verifier": verifier,
|
||||
}
|
||||
try:
|
||||
with _http_client(proxy) as client:
|
||||
response = client.post(
|
||||
token_endpoint,
|
||||
data=data,
|
||||
headers={"x-grok-client-version": XAI_CLIENT_VERSION},
|
||||
)
|
||||
except httpx.HTTPError as exc:
|
||||
raise XAIOAuthError(f"Could not exchange the xAI sign-in code: {type(exc).__name__}.") from exc
|
||||
if not response.is_success:
|
||||
raise _oauth_http_error(response, "token exchange")
|
||||
return _token_payload(response)
|
||||
|
||||
|
||||
def _refresh_token(token: XAIToken, proxy: str | None) -> XAIToken:
|
||||
data = {
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": token.refresh or "",
|
||||
"client_id": XAI_CLIENT_ID,
|
||||
}
|
||||
try:
|
||||
with _http_client(proxy) as client:
|
||||
response = client.post(
|
||||
f"{XAI_OAUTH_ISSUER}/oauth2/token",
|
||||
data=data,
|
||||
headers={"x-grok-client-version": XAI_CLIENT_VERSION},
|
||||
)
|
||||
except httpx.HTTPError as exc:
|
||||
raise XAIOAuthError(f"Could not refresh the xAI login: {type(exc).__name__}.") from exc
|
||||
if not response.is_success:
|
||||
raise _oauth_http_error(response, "token refresh")
|
||||
payload = _token_payload(response)
|
||||
return _token_from_response(
|
||||
payload,
|
||||
account_id=token.account_id,
|
||||
previous_refresh=token.refresh,
|
||||
)
|
||||
|
||||
|
||||
def _token_payload(response: httpx.Response) -> dict[str, Any]:
|
||||
try:
|
||||
payload = response.json()
|
||||
except ValueError as exc:
|
||||
raise XAIOAuthError("xAI sign-in returned an invalid token response.") from exc
|
||||
if not isinstance(payload, dict) or not isinstance(payload.get("access_token"), str):
|
||||
raise XAIOAuthError("xAI sign-in returned no access token.")
|
||||
return payload
|
||||
|
||||
|
||||
def _token_from_response(
|
||||
payload: dict[str, Any],
|
||||
*,
|
||||
account_id: str | None,
|
||||
previous_refresh: str | None = None,
|
||||
) -> XAIToken:
|
||||
try:
|
||||
expires_in = max(1, int(payload.get("expires_in") or _DEFAULT_TOKEN_TTL_S))
|
||||
except (TypeError, ValueError):
|
||||
expires_in = _DEFAULT_TOKEN_TTL_S
|
||||
refresh = payload.get("refresh_token")
|
||||
if not isinstance(refresh, str) or not refresh:
|
||||
refresh = previous_refresh
|
||||
return XAIToken(
|
||||
access=payload["access_token"],
|
||||
refresh=refresh,
|
||||
expires=_now_ms() + expires_in * 1000,
|
||||
account_id=account_id,
|
||||
)
|
||||
|
||||
|
||||
def _fetch_account(endpoint: str | None, access_token: str, proxy: str | None) -> str | None:
|
||||
if not endpoint:
|
||||
return None
|
||||
try:
|
||||
with _http_client(proxy) as client:
|
||||
response = client.get(endpoint, headers={"Authorization": f"Bearer {access_token}"})
|
||||
if not response.is_success:
|
||||
return None
|
||||
payload = response.json()
|
||||
except (httpx.HTTPError, ValueError):
|
||||
return None
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
for key in ("email", "preferred_username", "name", "sub"):
|
||||
value = payload.get(key)
|
||||
if isinstance(value, str) and value:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def _oauth_http_error(response: httpx.Response, action: str) -> XAIOAuthError:
|
||||
code: str | None = None
|
||||
description: str | None = None
|
||||
with suppress(ValueError):
|
||||
payload = response.json()
|
||||
if isinstance(payload, dict):
|
||||
raw_code = payload.get("error")
|
||||
raw_description = payload.get("error_description") or payload.get("message")
|
||||
code = raw_code[:80] if isinstance(raw_code, str) else None
|
||||
description = raw_description[:200] if isinstance(raw_description, str) else None
|
||||
detail = ": ".join(value for value in (code, description) if value)
|
||||
suffix = f" ({detail})" if detail else ""
|
||||
return XAIOAuthError(f"xAI OAuth {action} failed with HTTP {response.status_code}{suffix}.")
|
||||
|
||||
|
||||
def _http_client(proxy: str | None) -> httpx.Client:
|
||||
kwargs: dict[str, Any] = {"timeout": _HTTP_TIMEOUT_S}
|
||||
if proxy:
|
||||
kwargs.update(proxy=proxy, trust_env=False)
|
||||
return httpx.Client(**kwargs)
|
||||
|
||||
|
||||
def _token_lock() -> FileLock:
|
||||
path = get_xai_oauth_storage_path()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
return FileLock(str(path.with_suffix(".lock")), timeout=15)
|
||||
|
||||
|
||||
def _load_token() -> XAIToken | None:
|
||||
path = get_xai_oauth_storage_path()
|
||||
try:
|
||||
payload = json.loads(path.read_text(encoding="utf-8"))
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
except (OSError, ValueError, TypeError) as exc:
|
||||
logger.warning("Could not read xAI OAuth credentials: {}", type(exc).__name__)
|
||||
return None
|
||||
return XAIToken.from_dict(payload)
|
||||
|
||||
|
||||
def _write_token(token: XAIToken) -> None:
|
||||
path = get_xai_oauth_storage_path()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with suppress(OSError):
|
||||
os.chmod(path.parent, 0o700)
|
||||
_write_text_atomic(path, json.dumps(asdict(token), indent=2, ensure_ascii=False))
|
||||
with suppress(OSError):
|
||||
os.chmod(path, 0o600)
|
||||
|
||||
|
||||
def _token_is_fresh(token: XAIToken, min_ttl_ms: int) -> bool:
|
||||
return bool(token.access and token.expires > _now_ms() + max(0, min_ttl_ms))
|
||||
|
||||
|
||||
def _now_ms() -> int:
|
||||
return int(time.time() * 1000)
|
||||
Reference in New Issue
Block a user