| 1 | from __future__ import annotations |
| 2 | |
| 3 | import secrets |
| 4 | import time |
| 5 | import importlib |
| 6 | from pathlib import Path |
| 7 | from typing import Any |
| 8 | from urllib.parse import urlencode, urlparse |
| 9 | |
| 10 | from plugins._oauth.helpers.providers.base import ( |
| 11 | DUMMY_API_KEY, |
| 12 | XAI_GROK_PROVIDER_ID, |
| 13 | CallbackResult, |
| 14 | LoginPollResult, |
| 15 | LoginStartResult, |
| 16 | OAuthProviderMetadata, |
| 17 | ProviderError, |
| 18 | provider_auth_path, |
| 19 | read_json_file, |
| 20 | write_private_json, |
| 21 | ) |
| 22 | from plugins._oauth.helpers.providers.common import ( |
| 23 | as_int as _as_int, |
| 24 | as_optional_string as _as_optional_string, |
| 25 | error_message as _error_message, |
| 26 | expires_ms as _expires_ms, |
| 27 | json_payload as _json_payload, |
| 28 | latest_attempt, |
| 29 | models_from_payload as _models_from_payload, |
| 30 | parse_manual_callback, |
| 31 | ) |
| 32 | from plugins._oauth.helpers.state import get_attempt, pop_attempt, put_attempt |
| 33 | |
| 34 | |
| 35 | XAI_ISSUER = "https://auth.x.ai" |
| 36 | XAI_DISCOVERY_URL = f"{XAI_ISSUER}/.well-known/openid-configuration" |
| 37 | XAI_CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828" |
| 38 | XAI_SCOPE = "openid profile email offline_access grok-cli:access api:access" |
| 39 | XAI_REDIRECT_URI = "http://127.0.0.1:56121/callback" |
| 40 | XAI_API_BASE = "https://api.x.ai/v1" |
| 41 | CURATED_MODELS = [ |
| 42 | "grok-4.3", |
| 43 | "grok-4.20-0309-reasoning", |
| 44 | "grok-4.20-0309-non-reasoning", |
| 45 | "grok-4.20-multi-agent-0309", |
| 46 | "grok-code-fast-1", |
| 47 | ] |
| 48 | NOT_CONNECTED_MESSAGE = "xAI Grok OAuth is not connected yet." |
| 49 | OAUTH_TIER_WARNING = ( |
| 50 | "xAI Grok OAuth API access may be restricted by tier. " |
| 51 | "If OAuth token exchange is denied, the separate API-key `xai` provider may work." |
| 52 | ) |
| 53 | REFRESH_MARGIN_MS = 60_000 |
| 54 | |
| 55 | |
| 56 | class XaiGrokOAuthProvider: |
| 57 | provider_id = XAI_GROK_PROVIDER_ID |
| 58 | |
| 59 | def auth_path(self) -> Path: |
| 60 | return provider_auth_path("xai_grok") |
| 61 | |
| 62 | def read_auth(self) -> dict[str, Any]: |
| 63 | return read_json_file(self.auth_path()) |
| 64 | |
| 65 | def write_auth(self, data: dict[str, Any]) -> None: |
| 66 | write_private_json(self.auth_path(), data) |
| 67 | |
| 68 | def discovery(self) -> dict[str, str]: |
| 69 | import requests |
| 70 | |
| 71 | response = requests.get( |
| 72 | XAI_DISCOVERY_URL, |
| 73 | headers={"Accept": "application/json"}, |
| 74 | timeout=30, |
| 75 | ) |
| 76 | payload = _json_payload(response) |
| 77 | if not response.ok: |
| 78 | raise ProviderError( |
| 79 | _error_message(payload, f"xAI Grok discovery failed with status {response.status_code}."), |
| 80 | code="discovery_failed", |
| 81 | status=response.status_code, |
| 82 | ) |
| 83 | authorization_endpoint = str(payload.get("authorization_endpoint") or "").strip() |
| 84 | token_endpoint = str(payload.get("token_endpoint") or "").strip() |
| 85 | if not authorization_endpoint or not token_endpoint: |
| 86 | raise ProviderError( |
| 87 | "xAI Grok discovery response was missing OAuth endpoints.", |
| 88 | code="discovery_malformed", |
| 89 | status=502, |
| 90 | ) |
| 91 | _validate_xai_endpoint(authorization_endpoint) |
| 92 | _validate_xai_endpoint(token_endpoint) |
| 93 | return { |
| 94 | "authorization_endpoint": authorization_endpoint, |
| 95 | "token_endpoint": token_endpoint, |
| 96 | } |
| 97 | |
| 98 | def metadata(self) -> OAuthProviderMetadata: |
| 99 | return OAuthProviderMetadata( |
| 100 | provider_id=XAI_GROK_PROVIDER_ID, |
| 101 | display_name="xAI Grok", |
| 102 | short_name="Grok", |
| 103 | model_provider_id=XAI_GROK_PROVIDER_ID, |
| 104 | icon="xai", |
| 105 | auth_flow="browser_pkce", |
| 106 | default_model="grok-4.3", |
| 107 | default_models=list(CURATED_MODELS), |
| 108 | proxy_base_path="/oauth/xai-grok", |
| 109 | callback_path="/oauth/xai-grok/callback", |
| 110 | supports_manual_callback=True, |
| 111 | warning=OAUTH_TIER_WARNING, |
| 112 | ) |
| 113 | |
| 114 | def status(self) -> dict[str, Any]: |
| 115 | auth = self.read_auth() |
| 116 | access = str(auth.get("access") or "") |
| 117 | refresh = str(auth.get("refresh") or "") |
| 118 | result = { |
| 119 | **self.metadata().to_dict(), |
| 120 | "connected": bool(access and refresh), |
| 121 | "account_label": "xAI Grok" if access or refresh else "", |
| 122 | "base_url": str(auth.get("base_url") or XAI_API_BASE), |
| 123 | "auth_file_path": str(self.auth_path()), |
| 124 | } |
| 125 | warning = str(auth.get("warning") or "") |
| 126 | if warning: |
| 127 | result["warning"] = warning |
| 128 | elif access and refresh and _as_int(auth.get("expires"), 0) <= int(time.time() * 1000): |
| 129 | result["warning"] = "xAI Grok OAuth access token is expired and will be refreshed on the next request." |
| 130 | return result |
| 131 | |
| 132 | def start_login(self, input: dict[str, Any] | None = None, request: Any = None) -> LoginStartResult: |
| 133 | del input, request |
| 134 | try: |
| 135 | codex = importlib.import_module("plugins._oauth.helpers.codex") |
| 136 | metadata = self.discovery() |
| 137 | pkce = codex.generate_pkce() |
| 138 | state = codex.generate_state() |
| 139 | nonce = secrets.token_urlsafe(24) |
| 140 | attempt = put_attempt( |
| 141 | state, |
| 142 | pkce.verifier, |
| 143 | XAI_REDIRECT_URI, |
| 144 | provider_id=XAI_GROK_PROVIDER_ID, |
| 145 | extra={ |
| 146 | "nonce": nonce, |
| 147 | "code_challenge": pkce.challenge, |
| 148 | "token_endpoint": metadata["token_endpoint"], |
| 149 | }, |
| 150 | ) |
| 151 | query = { |
| 152 | "response_type": "code", |
| 153 | "client_id": XAI_CLIENT_ID, |
| 154 | "redirect_uri": XAI_REDIRECT_URI, |
| 155 | "scope": XAI_SCOPE, |
| 156 | "code_challenge": pkce.challenge, |
| 157 | "code_challenge_method": "S256", |
| 158 | "state": state, |
| 159 | "nonce": nonce, |
| 160 | "plan": "generic", |
| 161 | "referrer": "agent-zero", |
| 162 | } |
| 163 | auth_url = f'{metadata["authorization_endpoint"]}?{urlencode(query)}' |
| 164 | except Exception as exc: |
| 165 | return LoginStartResult( |
| 166 | ok=False, |
| 167 | provider_id=XAI_GROK_PROVIDER_ID, |
| 168 | flow="browser_pkce", |
| 169 | error=str(exc), |
| 170 | message=str(exc), |
| 171 | ) |
| 172 | |
| 173 | return LoginStartResult( |
| 174 | ok=True, |
| 175 | provider_id=XAI_GROK_PROVIDER_ID, |
| 176 | flow="browser_pkce", |
| 177 | auth_url=auth_url, |
| 178 | redirect_uri=XAI_REDIRECT_URI, |
| 179 | expires_at=attempt.expires_at, |
| 180 | ) |
| 181 | |
| 182 | def poll_login(self, input: dict[str, Any] | None = None, request: Any = None) -> LoginPollResult: |
| 183 | del input, request |
| 184 | return LoginPollResult( |
| 185 | ok=False, |
| 186 | provider_id=XAI_GROK_PROVIDER_ID, |
| 187 | error="xAI Grok uses browser callback login.", |
| 188 | ) |
| 189 | |
| 190 | def exchange_code( |
| 191 | self, |
| 192 | token_endpoint: str, |
| 193 | code: str, |
| 194 | redirect_uri: str, |
| 195 | code_verifier: str, |
| 196 | code_challenge: str, |
| 197 | ) -> dict[str, Any]: |
| 198 | import requests |
| 199 | |
| 200 | response = requests.post( |
| 201 | token_endpoint, |
| 202 | headers={ |
| 203 | "Accept": "application/json", |
| 204 | "Content-Type": "application/x-www-form-urlencoded", |
| 205 | }, |
| 206 | data={ |
| 207 | "grant_type": "authorization_code", |
| 208 | "code": code, |
| 209 | "redirect_uri": redirect_uri, |
| 210 | "client_id": XAI_CLIENT_ID, |
| 211 | "code_verifier": code_verifier, |
| 212 | "code_challenge": code_challenge, |
| 213 | "code_challenge_method": "S256", |
| 214 | }, |
| 215 | timeout=30, |
| 216 | ) |
| 217 | payload = _json_payload(response) |
| 218 | if not response.ok: |
| 219 | if response.status_code == 403: |
| 220 | raise ProviderError( |
| 221 | OAUTH_TIER_WARNING, |
| 222 | code="oauth_tier_restricted", |
| 223 | status=403, |
| 224 | ) |
| 225 | raise ProviderError( |
| 226 | _error_message(payload, f"xAI Grok token exchange failed with status {response.status_code}."), |
| 227 | code="token_exchange_failed", |
| 228 | status=response.status_code, |
| 229 | ) |
| 230 | _validate_token_payload(payload, require_refresh=True) |
| 231 | return payload |
| 232 | |
| 233 | def manual_callback(self, input: dict[str, Any], request: Any = None) -> LoginPollResult: |
| 234 | del request |
| 235 | raw = input.get("callback") |
| 236 | if raw is None: |
| 237 | raw = input.get("callback_url") |
| 238 | return self._complete_from_callback(parse_manual_callback(raw), allow_missing_state=True) |
| 239 | |
| 240 | def complete_callback( |
| 241 | self, |
| 242 | args: dict[str, Any], |
| 243 | request: Any = None, |
| 244 | ) -> CallbackResult: |
| 245 | del request |
| 246 | callback = { |
| 247 | "code": _as_optional_string(args.get("code")), |
| 248 | "state": _as_optional_string(args.get("state")), |
| 249 | "error": _as_optional_string(args.get("error")), |
| 250 | "error_description": _as_optional_string(args.get("error_description")), |
| 251 | } |
| 252 | result = self._complete_from_callback(callback, allow_missing_state=False) |
| 253 | return CallbackResult( |
| 254 | ok=result.ok, |
| 255 | provider_id=XAI_GROK_PROVIDER_ID, |
| 256 | account_label=result.account_label, |
| 257 | error=result.error, |
| 258 | ) |
| 259 | |
| 260 | def _complete_from_callback( |
| 261 | self, |
| 262 | callback: dict[str, str | None] | None, |
| 263 | *, |
| 264 | allow_missing_state: bool, |
| 265 | ) -> LoginPollResult: |
| 266 | if not callback: |
| 267 | return LoginPollResult(ok=False, provider_id=XAI_GROK_PROVIDER_ID, error="Missing OAuth callback.") |
| 268 | if callback.get("error"): |
| 269 | return LoginPollResult( |
| 270 | ok=False, |
| 271 | provider_id=XAI_GROK_PROVIDER_ID, |
| 272 | error=str(callback.get("error_description") or callback.get("error")), |
| 273 | ) |
| 274 | |
| 275 | code = str(callback.get("code") or "").strip() |
| 276 | if not code: |
| 277 | return LoginPollResult( |
| 278 | ok=False, |
| 279 | provider_id=XAI_GROK_PROVIDER_ID, |
| 280 | error="The OAuth callback did not include an authorization code.", |
| 281 | ) |
| 282 | |
| 283 | state = str(callback.get("state") or "").strip() |
| 284 | attempt = None |
| 285 | if state: |
| 286 | attempt = get_attempt(state) |
| 287 | if attempt is None: |
| 288 | if latest_attempt(XAI_GROK_PROVIDER_ID) is not None: |
| 289 | return LoginPollResult( |
| 290 | ok=False, |
| 291 | provider_id=XAI_GROK_PROVIDER_ID, |
| 292 | error="OAuth state mismatch. Return to Agent Zero and start a new xAI Grok connection.", |
| 293 | ) |
| 294 | return LoginPollResult( |
| 295 | ok=False, |
| 296 | provider_id=XAI_GROK_PROVIDER_ID, |
| 297 | expired=True, |
| 298 | error="OAuth sign-in expired. Return to Agent Zero and start a new xAI Grok connection.", |
| 299 | ) |
| 300 | if attempt.provider_id != XAI_GROK_PROVIDER_ID: |
| 301 | return LoginPollResult( |
| 302 | ok=False, |
| 303 | provider_id=XAI_GROK_PROVIDER_ID, |
| 304 | error="OAuth state mismatch. Return to Agent Zero and start a new xAI Grok connection.", |
| 305 | ) |
| 306 | elif allow_missing_state: |
| 307 | attempt = latest_attempt(XAI_GROK_PROVIDER_ID) |
| 308 | if attempt is None: |
| 309 | return LoginPollResult( |
| 310 | ok=False, |
| 311 | provider_id=XAI_GROK_PROVIDER_ID, |
| 312 | error="No active xAI Grok sign-in attempt was found.", |
| 313 | ) |
| 314 | state = attempt.state |
| 315 | else: |
| 316 | return LoginPollResult( |
| 317 | ok=False, |
| 318 | provider_id=XAI_GROK_PROVIDER_ID, |
| 319 | error="The OAuth callback did not include state.", |
| 320 | ) |
| 321 | |
| 322 | token_endpoint = str(attempt.extra.get("token_endpoint") or "") |
| 323 | if not token_endpoint: |
| 324 | token_endpoint = self.discovery()["token_endpoint"] |
| 325 | code_challenge = str(attempt.extra.get("code_challenge") or "") |
| 326 | try: |
| 327 | payload = self.exchange_code( |
| 328 | token_endpoint, |
| 329 | code, |
| 330 | attempt.redirect_uri, |
| 331 | attempt.verifier, |
| 332 | code_challenge, |
| 333 | ) |
| 334 | auth = _auth_from_token_payload(payload, token_endpoint) |
| 335 | self.write_auth(auth) |
| 336 | pop_attempt(state) |
| 337 | except Exception as exc: |
| 338 | return LoginPollResult( |
| 339 | ok=False, |
| 340 | provider_id=XAI_GROK_PROVIDER_ID, |
| 341 | error=str(exc), |
| 342 | ) |
| 343 | |
| 344 | return LoginPollResult( |
| 345 | ok=True, |
| 346 | provider_id=XAI_GROK_PROVIDER_ID, |
| 347 | completed=True, |
| 348 | account_label="xAI Grok", |
| 349 | ) |
| 350 | |
| 351 | def ensure_fresh_auth(self) -> dict[str, Any]: |
| 352 | auth = self.read_auth() |
| 353 | access = str(auth.get("access") or "") |
| 354 | refresh = str(auth.get("refresh") or "") |
| 355 | if not access or not refresh: |
| 356 | return auth |
| 357 | |
| 358 | expires = _as_int(auth.get("expires"), 0) |
| 359 | if expires and expires > int(time.time() * 1000) + REFRESH_MARGIN_MS: |
| 360 | return auth |
| 361 | |
| 362 | token_endpoint = str(auth.get("token_endpoint") or "") |
| 363 | if not token_endpoint: |
| 364 | token_endpoint = self.discovery()["token_endpoint"] |
| 365 | _validate_xai_endpoint(token_endpoint, code="invalid_token_endpoint") |
| 366 | try: |
| 367 | refreshed = self._refresh_tokens(token_endpoint, refresh, auth) |
| 368 | except ProviderError: |
| 369 | raise |
| 370 | except Exception as exc: |
| 371 | raise ProviderError( |
| 372 | f"xAI Grok OAuth refresh failed: {exc}", |
| 373 | code="auth_refresh_failed", |
| 374 | status=401, |
| 375 | ) from exc |
| 376 | self.write_auth(refreshed) |
| 377 | return refreshed |
| 378 | |
| 379 | def _refresh_tokens(self, token_endpoint: str, refresh: str, existing: dict[str, Any]) -> dict[str, Any]: |
| 380 | import requests |
| 381 | |
| 382 | response = requests.post( |
| 383 | token_endpoint, |
| 384 | headers={ |
| 385 | "Accept": "application/json", |
| 386 | "Content-Type": "application/x-www-form-urlencoded", |
| 387 | }, |
| 388 | data={ |
| 389 | "grant_type": "refresh_token", |
| 390 | "refresh_token": refresh, |
| 391 | "client_id": XAI_CLIENT_ID, |
| 392 | }, |
| 393 | timeout=30, |
| 394 | ) |
| 395 | payload = _json_payload(response) |
| 396 | if not response.ok: |
| 397 | if response.status_code == 403: |
| 398 | raise ProviderError( |
| 399 | OAUTH_TIER_WARNING, |
| 400 | code="oauth_tier_restricted", |
| 401 | status=403, |
| 402 | ) |
| 403 | raise ProviderError( |
| 404 | _error_message(payload, f"xAI Grok token refresh failed with status {response.status_code}."), |
| 405 | code="token_refresh_failed", |
| 406 | status=response.status_code, |
| 407 | ) |
| 408 | _validate_token_payload(payload, require_refresh=False) |
| 409 | merged = dict(existing) |
| 410 | merged.update(_auth_from_token_payload(payload, token_endpoint, fallback_refresh=refresh)) |
| 411 | if not payload.get("id_token") and existing.get("id_token"): |
| 412 | merged["id_token"] = existing["id_token"] |
| 413 | if not payload.get("token_type") and existing.get("token_type"): |
| 414 | merged["token_type"] = existing["token_type"] |
| 415 | return merged |
| 416 | |
| 417 | def models(self) -> list[str]: |
| 418 | if not self.read_auth(): |
| 419 | return list(CURATED_MODELS) |
| 420 | try: |
| 421 | auth = self.ensure_fresh_auth() |
| 422 | except Exception: |
| 423 | return list(CURATED_MODELS) |
| 424 | access = str(auth.get("access") or "") |
| 425 | if not access: |
| 426 | return list(CURATED_MODELS) |
| 427 | |
| 428 | base_url = safe_api_base_url(auth.get("base_url")) |
| 429 | try: |
| 430 | import requests |
| 431 | |
| 432 | response = requests.get( |
| 433 | f"{base_url}/models", |
| 434 | headers={ |
| 435 | "Accept": "application/json", |
| 436 | "Authorization": f"Bearer {access}", |
| 437 | }, |
| 438 | timeout=30, |
| 439 | ) |
| 440 | if not response.ok: |
| 441 | return list(CURATED_MODELS) |
| 442 | parsed = _models_from_payload(response.json()) |
| 443 | return parsed or list(CURATED_MODELS) |
| 444 | except Exception: |
| 445 | return list(CURATED_MODELS) |
| 446 | |
| 447 | def disconnect(self) -> dict[str, Any]: |
| 448 | path = self.auth_path() |
| 449 | existed = path.exists() |
| 450 | try: |
| 451 | path.unlink(missing_ok=True) |
| 452 | except FileNotFoundError: |
| 453 | pass |
| 454 | return { |
| 455 | "disconnected": existed, |
| 456 | "removed_auth_files": [str(path)] if existed else [], |
| 457 | } |
| 458 | |
| 459 | def api_key(self) -> str: |
| 460 | return DUMMY_API_KEY |
| 461 | |
| 462 | def register_routes(self, app: Any) -> None: |
| 463 | from plugins._oauth.helpers import routes |
| 464 | |
| 465 | route_defs = [ |
| 466 | ("/oauth/xai-grok/health", "oauth_xai_grok_health", routes.xai_grok_health, ["GET"]), |
| 467 | ("/oauth/xai-grok/callback", "oauth_xai_grok_callback", routes.xai_grok_callback, ["GET"]), |
| 468 | ( |
| 469 | "/oauth/xai-grok/v1/models", |
| 470 | "oauth_xai_grok_models", |
| 471 | routes.xai_grok_models, |
| 472 | ["GET", "OPTIONS"], |
| 473 | ), |
| 474 | ( |
| 475 | "/oauth/xai-grok/v1/chat/completions", |
| 476 | "oauth_xai_grok_chat_completions", |
| 477 | routes.xai_grok_chat_completions, |
| 478 | ["POST", "OPTIONS"], |
| 479 | ), |
| 480 | ( |
| 481 | "/oauth/xai-grok/v1/responses", |
| 482 | "oauth_xai_grok_responses", |
| 483 | routes.xai_grok_responses, |
| 484 | ["POST", "OPTIONS"], |
| 485 | ), |
| 486 | ] |
| 487 | for rule, endpoint, view_func, methods in route_defs: |
| 488 | if endpoint in app.view_functions: |
| 489 | continue |
| 490 | app.add_url_rule(rule, endpoint, view_func, methods=methods) |
| 491 | |
| 492 | |
| 493 | def _auth_from_token_payload( |
| 494 | payload: dict[str, Any], |
| 495 | token_endpoint: str, |
| 496 | *, |
| 497 | fallback_refresh: str = "", |
| 498 | ) -> dict[str, Any]: |
| 499 | return { |
| 500 | "provider": XAI_GROK_PROVIDER_ID, |
| 501 | "type": "oauth", |
| 502 | "access": str(payload.get("access_token") or ""), |
| 503 | "refresh": str(payload.get("refresh_token") or fallback_refresh or ""), |
| 504 | "expires": _expires_ms(payload), |
| 505 | "id_token": str(payload.get("id_token") or ""), |
| 506 | "token_type": str(payload.get("token_type") or "Bearer"), |
| 507 | "token_endpoint": token_endpoint, |
| 508 | "base_url": XAI_API_BASE, |
| 509 | } |
| 510 | |
| 511 | |
| 512 | def _validate_token_payload(payload: dict[str, Any], *, require_refresh: bool) -> None: |
| 513 | if not isinstance(payload, dict): |
| 514 | raise ProviderError("xAI Grok token endpoint returned a malformed response.", code="token_malformed", status=502) |
| 515 | missing = [] |
| 516 | if not str(payload.get("access_token") or ""): |
| 517 | missing.append("access_token") |
| 518 | if require_refresh and not str(payload.get("refresh_token") or ""): |
| 519 | missing.append("refresh_token") |
| 520 | if missing: |
| 521 | raise ProviderError( |
| 522 | f"xAI Grok token response is missing: {', '.join(missing)}", |
| 523 | code="token_malformed", |
| 524 | status=502, |
| 525 | ) |
| 526 | |
| 527 | |
| 528 | def safe_api_base_url(value: Any) -> str: |
| 529 | text = str(value or "").strip().rstrip("/") |
| 530 | if not text: |
| 531 | return XAI_API_BASE |
| 532 | parsed = urlparse(text) |
| 533 | host = (parsed.hostname or "").lower() |
| 534 | if parsed.scheme == "https" and (host == "api.x.ai" or host.endswith(".api.x.ai")): |
| 535 | return text |
| 536 | return XAI_API_BASE |
| 537 | |
| 538 | |
| 539 | def _validate_xai_endpoint(value: str, *, code: str = "discovery_invalid_endpoint") -> None: |
| 540 | parsed = urlparse(value) |
| 541 | host = (parsed.hostname or "").lower() |
| 542 | if parsed.scheme != "https" or not (host == "x.ai" or host.endswith(".x.ai")): |
| 543 | raise ProviderError( |
| 544 | "xAI Grok discovery returned an invalid OAuth endpoint.", |
| 545 | code=code, |
| 546 | status=502, |
| 547 | ) |