feat(xai-oauth): add xAI Grok OAuth (SuperGrok Subscription) provider

Adds a new authentication provider that lets SuperGrok subscribers sign
in to Hermes with their xAI account via the standard OAuth 2.0 PKCE
loopback flow, instead of pasting a raw API key from console.x.ai.

Highlights
----------
* OAuth 2.0 PKCE loopback login against accounts.x.ai with discovery,
  state/nonce, and a strict CORS-origin allowlist on the callback.
* Authorize URL carries `plan=generic` (required for non-allowlisted
  loopback clients) and `referrer=hermes-agent` for best-effort
  attribution in xAI's OAuth server logs.
* Token storage in `auth.json` with file-locked atomic writes; JWT
  `exp`-based expiry detection with skew; refresh-token rotation
  synced both ways between the singleton store and the credential
  pool so multi-process / multi-profile setups don't tear each other's
  refresh tokens.
* Reactive 401 retry: on a 401 from the xAI Responses API, the agent
  refreshes the token, swaps it back into `self.api_key`, and retries
  the call once. Guarded against silent account swaps when the active
  key was sourced from a different (manual) pool entry.
* Auxiliary tasks (curator, vision, embeddings, etc.) route through a
  dedicated xAI Responses-mode auxiliary client instead of falling back
  to OpenRouter billing.
* Direct HTTP tools (`tools/xai_http.py`, transcription, TTS, image-gen
  plugin) resolve credentials through a unified runtime → singleton →
  env-var fallback chain so xai-oauth users get them for free.
* `hermes auth add xai-oauth` and `hermes auth remove xai-oauth N` are
  wired through the standard auth-commands surface; remove cleans up
  the singleton loopback_pkce entry so it doesn't silently reinstate.
* `hermes model` provider picker shows
  "xAI Grok OAuth (SuperGrok Subscription)" and the model-flow falls
  back to pool credentials when the singleton is missing.

Hardening
---------
* Discovery and refresh responses validate the returned
  `token_endpoint` host against the same `*.x.ai` allowlist as the
  authorization endpoint, blocking MITM persistence of a hostile
  endpoint.
* Discovery / refresh / token-exchange `response.json()` calls are
  wrapped to raise typed `AuthError` on malformed bodies (captive
  portals, proxy error pages) instead of leaking JSONDecodeError
  tracebacks.
* `prompt_cache_key` is routed through `extra_body` on the codex
  transport (sending it as a top-level kwarg trips xAI's SDK with a
  TypeError).
* Credential-pool sync-back preserves `active_provider` so refreshing
  an OAuth entry doesn't silently flip the active provider out from
  under the running agent.

Testing
-------
* New `tests/hermes_cli/test_auth_xai_oauth_provider.py` (~63 tests)
  covers JWT expiry, OAuth URL params (plan + referrer), CORS origins,
  redirect URI validation, singleton↔pool sync, concurrency races,
  refresh error paths, runtime resolution, and malformed-JSON guards.
* Extended `test_credential_pool.py`, `test_codex_transport.py`, and
  `test_run_agent_codex_responses.py` cover the pool sync-back,
  `extra_body` routing, and 401 reactive refresh paths.
* 165 tests passing on this branch via `scripts/run_tests.sh`.
This commit is contained in:
Jaaneek 2026-05-15 16:10:38 +01:00 committed by Teknium
parent 9fb40e6a3d
commit b62c997973
27 changed files with 3843 additions and 131 deletions

View file

@ -31,7 +31,7 @@ from agent.image_gen_provider import (
save_b64_image,
success_response,
)
from tools.xai_http import hermes_xai_user_agent
from tools.xai_http import hermes_xai_user_agent, resolve_xai_http_credentials
logger = logging.getLogger(__name__)
@ -39,14 +39,17 @@ logger = logging.getLogger(__name__)
# Model catalog
# ---------------------------------------------------------------------------
API_MODEL = "grok-imagine-image"
_MODELS: Dict[str, Dict[str, Any]] = {
"grok-imagine-image": {
"display": "Grok Imagine Image",
"speed": "~5-10s",
"strengths": "Fast, high-quality",
},
"grok-imagine-image-quality": {
"display": "Grok Imagine Image (Quality)",
"speed": "~10-20s",
"strengths": "Higher fidelity / detail; slower than the standard model.",
},
}
DEFAULT_MODEL = "grok-imagine-image"
@ -127,7 +130,8 @@ class XAIImageGenProvider(ImageGenProvider):
return "xAI (Grok)"
def is_available(self) -> bool:
return bool(os.getenv("XAI_API_KEY"))
creds = resolve_xai_http_credentials()
return bool(creds.get("api_key"))
def list_models(self) -> List[Dict[str, Any]]:
return [
@ -141,17 +145,16 @@ class XAIImageGenProvider(ImageGenProvider):
]
def get_setup_schema(self) -> Dict[str, Any]:
# Auth resolution is delegated to the shared ``xai_grok`` post_setup
# hook (``hermes_cli/tools_config.py``); identical to the TTS / video
# gen entries so users see the same OAuth-or-API-key choice for every
# xAI service.
return {
"name": "xAI (Grok)",
"name": "xAI Grok Imagine (image)",
"badge": "paid",
"tag": "Native xAI image generation via grok-imagine-image",
"env_vars": [
{
"key": "XAI_API_KEY",
"prompt": "xAI API key",
"url": "https://console.x.ai/",
},
],
"tag": "grok-imagine-image — text-to-image; uses xAI Grok OAuth or XAI_API_KEY",
"env_vars": [],
"post_setup": "xai_grok",
}
def generate(
@ -161,12 +164,14 @@ class XAIImageGenProvider(ImageGenProvider):
**kwargs: Any,
) -> Dict[str, Any]:
"""Generate an image using xAI's grok-imagine-image."""
api_key = os.getenv("XAI_API_KEY", "").strip()
creds = resolve_xai_http_credentials()
api_key = str(creds.get("api_key") or "").strip()
provider_name = str(creds.get("provider") or "xai").strip() or "xai"
if not api_key:
return error_response(
error="XAI_API_KEY not set. Get one at https://console.x.ai/",
error="No xAI credentials found. Configure xAI OAuth in `hermes model` or set XAI_API_KEY.",
error_type="missing_api_key",
provider="xai",
provider=provider_name,
aspect_ratio=aspect_ratio,
)
@ -177,7 +182,7 @@ class XAIImageGenProvider(ImageGenProvider):
xai_res = resolution if resolution in _XAI_RESOLUTIONS else DEFAULT_RESOLUTION
payload: Dict[str, Any] = {
"model": API_MODEL,
"model": model_id,
"prompt": prompt,
"aspect_ratio": xai_ar,
"resolution": xai_res,
@ -189,7 +194,7 @@ class XAIImageGenProvider(ImageGenProvider):
"User-Agent": hermes_xai_user_agent(),
}
base_url = (os.getenv("XAI_BASE_URL") or "https://api.x.ai/v1").strip().rstrip("/")
base_url = str(creds.get("base_url") or "https://api.x.ai/v1").strip().rstrip("/")
try:
response = requests.post(
@ -210,7 +215,7 @@ class XAIImageGenProvider(ImageGenProvider):
return error_response(
error=f"xAI image generation failed ({status}): {err_msg}",
error_type="api_error",
provider="xai",
provider=provider_name,
model=model_id,
prompt=prompt,
aspect_ratio=aspect,
@ -219,7 +224,7 @@ class XAIImageGenProvider(ImageGenProvider):
return error_response(
error="xAI image generation timed out (120s)",
error_type="timeout",
provider="xai",
provider=provider_name,
model=model_id,
prompt=prompt,
aspect_ratio=aspect,
@ -228,7 +233,7 @@ class XAIImageGenProvider(ImageGenProvider):
return error_response(
error=f"xAI connection error: {exc}",
error_type="connection_error",
provider="xai",
provider=provider_name,
model=model_id,
prompt=prompt,
aspect_ratio=aspect,
@ -240,7 +245,7 @@ class XAIImageGenProvider(ImageGenProvider):
return error_response(
error=f"xAI returned invalid JSON: {exc}",
error_type="invalid_response",
provider="xai",
provider=provider_name,
model=model_id,
prompt=prompt,
aspect_ratio=aspect,
@ -252,7 +257,7 @@ class XAIImageGenProvider(ImageGenProvider):
return error_response(
error="xAI returned no image data",
error_type="empty_response",
provider="xai",
provider=provider_name,
model=model_id,
prompt=prompt,
aspect_ratio=aspect,

View file

@ -10,8 +10,12 @@ Originally salvaged from PR #10600 by @Jaaneek; reshaped into the
:class:`VideoGenProvider` plugin interface and trimmed to the
generate-only surface.
Authentication via ``XAI_API_KEY``. Output is an HTTPS URL from xAI's
CDN; the gateway downloads and delivers it.
Authentication: xAI Grok OAuth tokens (preferred billed against the
user's SuperGrok subscription) or ``XAI_API_KEY``. Both routes are
resolved through ``tools.xai_http.resolve_xai_http_credentials`` so a
single login covers chat + TTS + image gen + video gen + transcription.
Output is an HTTPS URL from xAI's CDN; the gateway downloads and
delivers it.
"""
from __future__ import annotations
@ -20,7 +24,7 @@ import asyncio
import logging
import os
import uuid
from typing import Any, Dict, List, Optional
from typing import Any, Dict, List, Optional, Tuple
import httpx
@ -66,24 +70,44 @@ _MODELS: Dict[str, Dict[str, Any]] = {
# ---------------------------------------------------------------------------
def _xai_base_url() -> str:
return (os.getenv("XAI_BASE_URL") or DEFAULT_XAI_BASE_URL).strip().rstrip("/")
def _resolve_xai_credentials() -> Tuple[str, str]:
"""Return ``(api_key, base_url)`` from the shared xAI credential resolver.
Order: runtime provider (xai-oauth pool entry) singleton ``auth.json``
OAuth tokens ``XAI_API_KEY`` env var. ``api_key`` is empty when no
credential source is available; callers must check before using it.
"""
try:
from tools.xai_http import resolve_xai_http_credentials
creds = resolve_xai_http_credentials() or {}
except Exception as exc:
logger.debug("xAI credential resolver failed: %s", exc)
creds = {}
api_key = str(creds.get("api_key") or os.getenv("XAI_API_KEY", "")).strip()
base_url = str(
creds.get("base_url")
or os.getenv("XAI_BASE_URL")
or DEFAULT_XAI_BASE_URL
).strip().rstrip("/")
return api_key, base_url
def _xai_headers() -> Dict[str, str]:
api_key = os.getenv("XAI_API_KEY", "").strip()
if not api_key:
raise ValueError("XAI_API_KEY not set. Get one at https://console.x.ai/")
def _xai_user_agent() -> str:
try:
from tools.xai_http import hermes_xai_user_agent
ua = hermes_xai_user_agent()
return hermes_xai_user_agent()
except Exception:
ua = "hermes-agent/video_gen"
return "hermes-agent/video_gen"
def _xai_headers(api_key: str) -> Dict[str, str]:
return {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"User-Agent": ua,
"User-Agent": _xai_user_agent(),
}
@ -110,12 +134,15 @@ def _clamp_duration(duration: Optional[int], has_reference_images: bool) -> int:
async def _submit(
client: httpx.AsyncClient,
payload: Dict[str, Any],
*,
api_key: str,
base_url: str,
) -> str:
"""POST to /videos/generations — xAI's only public endpoint for our
text-to-video and image-to-video surface."""
response = await client.post(
f"{_xai_base_url()}/videos/generations",
headers={**_xai_headers(), "x-idempotency-key": str(uuid.uuid4())},
f"{base_url}/videos/generations",
headers={**_xai_headers(api_key), "x-idempotency-key": str(uuid.uuid4())},
json=payload,
timeout=60,
)
@ -131,6 +158,8 @@ async def _poll(
client: httpx.AsyncClient,
request_id: str,
*,
api_key: str,
base_url: str,
timeout_seconds: int,
poll_interval: int,
) -> Dict[str, Any]:
@ -138,8 +167,8 @@ async def _poll(
last_status = "queued"
while elapsed < timeout_seconds:
response = await client.get(
f"{_xai_base_url()}/videos/{request_id}",
headers=_xai_headers(),
f"{base_url}/videos/{request_id}",
headers=_xai_headers(api_key),
timeout=30,
)
response.raise_for_status()
@ -174,7 +203,8 @@ class XAIVideoGenProvider(VideoGenProvider):
return "xAI"
def is_available(self) -> bool:
return bool(os.environ.get("XAI_API_KEY", "").strip())
api_key, _ = _resolve_xai_credentials()
return bool(api_key)
def list_models(self) -> List[Dict[str, Any]]:
return [{"id": mid, **meta} for mid, meta in _MODELS.items()]
@ -183,17 +213,18 @@ class XAIVideoGenProvider(VideoGenProvider):
return DEFAULT_MODEL
def get_setup_schema(self) -> Dict[str, Any]:
# Auth resolution lives entirely in the shared ``xai_grok`` post_setup
# hook (``hermes_cli/tools_config.py``) so the picker doesn't blindly
# prompt for an API key when the user is already signed in via xAI
# Grok OAuth (SuperGrok Subscription) — TTS / image gen / video gen
# all share the same credential resolver. The hook offers an
# OAuth-vs-API-key choice when neither is configured.
return {
"name": "xAI",
"name": "xAI Grok Imagine",
"badge": "paid",
"tag": "grok-imagine-video — text-to-video & image-to-video with reference images",
"env_vars": [
{
"key": "XAI_API_KEY",
"prompt": "xAI API key",
"url": "https://console.x.ai/",
},
],
"tag": "grok-imagine-video — text-to-video & image-to-video; uses xAI Grok OAuth or XAI_API_KEY",
"env_vars": [],
"post_setup": "xai_grok",
}
def capabilities(self) -> Dict[str, Any]:
@ -259,9 +290,14 @@ class XAIVideoGenProvider(VideoGenProvider):
aspect_ratio: str,
resolution: str,
) -> Dict[str, Any]:
if not os.environ.get("XAI_API_KEY", "").strip():
api_key, base_url = _resolve_xai_credentials()
if not api_key:
return error_response(
error="XAI_API_KEY not set. Get one at https://console.x.ai/",
error=(
"No xAI credentials found. Sign in via `hermes auth add xai-oauth` "
"(SuperGrok subscription) or set XAI_API_KEY from "
"https://console.x.ai/."
),
error_type="auth_required",
provider="xai", prompt=prompt,
)
@ -317,7 +353,9 @@ class XAIVideoGenProvider(VideoGenProvider):
async with httpx.AsyncClient() as client:
try:
request_id = await _submit(client, payload)
request_id = await _submit(
client, payload, api_key=api_key, base_url=base_url
)
except httpx.HTTPStatusError as exc:
detail = ""
try:
@ -334,6 +372,7 @@ class XAIVideoGenProvider(VideoGenProvider):
poll_result = await _poll(
client, request_id,
api_key=api_key, base_url=base_url,
timeout_seconds=DEFAULT_TIMEOUT_SECONDS,
poll_interval=DEFAULT_POLL_INTERVAL_SECONDS,
)