mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-05-19 04:52:06 +00:00
b62c997973
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`.
282 lines
10 KiB
Python
282 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""Tests for xAI image generation provider."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _fake_api_key(monkeypatch):
|
|
"""Ensure XAI_API_KEY is set for all tests."""
|
|
monkeypatch.setenv("XAI_API_KEY", "test-key-12345")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Provider class tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestXAIImageGenProvider:
|
|
def test_name(self):
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
provider = XAIImageGenProvider()
|
|
assert provider.name == "xai"
|
|
|
|
def test_display_name(self):
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
provider = XAIImageGenProvider()
|
|
assert provider.display_name == "xAI (Grok)"
|
|
|
|
def test_is_available_with_key(self, monkeypatch):
|
|
monkeypatch.setenv("XAI_API_KEY", "sk-xxx")
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
provider = XAIImageGenProvider()
|
|
assert provider.is_available() is True
|
|
|
|
def test_is_available_without_key(self, monkeypatch):
|
|
monkeypatch.delenv("XAI_API_KEY", raising=False)
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
provider = XAIImageGenProvider()
|
|
assert provider.is_available() is False
|
|
|
|
def test_list_models(self):
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
provider = XAIImageGenProvider()
|
|
models = provider.list_models()
|
|
assert len(models) >= 1
|
|
assert models[0]["id"] == "grok-imagine-image"
|
|
|
|
def test_default_model(self):
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
provider = XAIImageGenProvider()
|
|
assert provider.default_model() == "grok-imagine-image"
|
|
|
|
def test_get_setup_schema(self):
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
provider = XAIImageGenProvider()
|
|
schema = provider.get_setup_schema()
|
|
assert schema["name"] == "xAI Grok Imagine (image)"
|
|
assert schema["badge"] == "paid"
|
|
# Auth resolution is delegated to the shared "xai_grok" post_setup
|
|
# hook so the picker doesn't blindly prompt for XAI_API_KEY when the
|
|
# user is already signed in via xAI Grok OAuth.
|
|
assert schema["env_vars"] == []
|
|
assert schema["post_setup"] == "xai_grok"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Config tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestConfig:
|
|
def test_default_model(self):
|
|
from plugins.image_gen.xai import _resolve_model
|
|
|
|
model_id, meta = _resolve_model()
|
|
assert model_id == "grok-imagine-image"
|
|
|
|
def test_default_resolution(self):
|
|
from plugins.image_gen.xai import _resolve_resolution
|
|
|
|
assert _resolve_resolution() == "1k"
|
|
|
|
def test_custom_model(self, monkeypatch):
|
|
monkeypatch.setenv("XAI_IMAGE_MODEL", "grok-imagine-image")
|
|
from plugins.image_gen.xai import _resolve_model
|
|
|
|
model_id, _ = _resolve_model()
|
|
assert model_id == "grok-imagine-image"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Generate tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGenerate:
|
|
def test_missing_api_key(self, monkeypatch):
|
|
monkeypatch.delenv("XAI_API_KEY", raising=False)
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
provider = XAIImageGenProvider()
|
|
result = provider.generate(prompt="test")
|
|
assert result["success"] is False
|
|
assert "XAI_API_KEY" in result["error"]
|
|
|
|
def test_successful_generation(self):
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 200
|
|
mock_resp.raise_for_status = MagicMock()
|
|
mock_resp.json.return_value = {
|
|
"data": [{"b64_json": "dGVzdC1pbWFnZS1kYXRh"}], # base64 "test-image-data"
|
|
}
|
|
|
|
with patch("plugins.image_gen.xai.requests.post", return_value=mock_resp):
|
|
with patch("plugins.image_gen.xai.save_b64_image", return_value="/tmp/test.png"):
|
|
provider = XAIImageGenProvider()
|
|
result = provider.generate(prompt="A cat playing piano")
|
|
|
|
assert result["success"] is True
|
|
assert result["image"] == "/tmp/test.png"
|
|
assert result["provider"] == "xai"
|
|
assert result["model"] == "grok-imagine-image"
|
|
|
|
def test_successful_url_response(self):
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 200
|
|
mock_resp.raise_for_status = MagicMock()
|
|
mock_resp.json.return_value = {
|
|
"data": [{"url": "https://xai.image/result.png"}],
|
|
}
|
|
|
|
with patch("plugins.image_gen.xai.requests.post", return_value=mock_resp):
|
|
provider = XAIImageGenProvider()
|
|
result = provider.generate(prompt="A cat playing piano")
|
|
|
|
assert result["success"] is True
|
|
assert result["image"] == "https://xai.image/result.png"
|
|
|
|
def test_api_error(self):
|
|
import requests as req_lib
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 401
|
|
mock_resp.text = "Unauthorized"
|
|
mock_resp.json.return_value = {"error": {"message": "Invalid API key"}}
|
|
mock_resp.raise_for_status.side_effect = req_lib.HTTPError(response=mock_resp)
|
|
|
|
with patch("plugins.image_gen.xai.requests.post", return_value=mock_resp):
|
|
provider = XAIImageGenProvider()
|
|
result = provider.generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "api_error"
|
|
|
|
def test_api_error_preserves_real_response_status(self):
|
|
import requests as req_lib
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
response = req_lib.Response()
|
|
response.status_code = 401
|
|
response._content = json.dumps({"error": {"message": "Invalid API key"}}).encode()
|
|
response.headers["Content-Type"] = "application/json"
|
|
|
|
response.raise_for_status = MagicMock(
|
|
side_effect=req_lib.HTTPError(response=response)
|
|
)
|
|
|
|
with patch("plugins.image_gen.xai.requests.post", return_value=response):
|
|
provider = XAIImageGenProvider()
|
|
result = provider.generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "api_error"
|
|
assert "xAI image generation failed (401): Invalid API key" in result["error"]
|
|
|
|
def test_timeout(self):
|
|
import requests as req_lib
|
|
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
with patch("plugins.image_gen.xai.requests.post", side_effect=req_lib.Timeout()):
|
|
provider = XAIImageGenProvider()
|
|
result = provider.generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "timeout"
|
|
|
|
def test_empty_response(self):
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 200
|
|
mock_resp.raise_for_status = MagicMock()
|
|
mock_resp.json.return_value = {"data": []}
|
|
|
|
with patch("plugins.image_gen.xai.requests.post", return_value=mock_resp):
|
|
provider = XAIImageGenProvider()
|
|
result = provider.generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "empty_response"
|
|
|
|
def test_auth_header(self):
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 200
|
|
mock_resp.raise_for_status = MagicMock()
|
|
mock_resp.json.return_value = {
|
|
"data": [{"url": "https://xai.image/test.png"}],
|
|
}
|
|
|
|
with patch("plugins.image_gen.xai.requests.post", return_value=mock_resp) as mock_post:
|
|
provider = XAIImageGenProvider()
|
|
provider.generate(prompt="test")
|
|
|
|
call_args = mock_post.call_args
|
|
headers = call_args.kwargs.get("headers") or call_args[1].get("headers")
|
|
assert "Bearer test-key-12345" in headers["Authorization"]
|
|
assert "Hermes-Agent" in headers["User-Agent"]
|
|
|
|
def test_payload_resolution_is_literal_1k_or_2k(self):
|
|
"""Regression: xAI API rejects numeric resolutions ("1024"/"2048") with 422.
|
|
|
|
The endpoint expects the literal strings "1k" or "2k". Ensure the wire
|
|
payload carries that literal — not a numeric mapping. See PR #18678.
|
|
"""
|
|
from plugins.image_gen.xai import XAIImageGenProvider
|
|
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 200
|
|
mock_resp.raise_for_status = MagicMock()
|
|
mock_resp.json.return_value = {"data": [{"url": "https://xai.image/test.png"}]}
|
|
|
|
with patch("plugins.image_gen.xai.requests.post", return_value=mock_resp) as mock_post:
|
|
provider = XAIImageGenProvider()
|
|
provider.generate(prompt="test")
|
|
|
|
payload = mock_post.call_args.kwargs.get("json") or mock_post.call_args[1].get("json")
|
|
assert payload["resolution"] in {"1k", "2k"}, (
|
|
f"resolution must be the literal '1k' or '2k', got {payload['resolution']!r}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Registration test
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestRegistration:
|
|
def test_register(self):
|
|
from plugins.image_gen.xai import XAIImageGenProvider, register
|
|
|
|
mock_ctx = MagicMock()
|
|
register(mock_ctx)
|
|
mock_ctx.register_image_gen_provider.assert_called_once()
|
|
provider = mock_ctx.register_image_gen_provider.call_args[0][0]
|
|
assert isinstance(provider, XAIImageGenProvider)
|
|
assert provider.name == "xai"
|