diff --git a/pyproject.toml b/pyproject.toml index c9ea8bc1..40828692 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -217,6 +217,7 @@ ignore = [ # args they intentionally ignore. "tests/test_viewer_auth.py" = ["S105", "S106", "ARG001"] "tests/test_codex_auth.py" = ["S105", "S106", "SLF001"] +"tests/test_grok_auth.py" = ["S105", "S106", "SLF001"] # Stdlib HTTP handler overrides (do_GET/do_POST). "strix/interface/auth_cli.py" = ["N802"] "tests/test_codex_streaming.py" = ["N802"] @@ -268,6 +269,7 @@ ignore = [ # Heavy inference deps (httpx, openai) imported lazily so auth-status checks # don't pull them in. "strix/config/codex.py" = ["PLC0415"] +"strix/config/grok.py" = ["PLC0415"] # Interface utility branches per scope-mode / target-type combination; # splitting would obscure the decision tree without simplifying it. "strix/interface/utils.py" = ["PLR0912", "BLE001", "PLC0415"] diff --git a/strix/config/grok.py b/strix/config/grok.py new file mode 100644 index 00000000..2350bc7d --- /dev/null +++ b/strix/config/grok.py @@ -0,0 +1,344 @@ +"""Grok (xAI) subscription auth: OAuth login, token refresh, and the OpenAI +client that routes inference through xAI's API. + +Mirrors xAI's Grok CLI: OAuth 2.0 + PKCE against ``auth.x.ai``, with the access +token sent as a ``Bearer`` token to ``api.x.ai/v1`` (OpenAI-compatible, so the +subscription and a metered API key share one endpoint — only the bearer differs). +Using a Grok/SuperGrok subscription outside xAI's own products is not officially +supported by xAI; the user chooses this path knowingly. The OAuth constants are +xAI's own Grok CLI values (the backend only accepts that client). +""" + +from __future__ import annotations + +import base64 +import contextlib +import hashlib +import json +import logging +import secrets +import threading +import time +import urllib.parse +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import requests + + +if TYPE_CHECKING: + from collections.abc import Iterator + + from openai import AsyncOpenAI + + +logger = logging.getLogger(__name__) + + +PROVIDER = "grok" + +CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828" +AUTHORIZE_URL = "https://auth.x.ai/oauth2/authorize" +TOKEN_URL = "https://auth.x.ai/oauth2/token" # noqa: S105 # nosec B105 - URL, not a secret +CALLBACK_HOST = "127.0.0.1" +CALLBACK_PORT = 56121 +CALLBACK_PATH = "/callback" +REDIRECT_URI = f"http://{CALLBACK_HOST}:{CALLBACK_PORT}{CALLBACK_PATH}" +SCOPE = "openid profile email offline_access grok-cli:access api:access" + +XAI_BASE_URL = "https://api.x.ai/v1" + +_TOKEN_TIMEOUT = 30 +_EXPIRY_SKEW_S = 300 + +_refresh_lock = threading.Lock() + +# Shared with the other subscription providers; kept separate from cli-config.json +# so OAuth tokens never land in the env-var config. +AUTH_PATH = Path.home() / ".strix" / "subscription-auth.json" + + +def _read_store() -> dict[str, Any]: + try: + data = json.loads(AUTH_PATH.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return {} + return data if isinstance(data, dict) else {} + + +def _write_store(data: dict[str, Any]) -> None: + AUTH_PATH.parent.mkdir(parents=True, exist_ok=True) + tmp = AUTH_PATH.with_suffix(".json.tmp") + tmp.write_text(json.dumps(data, indent=2), encoding="utf-8") + with contextlib.suppress(OSError): + tmp.chmod(0o600) + tmp.replace(AUTH_PATH) + with contextlib.suppress(OSError): + AUTH_PATH.chmod(0o600) + + +def read_record() -> dict[str, Any] | None: + record = _read_store().get(PROVIDER) + if not isinstance(record, dict) or record.get("type") != "oauth": + return None + if not (record.get("access") and record.get("refresh")): + return None + return record + + +def is_authenticated() -> bool: + return read_record() is not None + + +def save_record(record: dict[str, Any]) -> None: + data = _read_store() + data[PROVIDER] = record + _write_store(data) + + +def logout() -> None: + data = _read_store() + if PROVIDER not in data: + return + del data[PROVIDER] + if data: + _write_store(data) + return + with contextlib.suppress(OSError): + AUTH_PATH.unlink() + + +@contextlib.contextmanager +def _refresh_guard() -> Iterator[None]: + """Serialize token refresh within (lock) and across (flock) Strix processes, + so concurrent runs can't both spend the single-use refresh token.""" + with _refresh_lock: + try: + import fcntl + + lock_path = AUTH_PATH.with_suffix(".lock") + lock_path.parent.mkdir(parents=True, exist_ok=True) + handle = lock_path.open("w") + except (ImportError, OSError): + yield + return + try: + with contextlib.suppress(OSError): + fcntl.flock(handle.fileno(), fcntl.LOCK_EX) + yield + finally: + with contextlib.suppress(OSError): + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + handle.close() + + +class GrokAuthError(Exception): + def __init__(self, code: str, message: str | None = None) -> None: + self.code = code + super().__init__(message or code) + + +def _b64url(raw: bytes) -> str: + return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii") + + +def generate_pkce() -> tuple[str, str]: + verifier = _b64url(secrets.token_bytes(64)) + challenge = _b64url(hashlib.sha256(verifier.encode("ascii")).digest()) + return verifier, challenge + + +def create_state() -> str: + return secrets.token_hex(16) + + +def build_authorize_url(challenge: str, state: str) -> str: + params = { + "response_type": "code", + "client_id": CLIENT_ID, + "redirect_uri": REDIRECT_URI, + "scope": SCOPE, + "code_challenge": challenge, + "code_challenge_method": "S256", + "state": state, + } + return f"{AUTHORIZE_URL}?{urllib.parse.urlencode(params)}" + + +def parse_redirect_input(value: str) -> tuple[str | None, str | None]: + """Extract ``(code, state)`` from a pasted redirect URL, ``code#state``, + query string, or bare code.""" + value = (value or "").strip() + if not value: + return None, None + with contextlib.suppress(ValueError): + parsed = urllib.parse.urlparse(value) + if parsed.scheme and parsed.query: + query = urllib.parse.parse_qs(parsed.query) + return _first(query, "code"), _first(query, "state") + if "#" in value: + code, _, state = value.partition("#") + return code or None, state or None + if "code=" in value: + query = urllib.parse.parse_qs(value) + return _first(query, "code"), _first(query, "state") + return value, None + + +def _first(query: dict[str, list[str]], key: str) -> str | None: + values = query.get(key) + return values[0] if values else None + + +def _post_form(payload: dict[str, str]) -> dict[str, Any]: + try: + response = requests.post( + TOKEN_URL, + data=payload, + headers={"Accept": "application/json"}, + timeout=_TOKEN_TIMEOUT, + ) + except requests.RequestException as exc: + raise GrokAuthError("unavailable", str(exc)) from exc + if response.status_code >= 400: + detail = response.text[:300] + raise GrokAuthError("token_http_error", f"HTTP {response.status_code}: {detail}") + data = json.loads(response.content or b"{}") + if not isinstance(data, dict): + raise GrokAuthError("bad_response", "token endpoint returned non-object") + return data + + +def _record_from_token_response( + data: dict[str, Any], refresh_fallback: str | None = None +) -> dict[str, Any]: + access = data.get("access_token") + # A refresh response may omit refresh_token when it isn't rotated; keep the old one. + refresh = data.get("refresh_token") or refresh_fallback + expires_in = data.get("expires_in") + if not isinstance(access, str) or not access: + raise GrokAuthError("bad_response", "token response missing access_token") + if not isinstance(refresh, str) or not refresh: + raise GrokAuthError("bad_response", "token response missing refresh_token") + ttl = expires_in if isinstance(expires_in, int | float) else 3600 + return { + "type": "oauth", + "provider": PROVIDER, + "access": access, + "refresh": refresh, + "expires_at": time.time() + ttl, + } + + +def exchange_code(code: str, verifier: str) -> dict[str, Any]: + data = _post_form( + { + "grant_type": "authorization_code", + "client_id": CLIENT_ID, + "code": code, + "code_verifier": verifier, + "redirect_uri": REDIRECT_URI, + } + ) + return _record_from_token_response(data) + + +def refresh_tokens(refresh_token: str) -> dict[str, Any]: + data = _post_form( + { + "grant_type": "refresh_token", + "client_id": CLIENT_ID, + "refresh_token": refresh_token, + } + ) + return _record_from_token_response(data, refresh_fallback=refresh_token) + + +def _access_token(record: dict[str, Any]) -> str: + access = record["access"] + if not isinstance(access, str) or not access: + raise GrokAuthError("bad_response", "stored access token is missing or malformed") + return access + + +def _near_expiry(record: dict[str, Any]) -> bool: + expires_at = record.get("expires_at") + if not isinstance(expires_at, int | float): + return True + return expires_at - _EXPIRY_SKEW_S <= time.time() + + +def get_valid_token() -> str: + """Return a valid access token, refreshing under the cross-process guard if + near expiry.""" + record = read_record() + if record is None: + raise GrokAuthError("not_authenticated", "not signed in; run: strix auth login grok") + if not _near_expiry(record): + return _access_token(record) + with _refresh_guard(): + record = read_record() + if record is None: + raise GrokAuthError("not_authenticated", "not signed in; run: strix auth login grok") + if not _near_expiry(record): + return _access_token(record) + try: + refreshed = refresh_tokens(record["refresh"]) + except GrokAuthError: + # A peer process may have already spent this single-use refresh token. + latest = read_record() + if latest and latest["refresh"] != record["refresh"] and not _near_expiry(latest): + return _access_token(latest) + raise + save_record(refreshed) + return _access_token(refreshed) + + +def build_openai_client() -> AsyncOpenAI: + """An ``AsyncOpenAI`` for xAI's API. A per-request hook re-stamps a fresh + bearer token so long scans survive token expiry.""" + import asyncio + + import httpx + from openai import AsyncOpenAI + + get_valid_token() # fail fast at configure time if the sign-in is dead + + async def _auth_hook(request: httpx.Request) -> None: + access = await asyncio.to_thread(get_valid_token) + request.headers["Authorization"] = f"Bearer {access}" + + http_client = httpx.AsyncClient( + timeout=httpx.Timeout(600.0, connect=30.0), + event_hooks={"request": [_auth_hook]}, + ) + return AsyncOpenAI( + api_key="strix-grok-oauth", # placeholder; the hook overwrites Authorization + base_url=XAI_BASE_URL, + http_client=http_client, + ) + + +_subscription_client: AsyncOpenAI | None = None + + +def get_subscription_client() -> AsyncOpenAI: + global _subscription_client # noqa: PLW0603 + if _subscription_client is None: + _subscription_client = build_openai_client() + return _subscription_client + + +SUBSCRIPTION_PREFIX = "grok/" + + +def subscription_model(model_name: str | None) -> str | None: + """The model slug behind a ``grok/`` STRIX_LLM, or None.""" + name = (model_name or "").strip() + if not name.lower().startswith(SUBSCRIPTION_PREFIX): + return None + return name[len(SUBSCRIPTION_PREFIX) :] or None + + +def auth_mode(model_name: str | None) -> str: + return "subscription" if subscription_model(model_name) else "api_key" diff --git a/strix/config/models.py b/strix/config/models.py index b6e6b0f0..ee8d0d68 100644 --- a/strix/config/models.py +++ b/strix/config/models.py @@ -14,6 +14,7 @@ from agents import ( ) from agents.model_settings import ModelSettings from agents.models.multi_provider import MultiProvider +from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel from agents.models.openai_responses import OpenAIResponsesModel from agents.retry import ( ModelRetryBackoffSettings, @@ -23,7 +24,7 @@ from agents.retry import ( ) from openai.types.shared import Reasoning -from strix.config import codex +from strix.config import codex, grok from strix.config.loader import load_settings @@ -166,6 +167,11 @@ class StrixProvider(MultiProvider): codex.get_subscription_client(), reasoning_effort=load_settings().llm.reasoning_effort, ) + grok_slug = grok.subscription_model(model_name) + if grok_slug: + # xAI's API is OpenAI chat-completions compatible; the subscription + # bearer is stamped per-request by the client's auth hook. + return OpenAIChatCompletionsModel(grok_slug, grok.get_subscription_client()) return super().get_model(model_name) @@ -229,7 +235,7 @@ def configure_sdk_model_defaults(settings: Settings) -> None: """Apply Strix config to SDK-native defaults.""" llm = settings.llm set_tracing_disabled(True) - if codex.subscription_model(llm.model): + if codex.subscription_model(llm.model) or grok.subscription_model(llm.model): return _configure_litellm_compatibility() _configure_openrouter_attribution(llm.model) diff --git a/strix/config/subscription.py b/strix/config/subscription.py new file mode 100644 index 00000000..7cd578ca --- /dev/null +++ b/strix/config/subscription.py @@ -0,0 +1,33 @@ +"""Shared helpers across model-subscription providers (ChatGPT/Codex and Grok). + +Each provider module (:mod:`strix.config.codex`, :mod:`strix.config.grok`) +exposes the same small surface — ``subscription_model``, ``auth_mode``, +``is_authenticated`` — so callers that only care "is this run on a subscription, +and which provider?" can stay provider-agnostic. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from strix.config import codex, grok + + +if TYPE_CHECKING: + from types import ModuleType + + +_PROVIDERS: tuple[ModuleType, ...] = (codex, grok) + + +def provider_for_model(model_name: str | None) -> ModuleType | None: + """Return the subscription provider module that owns ``model_name``'s prefix, + or None when the model isn't a subscription model.""" + for provider in _PROVIDERS: + if provider.subscription_model(model_name): + return provider + return None + + +def auth_mode(model_name: str | None) -> str: + return "subscription" if provider_for_model(model_name) is not None else "api_key" diff --git a/strix/interface/auth_cli.py b/strix/interface/auth_cli.py index 51e6b9fe..90eff0f5 100644 --- a/strix/interface/auth_cli.py +++ b/strix/interface/auth_cli.py @@ -1,8 +1,8 @@ -"""`strix auth` — ChatGPT subscription sign-in (login / status / logout). +"""`strix auth` — model-subscription sign-in (login / status / logout). Signing in only stores OAuth tokens (``~/.strix/subscription-auth.json``); model -selection stays with ``STRIX_LLM``. A ``chatgpt/`` STRIX_LLM runs on the -subscription. +selection stays with ``STRIX_LLM``. A ``chatgpt/`` STRIX_LLM runs on a +ChatGPT subscription and a ``grok/`` one on a Grok/SuperGrok subscription. """ from __future__ import annotations @@ -12,6 +12,7 @@ import base64 import logging import threading import webbrowser +from dataclasses import dataclass from http.server import BaseHTTPRequestHandler, HTTPServer from pathlib import Path from typing import TYPE_CHECKING, Any @@ -21,24 +22,76 @@ from rich.console import Console from rich.panel import Panel from rich.text import Text -from strix.config import codex, load_settings +from strix.config import codex, grok, load_settings if TYPE_CHECKING: from collections.abc import Callable + from types import ModuleType logger = logging.getLogger(__name__) _CALLBACK_TIMEOUT_S = 300 -# CLI-facing name for the login provider. Internally this is the Codex OAuth -# flow (``codex.PROVIDER``), but users know it as ChatGPT, so that's what the -# command and messaging say. ``codex`` is accepted as an alias. -LOGIN_PROVIDER = "chatgpt" -_ACCEPTED_PROVIDERS = frozenset({LOGIN_PROVIDER, codex.PROVIDER}) -_USAGE = "Usage:\n strix auth login chatgpt [--manual]\n strix auth status\n strix auth logout" +@dataclass(frozen=True) +class _Provider: + """A model-subscription provider the ``strix auth`` command can sign into. + + ``module`` is the provider's OAuth module (:mod:`strix.config.codex` or + :mod:`strix.config.grok`); both expose the same login surface. ``error`` is + that module's auth-error class, caught to report a clean failure. + """ + + name: str + module: ModuleType + error: type[Exception] + display: str + example_model: str + blurb: str + + +_PROVIDERS: dict[str, _Provider] = { + "chatgpt": _Provider( + name="chatgpt", + module=codex, + error=codex.CodexAuthError, + display="ChatGPT", + example_model="chatgpt/gpt-5.4", + blurb="This uses your ChatGPT Plus/Pro plan for inference instead of a metered API key.", + ), + "grok": _Provider( + name="grok", + module=grok, + error=grok.GrokAuthError, + display="Grok", + example_model="grok/grok-4", + blurb="This uses your Grok/SuperGrok plan for inference instead of a metered API key.", + ), +} + +# Internal OAuth provider ids and common vendor names accepted as aliases. +_PROVIDER_ALIASES: dict[str, str] = { + codex.PROVIDER: "chatgpt", + grok.PROVIDER: "grok", + "xai": "grok", + "supergrok": "grok", +} + +_DEFAULT_PROVIDER = "chatgpt" + +_USAGE = ( + "Usage:\n" + " strix auth login [chatgpt|grok] [--manual]\n" + " strix auth status\n" + " strix auth logout [chatgpt|grok]" +) + + +def _resolve_provider(name: str) -> _Provider | None: + key = _PROVIDER_ALIASES.get(name.lower(), name.lower()) + return _PROVIDERS.get(key) def run_auth(argv: list[str]) -> int: @@ -49,20 +102,20 @@ def run_auth(argv: list[str]) -> int: rest = argv[1:] if subcommand in ("-h", "--help", "help"): - console.print(_USAGE) + console.print(_USAGE, markup=False) return 0 handlers: dict[str, Callable[[], int]] = { "login": lambda: _login(console, rest), "status": lambda: _status(console), - "logout": lambda: _logout(console), + "logout": lambda: _logout(console, rest), } handler = handlers.get(subcommand) if handler is not None: return handler() console.print(f"[red]Unknown auth command:[/] {subcommand}\n") - console.print(_USAGE) + console.print(_USAGE, markup=False) return 2 @@ -71,8 +124,8 @@ def _login(console: Console, argv: list[str]) -> int: parser.add_argument( "provider", nargs="?", - default=LOGIN_PROVIDER, - help="Model provider to sign in with (default: chatgpt).", + default=_DEFAULT_PROVIDER, + help="Model provider to sign in with (chatgpt or grok; default: chatgpt).", ) parser.add_argument( "--manual", @@ -84,39 +137,42 @@ def _login(console: Console, argv: list[str]) -> int: except SystemExit as exc: # argparse already printed the message return int(exc.code or 2) - if args.provider.lower() not in _ACCEPTED_PROVIDERS: - console.print( - f"[red]Unsupported provider:[/] {args.provider}. " - f"Only '{LOGIN_PROVIDER}' (ChatGPT subscription) is supported." - ) + provider = _resolve_provider(args.provider) + if provider is None: + supported = ", ".join(f"'{name}'" for name in _PROVIDERS) + console.print(f"[red]Unsupported provider:[/] {args.provider}. Supported: {supported}.") return 2 - verifier, challenge = codex.generate_pkce() - state = codex.create_state() - authorize_url = codex.build_authorize_url(challenge, state) + module = provider.module + verifier, challenge = module.generate_pkce() + state = module.create_state() + authorize_url = module.build_authorize_url(challenge, state) console.print() - console.print("[bold]Signing in with ChatGPT[/] [dim](provider: chatgpt)[/]") console.print( - "[dim]This uses your ChatGPT Plus/Pro plan for inference instead of a metered API key.[/]" + f"[bold]Signing in with {provider.display}[/] [dim](provider: {provider.name})[/]" ) + console.print(f"[dim]{provider.blurb}[/]") console.print() try: - record = _run_oauth_flow(console, authorize_url, verifier, state, manual=args.manual) - except codex.CodexAuthError as exc: + record = _run_oauth_flow( + console, provider, authorize_url, verifier, state, manual=args.manual + ) + except provider.error as exc: return _fail(console, exc) except KeyboardInterrupt: console.print("\n[yellow]Sign-in cancelled.[/]") return 130 - codex.save_record(record) - _print_success(console) + module.save_record(record) + _print_success(console, provider) return 0 def _run_oauth_flow( console: Console, + provider: _Provider, authorize_url: str, verifier: str, state: str, @@ -124,7 +180,10 @@ def _run_oauth_flow( manual: bool, ) -> dict[str, Any]: """Drive the browser (or manual) OAuth flow and return a token record.""" - server = None if manual else _try_start_callback_server() + module = provider.module + server = ( + None if manual else _try_start_callback_server(module.CALLBACK_PORT, module.CALLBACK_PATH) + ) console.print("Open this URL in your browser to authorize:") console.print(f"[cyan]{authorize_url}[/]") @@ -142,8 +201,8 @@ def _run_oauth_flow( if result is not None: code, returned_state, error = result if error: - raise codex.CodexAuthError("oauth_error", error) - return _finish(code, returned_state, verifier, state, require_state=True) + raise provider.error("oauth_error", error) + return _finish(provider, code, returned_state, verifier, state, require_state=True) console.print("[yellow]Timed out waiting for the browser. Falling back to manual paste.[/]") # Manual fallback: the user completes sign-in and pastes the redirect URL @@ -153,12 +212,13 @@ def _run_oauth_flow( try: pasted = console.input("Paste the full redirect URL (or code#state): ").strip() except EOFError as exc: - raise codex.CodexAuthError("no_input", "no redirect URL provided") from exc - code, returned_state = codex.parse_redirect_input(pasted) - return _finish(code, returned_state, verifier, state, require_state=False) + raise provider.error("no_input", "no redirect URL provided") from exc + code, returned_state = module.parse_redirect_input(pasted) + return _finish(provider, code, returned_state, verifier, state, require_state=False) def _finish( + provider: _Provider, code: str | None, returned_state: str | None, verifier: str, @@ -167,16 +227,17 @@ def _finish( require_state: bool, ) -> dict[str, Any]: if not code: - raise codex.CodexAuthError("no_code", "no authorization code found in the redirect") - # The loopback callback from OpenAI always carries state, so a missing or - # mismatched value there is forged (CSRF) and must be rejected. Manual paste - # is user-initiated (the user copies their own redirect), so state is only - # validated when the pasted value includes it. + raise provider.error("no_code", "no authorization code found in the redirect") + # The loopback callback from the provider always carries state, so a missing + # or mismatched value there is forged (CSRF) and must be rejected. Manual + # paste is user-initiated (the user copies their own redirect), so state is + # only validated when the pasted value includes it. if require_state and returned_state is None: - raise codex.CodexAuthError("state_mismatch", "missing state in callback; possible CSRF") + raise provider.error("state_mismatch", "missing state in callback; possible CSRF") if returned_state is not None and returned_state != expected_state: - raise codex.CodexAuthError("state_mismatch", "state did not match; possible CSRF") - return codex.exchange_code(code, verifier) + raise provider.error("state_mismatch", "state did not match; possible CSRF") + record: dict[str, Any] = provider.module.exchange_code(code, verifier) + return record class _CallbackServer: @@ -203,7 +264,7 @@ class _CallbackServer: self._httpd.server_close() -def _try_start_callback_server() -> _CallbackServer | None: +def _try_start_callback_server(port: int, path: str) -> _CallbackServer | None: event = threading.Event() holder: dict[str, Any] = {} @@ -213,7 +274,7 @@ def _try_start_callback_server() -> _CallbackServer | None: def do_GET(self) -> None: parsed = urlparse(self.path) - if parsed.path != codex.CALLBACK_PATH: + if parsed.path != path: self.send_response(404) self.end_headers() return @@ -230,9 +291,9 @@ def _try_start_callback_server() -> _CallbackServer | None: event.set() try: - httpd = HTTPServer(("127.0.0.1", codex.CALLBACK_PORT), Handler) + httpd = HTTPServer(("127.0.0.1", port), Handler) except OSError: - logger.debug("could not bind callback port %d", codex.CALLBACK_PORT, exc_info=True) + logger.debug("could not bind callback port %d", port, exc_info=True) return None return _CallbackServer(httpd, event, holder) @@ -243,30 +304,64 @@ def _first(query: dict[str, list[str]], key: str) -> str | None: def _status(console: Console) -> int: - record = codex.read_record() - if record is None: - console.print("[yellow]Not signed in.[/] Run [cyan]strix auth login chatgpt[/] to sign in.") - return 1 settings = load_settings() - console.print("[green]Signed in[/] with a ChatGPT subscription.") - console.print(f" Account: [bold]{record.get('account_id')}[/]") - if codex.subscription_model(settings.llm.model): - console.print(f" Runs use the subscription (STRIX_LLM=[bold]{settings.llm.model}[/]).") - else: + active_model = settings.llm.model + signed_in_any = False + for provider in _PROVIDERS.values(): + record = provider.module.read_record() + if record is None: + continue + signed_in_any = True + console.print(f"[green]Signed in[/] with a {provider.display} subscription.") + account_id = record.get("account_id") + if account_id: + console.print(f" Account: [bold]{account_id}[/]") + if provider.module.subscription_model(active_model): + console.print(f" Runs use the subscription (STRIX_LLM=[bold]{active_model}[/]).") + else: + console.print( + f" [yellow]Note:[/] set [cyan]STRIX_LLM[/] to e.g. " + f"[cyan]{provider.example_model}[/] to run on this subscription." + ) + if not signed_in_any: console.print( - " [yellow]Note:[/] set [cyan]STRIX_LLM[/] to e.g. [cyan]chatgpt/gpt-5.4[/] " - "to run on the subscription." + "[yellow]Not signed in.[/] Run [cyan]strix auth login chatgpt[/] " + "or [cyan]strix auth login grok[/] to sign in." ) + return 1 return 0 -def _logout(console: Console) -> int: - codex.logout() - console.print("[green]Signed out.[/] Stored subscription credentials removed.") +def _logout(console: Console, argv: list[str]) -> int: + parser = argparse.ArgumentParser(prog="strix auth logout", add_help=True) + parser.add_argument( + "provider", + nargs="?", + default=None, + help="Provider to sign out of (chatgpt or grok; default: all).", + ) + try: + args = parser.parse_args(argv) + except SystemExit as exc: + return int(exc.code or 2) + + if args.provider is None: + for provider in _PROVIDERS.values(): + provider.module.logout() + console.print("[green]Signed out.[/] Stored subscription credentials removed.") + return 0 + + target = _resolve_provider(args.provider) + if target is None: + supported = ", ".join(f"'{name}'" for name in _PROVIDERS) + console.print(f"[red]Unsupported provider:[/] {args.provider}. Supported: {supported}.") + return 2 + target.module.logout() + console.print(f"[green]Signed out of {target.display}.[/] Stored credentials removed.") return 0 -def _fail(console: Console, exc: codex.CodexAuthError) -> int: +def _fail(console: Console, exc: Exception) -> int: error_text = Text() error_text.append("SIGN-IN FAILED", style="bold red") error_text.append("\n\n", style="white") @@ -284,17 +379,18 @@ def _fail(console: Console, exc: codex.CodexAuthError) -> int: return 1 -def _print_success(console: Console) -> None: +def _print_success(console: Console, provider: _Provider) -> None: + prefix = provider.module.SUBSCRIPTION_PREFIX text = Text() - text.append("Signed in with your ChatGPT subscription", style="bold #22c55e") + text.append(f"Signed in with your {provider.display} subscription", style="bold #22c55e") text.append("\n\n", style="white") text.append("Set ", style="white") text.append("STRIX_LLM", style="bold white") text.append(" to a ", style="white") - text.append("chatgpt/", style="bold cyan") + text.append(prefix, style="bold cyan") text.append(" model (e.g. ", style="white") - text.append("chatgpt/gpt-5.4", style="bold cyan") - text.append(") — runs are billed to your ChatGPT plan.", style="white") + text.append(provider.example_model, style="bold cyan") + text.append(f") — runs are billed to your {provider.display} plan.", style="white") text.append("\n\n", style="white") text.append("Run a scan as usual, e.g. ", style="white") text.append("strix --target https://example.com", style="bold cyan") diff --git a/strix/interface/main.py b/strix/interface/main.py index 0fbcbf48..39ecfa1e 100644 --- a/strix/interface/main.py +++ b/strix/interface/main.py @@ -21,8 +21,10 @@ from rich.text import Text from strix.config import ( apply_config_override, codex, + grok, load_settings, persist_current, + subscription, ) from strix.config.models import ( RECOMMENDED_MODEL_NAMES, @@ -104,6 +106,16 @@ def validate_environment() -> None: logger.info("Environment OK (ChatGPT subscription)") return + if grok.subscription_model(settings.llm.model): + if not grok.is_authenticated(): + console.print( + f"[red]STRIX_LLM={settings.llm.model} uses your Grok subscription, " + "but you're not signed in.[/] Run [cyan]strix auth login grok[/] first." + ) + sys.exit(1) + logger.info("Environment OK (Grok subscription)") + return + if not settings.llm.model: missing_required_vars.append("STRIX_LLM") @@ -784,7 +796,7 @@ def _persist_run_record(args: argparse.Namespace) -> None: "status": "running", "start_time": datetime.now(UTC).isoformat(), "end_time": None, - "auth_mode": codex.auth_mode(load_settings().llm.model), + "auth_mode": subscription.auth_mode(load_settings().llm.model), "targets_info": args.targets_info, "scan_mode": args.scan_mode, "instruction": args.instruction, @@ -1059,7 +1071,7 @@ def main() -> None: _telemetry_start_kwargs = { "model": load_settings().llm.model, - "auth_mode": codex.auth_mode(load_settings().llm.model), + "auth_mode": subscription.auth_mode(load_settings().llm.model), "scan_mode": args.scan_mode, "is_whitebox": is_whitebox_scan(args.targets_info), "interactive": not args.non_interactive, diff --git a/strix/interface/utils.py b/strix/interface/utils.py index 872c77fe..8aa3c0b3 100644 --- a/strix/interface/utils.py +++ b/strix/interface/utils.py @@ -261,9 +261,18 @@ def _is_subscription(report_state: Any) -> bool: record = getattr(report_state, "run_record", None) if isinstance(record, dict) and record.get("auth_mode"): return record.get("auth_mode") == "subscription" - from strix.config import codex + from strix.config import subscription - return codex.auth_mode(load_settings().llm.model) == "subscription" + return subscription.auth_mode(load_settings().llm.model) == "subscription" + + +def _subscription_label() -> str: + """Human label for the active model subscription (e.g. "Grok subscription").""" + from strix.config import grok + + if grok.subscription_model(load_settings().llm.model): + return "Grok subscription" + return "ChatGPT subscription" def _int_stat(usage: dict[str, Any], key: str) -> int: @@ -362,7 +371,7 @@ def build_live_stats_text(report_state: Any) -> Text: stats_text.append(str(model), style="white") if _is_subscription(report_state): stats_text.append(" · ", style="dim white") - stats_text.append("ChatGPT subscription", style="#22c55e") + stats_text.append(_subscription_label(), style="#22c55e") stats_text.append("\n") vuln_count = len(report_state.vulnerability_reports) @@ -408,7 +417,7 @@ def build_tui_stats_text(report_state: Any) -> Text: subscription = _is_subscription(report_state) if subscription: stats_text.append("\n") - stats_text.append("ChatGPT subscription", style="#22c55e") + stats_text.append(_subscription_label(), style="#22c55e") usage = _llm_usage(report_state) if usage and _int_stat(usage, "total_tokens") > 0: diff --git a/strix/report/state.py b/strix/report/state.py index 1475ccf3..78798d97 100644 --- a/strix/report/state.py +++ b/strix/report/state.py @@ -10,7 +10,7 @@ from uuid import uuid4 from agents.usage import Usage -from strix.config import codex +from strix.config import subscription from strix.config.loader import load_settings from strix.core.paths import run_dir_for from strix.report.sarif import write_sarif @@ -119,7 +119,7 @@ class ReportState: self.scan_results: dict[str, Any] | None = None self.scan_config: dict[str, Any] | None = None self._llm_usage = LLMUsageLedger() - auth_mode = codex.auth_mode(load_settings().llm.model) + auth_mode = subscription.auth_mode(load_settings().llm.model) self._llm_usage.zero_cost = auth_mode == "subscription" self.run_record: dict[str, Any] = { "run_id": self.run_id, diff --git a/tests/test_auth_cli.py b/tests/test_auth_cli.py index 9b66a2db..190b8741 100644 --- a/tests/test_auth_cli.py +++ b/tests/test_auth_cli.py @@ -6,23 +6,34 @@ from typing import TYPE_CHECKING, Any import pytest -from strix.config import codex +from strix.config import codex, grok from strix.interface import auth_cli if TYPE_CHECKING: from pathlib import Path +_CHATGPT = auth_cli._PROVIDERS["chatgpt"] + @pytest.fixture(autouse=True) def _tmp_store(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(codex, "AUTH_PATH", tmp_path / "home" / ".strix" / "subscription-auth.json") + store = tmp_path / "home" / ".strix" / "subscription-auth.json" + monkeypatch.setattr(codex, "AUTH_PATH", store) + monkeypatch.setattr(grok, "AUTH_PATH", store) -def test_login_provider_is_chatgpt() -> None: - assert auth_cli.LOGIN_PROVIDER == "chatgpt" - assert codex.PROVIDER in auth_cli._ACCEPTED_PROVIDERS - assert "chatgpt" in auth_cli._ACCEPTED_PROVIDERS +def test_default_provider_is_chatgpt() -> None: + assert auth_cli._DEFAULT_PROVIDER == "chatgpt" + assert set(auth_cli._PROVIDERS) == {"chatgpt", "grok"} + + +def test_provider_aliases_resolve() -> None: + assert auth_cli._resolve_provider(codex.PROVIDER) is _CHATGPT + assert auth_cli._resolve_provider("ChatGPT") is _CHATGPT + assert auth_cli._resolve_provider("grok") is auth_cli._PROVIDERS["grok"] + assert auth_cli._resolve_provider("xai") is auth_cli._PROVIDERS["grok"] + assert auth_cli._resolve_provider("gemini") is None def test_unknown_subcommand_returns_usage_error() -> None: @@ -51,32 +62,32 @@ def test_finish_requires_state_on_loopback(monkeypatch: pytest.MonkeyPatch) -> N # Loopback (require_state=True): missing or mismatched state is rejected. with pytest.raises(codex.CodexAuthError) as missing: - auth_cli._finish("code", None, "verifier", "expected", require_state=True) + auth_cli._finish(_CHATGPT, "code", None, "verifier", "expected", require_state=True) assert missing.value.code == "state_mismatch" with pytest.raises(codex.CodexAuthError) as mismatch: - auth_cli._finish("code", "wrong", "verifier", "expected", require_state=True) + auth_cli._finish(_CHATGPT, "code", "wrong", "verifier", "expected", require_state=True) assert mismatch.value.code == "state_mismatch" # Matching state proceeds to the exchange. - assert auth_cli._finish("code", "expected", "verifier", "expected", require_state=True) == { - "ok": True - } + assert auth_cli._finish( + _CHATGPT, "code", "expected", "verifier", "expected", require_state=True + ) == {"ok": True} def test_finish_manual_paste_allows_absent_state(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(codex, "exchange_code", lambda *_: {"ok": True}) # Manual paste (require_state=False): a bare code with no state is accepted, # but a present-and-wrong state is still rejected. - assert auth_cli._finish("code", None, "verifier", "expected", require_state=False) == { - "ok": True - } + assert auth_cli._finish( + _CHATGPT, "code", None, "verifier", "expected", require_state=False + ) == {"ok": True} with pytest.raises(codex.CodexAuthError): - auth_cli._finish("code", "wrong", "verifier", "expected", require_state=False) + auth_cli._finish(_CHATGPT, "code", "wrong", "verifier", "expected", require_state=False) def test_finish_rejects_missing_code() -> None: with pytest.raises(codex.CodexAuthError) as exc: - auth_cli._finish(None, "expected", "verifier", "expected", require_state=True) + auth_cli._finish(_CHATGPT, None, "expected", "verifier", "expected", require_state=True) assert exc.value.code == "no_code" diff --git a/tests/test_grok_auth.py b/tests/test_grok_auth.py new file mode 100644 index 00000000..a8c75d43 --- /dev/null +++ b/tests/test_grok_auth.py @@ -0,0 +1,265 @@ +"""Tests for Grok (xAI) subscription auth: PKCE, token handling, store.""" + +from __future__ import annotations + +import base64 +import hashlib +import json +import time +from typing import TYPE_CHECKING, Any +from unittest import mock + +import pytest +import requests + +from strix.config import grok + + +if TYPE_CHECKING: + from pathlib import Path + + +@pytest.fixture(autouse=True) +def _tmp_store(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + path = tmp_path / "home" / ".strix" / "subscription-auth.json" + monkeypatch.setattr(grok, "AUTH_PATH", path) + return path + + +def test_pkce_challenge_matches_verifier_and_is_unpadded() -> None: + verifier, challenge = grok.generate_pkce() + expected = ( + base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode() + ) + assert challenge == expected + assert "=" not in verifier + assert "=" not in challenge + + +def test_authorize_url_carries_pkce_client_and_grok_scope() -> None: + url = grok.build_authorize_url("chal", "st8") + assert grok.AUTHORIZE_URL in url + assert "code_challenge=chal" in url + assert "code_challenge_method=S256" in url + assert f"client_id={grok.CLIENT_ID}" in url + assert "state=st8" in url + # The Grok-CLI scope is what unlocks subscription inference. + assert "grok-cli%3Aaccess" in url + assert "offline_access" in url + + +def test_redirect_uri_is_loopback() -> None: + assert grok.REDIRECT_URI == "http://127.0.0.1:56121/callback" + + +def test_post_form_returns_parsed_body() -> None: + resp = mock.MagicMock() + resp.status_code = 200 + resp.content = b'{"access_token": "tok"}' + + with mock.patch.object(requests, "post", return_value=resp) as post: + data = grok._post_form({"grant_type": "refresh_token"}) + + assert data == {"access_token": "tok"} + assert post.call_args.kwargs["timeout"] == grok._TOKEN_TIMEOUT + + +def test_post_form_raises_on_http_error() -> None: + resp = mock.MagicMock() + resp.status_code = 400 + resp.text = "invalid_grant" + + with ( + mock.patch.object(requests, "post", return_value=resp), + pytest.raises(grok.GrokAuthError) as exc, + ): + grok._post_form({"grant_type": "refresh_token"}) + assert exc.value.code == "token_http_error" + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("http://127.0.0.1:56121/callback?code=AAA&state=BBB", ("AAA", "BBB")), + ("AAA#BBB", ("AAA", "BBB")), + ("code=AAA&state=BBB", ("AAA", "BBB")), + ("AAA", ("AAA", None)), + ("", (None, None)), + ], +) +def test_parse_redirect_input(value: str, expected: tuple[str | None, str | None]) -> None: + assert grok.parse_redirect_input(value) == expected + + +@pytest.mark.parametrize( + ("model", "expected"), + [ + ("grok/grok-4", "grok-4"), + ("Grok/Grok-4", "Grok-4"), + (" grok/grok-4 ", "grok-4"), + ("xai/grok-4", None), # metered API path + ("chatgpt/gpt-5.4", None), + ("grok-4", None), + ("grok/", None), + ("", None), + (None, None), + ], +) +def test_subscription_model(model: str | None, expected: str | None) -> None: + assert grok.subscription_model(model) == expected + + +def test_auth_mode() -> None: + assert grok.auth_mode("grok/grok-4") == "subscription" + assert grok.auth_mode("xai/grok-4") == "api_key" + assert grok.auth_mode("chatgpt/gpt-5.4") == "api_key" + assert grok.auth_mode(None) == "api_key" + + +def _record(access: str, refresh: str, expires_at: float) -> dict[str, Any]: + return { + "type": "oauth", + "provider": "grok", + "access": access, + "refresh": refresh, + "expires_at": expires_at, + } + + +def test_store_roundtrip_and_logout() -> None: + assert grok.read_record() is None + assert grok.is_authenticated() is False + + grok.save_record(_record("a1", "r1", time.time() + 3600)) + record = grok.read_record() + assert record is not None + assert record["access"] == "a1" + assert grok.is_authenticated() is True + + grok.logout() + assert grok.read_record() is None + grok.logout() # no-op when already gone + + +def test_store_file_permissions_are_owner_only(_tmp_store: Path) -> None: + grok.save_record(_record("a1", "r1", time.time() + 3600)) + assert (_tmp_store.stat().st_mode & 0o777) == 0o600 + + +def test_store_shares_file_with_other_providers(_tmp_store: Path) -> None: + # Grok must not clobber a co-resident ChatGPT record in the shared store. + _tmp_store.parent.mkdir(parents=True, exist_ok=True) + _tmp_store.write_text(json.dumps({"codex": {"type": "oauth", "access": "x"}})) + + grok.save_record(_record("a1", "r1", time.time() + 3600)) + on_disk = json.loads(_tmp_store.read_text()) + assert on_disk["codex"] == {"type": "oauth", "access": "x"} + assert on_disk["grok"]["access"] == "a1" + + grok.logout() + # Removing grok leaves the other provider's record and the file intact. + assert json.loads(_tmp_store.read_text()) == {"codex": {"type": "oauth", "access": "x"}} + + +def test_read_record_rejects_incomplete_records() -> None: + grok.save_record({"type": "oauth", "access": "a"}) # missing refresh + assert grok.read_record() is None + assert grok.is_authenticated() is False + + +def test_get_valid_token_returns_stored_when_fresh(monkeypatch: pytest.MonkeyPatch) -> None: + def _boom(_payload: dict[str, str]) -> dict[str, Any]: + msg = "should not refresh a fresh token" + raise AssertionError(msg) + + monkeypatch.setattr(grok, "_post_form", _boom) + grok.save_record(_record("access-fresh", "r1", time.time() + 3600)) + assert grok.get_valid_token() == "access-fresh" + + +def test_get_valid_token_refreshes_and_persists_rotation(monkeypatch: pytest.MonkeyPatch) -> None: + calls = {"n": 0} + + def _fake_post(payload: dict[str, str]) -> dict[str, Any]: + calls["n"] += 1 + assert payload["grant_type"] == "refresh_token" + assert payload["refresh_token"] == "r1" + return {"access_token": "access-new", "refresh_token": "r2", "expires_in": 3600} + + monkeypatch.setattr(grok, "_post_form", _fake_post) + grok.save_record(_record("stale", "r1", time.time() - 10)) # already expired + + assert grok.get_valid_token() == "access-new" + assert calls["n"] == 1 + record = grok.read_record() + assert record is not None + assert record["refresh"] == "r2" # rotated refresh written back + + +def test_refresh_keeps_old_refresh_when_response_omits_it(monkeypatch: pytest.MonkeyPatch) -> None: + def _fake_post(_payload: dict[str, str]) -> dict[str, Any]: + return {"access_token": "access-new", "expires_in": 3600} # no refresh_token + + monkeypatch.setattr(grok, "_post_form", _fake_post) + grok.save_record(_record("stale", "r1", time.time() - 10)) + + assert grok.get_valid_token() == "access-new" + record = grok.read_record() + assert record is not None + assert record["refresh"] == "r1" # fell back to the prior refresh token + + +def test_get_valid_token_uses_token_rotated_by_another_process( + monkeypatch: pytest.MonkeyPatch, +) -> None: + records = [ + _record("stale", "r1", time.time() - 10), + _record("fresh-from-other-process", "r2", time.time() + 3600), + ] + calls = {"n": 0} + + def _fake_read() -> dict[str, Any]: + record = records[min(calls["n"], len(records) - 1)] + calls["n"] += 1 + return record + + def _boom(_payload: dict[str, str]) -> dict[str, Any]: + msg = "must not refresh a token another process already rotated" + raise AssertionError(msg) + + monkeypatch.setattr(grok, "read_record", _fake_read) + monkeypatch.setattr(grok, "_post_form", _boom) + + assert grok.get_valid_token() == "fresh-from-other-process" + + +def test_get_valid_token_recovers_when_refresh_loses_race( + monkeypatch: pytest.MonkeyPatch, +) -> None: + grok.save_record(_record("stale", "r1", time.time() - 10)) + + def _fake_post(_payload: dict[str, str]) -> dict[str, Any]: + grok.save_record(_record("fresh-from-peer", "r2", time.time() + 3600)) + raise grok.GrokAuthError("token_http_error", "HTTP 400: invalid_grant") + + monkeypatch.setattr(grok, "_post_form", _fake_post) + assert grok.get_valid_token() == "fresh-from-peer" + + +def test_get_valid_token_reraises_refresh_error_without_rotation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + grok.save_record(_record("stale", "r1", time.time() - 10)) + + def _fake_post(_payload: dict[str, str]) -> dict[str, Any]: + raise grok.GrokAuthError("token_http_error", "HTTP 400: invalid_grant") + + monkeypatch.setattr(grok, "_post_form", _fake_post) + with pytest.raises(grok.GrokAuthError): + grok.get_valid_token() + + +def test_get_valid_token_raises_when_not_signed_in() -> None: + with pytest.raises(grok.GrokAuthError) as exc: + grok.get_valid_token() + assert exc.value.code == "not_authenticated" diff --git a/tests/test_grok_routing.py b/tests/test_grok_routing.py new file mode 100644 index 00000000..69edc695 --- /dev/null +++ b/tests/test_grok_routing.py @@ -0,0 +1,34 @@ +"""Grok subscription routing through StrixProvider.get_model.""" + +from __future__ import annotations + +from unittest import mock + +from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel + +from strix.config import grok +from strix.config.models import StrixProvider + + +def test_grok_prefix_routes_to_chat_completions(monkeypatch) -> None: # type: ignore[no-untyped-def] + client = mock.MagicMock() + monkeypatch.setattr(grok, "get_subscription_client", lambda: client) + + model = StrixProvider().get_model("grok/grok-4") + + assert isinstance(model, OpenAIChatCompletionsModel) + # The provider strips the grok/ prefix and passes xAI's bare model slug. + assert model.model == "grok-4" + + +def test_non_subscription_model_is_not_hijacked_by_grok(monkeypatch) -> None: # type: ignore[no-untyped-def] + def _boom() -> object: + msg = "grok client must not be built for a non-grok model" + raise AssertionError(msg) + + monkeypatch.setattr(grok, "get_subscription_client", _boom) + + # A metered xai/* key model must fall through to the normal provider path, + # not the subscription route. + model = StrixProvider().get_model("xai/grok-4") + assert not (isinstance(model, OpenAIChatCompletionsModel) and model.model == "grok-4")