From 34837597d298ee9b67be5a47409d7e24b6751ec7 Mon Sep 17 00:00:00 2001 From: cresslank <9219265+cresslank@users.noreply.github.com> Date: Fri, 10 Jul 2026 13:09:13 -0500 Subject: [PATCH] fix(auth): make xAI OAuth pools multi-account resilient Keep each xAI OAuth auth-add login as an independent manual device-code pool entry and recognize xAI personal-team spending-limit 403 responses as billing exhaustion. Preserve the structured top-level error message so the failed credential is quarantined and the next healthy account is selected without attempting a pointless token refresh. Route direct xAI HTTP consumers through the credential pool as well. Proactive and 401-reactive refreshes update the exact issuing manual entry, preserve validated xAI base URL overrides, and serialize single-use refresh-token rotation across concurrent pool instances. --- agent/agent_runtime_helpers.py | 13 +- agent/credential_pool.py | 88 +++++++++++- agent/error_classifier.py | 35 +++-- hermes_cli/auth_commands.py | 48 +++++-- plugins/web/xai/provider.py | 5 +- tests/agent/test_credential_pool.py | 72 ++++++++++ tests/agent/test_credential_pool_routing.py | 8 +- tests/agent/test_error_classifier.py | 34 +++++ tests/hermes_cli/test_auth_commands.py | 114 +++++++++++++-- .../test_codex_xai_oauth_recovery.py | 63 ++++++++ tests/run_agent/test_run_agent.py | 18 ++- tests/tools/test_web_providers_xai.py | 135 +++++++++++++++--- tools/xai_http.py | 89 ++++++++---- 13 files changed, 629 insertions(+), 93 deletions(-) diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index 78d12d007faa9..09c4aec5cb067 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -837,7 +837,14 @@ def recover_with_credential_pool( if effective_reason == FailoverReason.billing: rotate_status = status_code if status_code is not None else 402 - next_entry = pool.mark_exhausted_and_rotate(status_code=rotate_status, error_context=error_context) + next_entry = pool.mark_exhausted_and_rotate( + status_code=rotate_status, + error_context=error_context, + # Runtime credentials can be resolved by a separate pool instance, + # leaving this recovery pool without ``current_id``. Match the key + # that actually failed instead of quarantining a different account. + api_key_hint=getattr(agent, "api_key", None), + ) if next_entry is not None: _ra().logger.info( "Credential %s (billing) — rotated to pool entry %s", @@ -3134,6 +3141,10 @@ def extract_api_error_context(error: Exception) -> Dict[str, Any]: if isinstance(reason, str) and reason.strip(): context["reason"] = reason.strip() message = payload.get("message") or payload.get("error_description") + if not message and isinstance(payload.get("error"), str): + # xAI uses a top-level string ``error`` beside a structured + # ``code`` (for example personal-team-blocked:spending-limit). + message = payload.get("error") if isinstance(message, str) and message.strip(): context["message"] = message.strip() for key in ("resets_at", "reset_at"): diff --git a/agent/credential_pool.py b/agent/credential_pool.py index 8e24ffe498de1..95bfca5b5d06c 100644 --- a/agent/credential_pool.py +++ b/agent/credential_pool.py @@ -824,6 +824,45 @@ class CredentialPool: logger.debug("Failed to sync xAI OAuth entry from auth.json: %s", exc) return entry + def _sync_xai_oauth_entry_from_pool_store( + self, entry: PooledCredential + ) -> PooledCredential: + """Adopt a token pair rotated by another pool instance. + + Direct xAI integrations load a fresh ``CredentialPool`` for each + request. Their in-memory locks therefore cannot protect xAI's + single-use refresh token across concurrent requests or processes. + This helper is called while the shared auth-store lock is held and + re-reads the exact persisted row before a refresh POST is attempted. + """ + if self.provider != "xai-oauth": + return entry + try: + persisted = next( + ( + payload + for payload in read_credential_pool(self.provider) + if isinstance(payload, dict) and payload.get("id") == entry.id + ), + None, + ) + if not isinstance(persisted, dict): + return entry + stored = PooledCredential.from_dict(self.provider, persisted) + if ( + stored.access_token != entry.access_token + or stored.refresh_token != entry.refresh_token + ): + logger.debug( + "Pool entry %s: adopting xAI OAuth tokens rotated by another pool instance", + entry.id, + ) + self._replace_entry(entry, stored) + return stored + except Exception as exc: + logger.debug("Failed to sync xAI OAuth entry from credential pool: %s", exc) + return entry + def _sync_nous_entry_from_auth_store(self, entry: PooledCredential) -> PooledCredential: """Sync a Nous pool entry from auth.json if tokens differ. @@ -1040,6 +1079,22 @@ class CredentialPool: if not force and not self._entry_needs_refresh(entry): return entry return self._refresh_entry_impl(entry, force=force) + if self.provider == "xai-oauth": + refresh_timeout_seconds = auth_mod.env_float( + "HERMES_XAI_REFRESH_TIMEOUT_SECONDS", 20 + ) + lock_timeout = max( + float(auth_mod.AUTH_LOCK_TIMEOUT_SECONDS), + float(refresh_timeout_seconds) + 5.0, + ) + with _auth_store_lock(timeout_seconds=lock_timeout): + synced = self._sync_xai_oauth_entry_from_pool_store(entry) + if ( + synced.access_token != entry.access_token + or synced.refresh_token != entry.refresh_token + ): + return synced + return self._refresh_entry_impl(synced, force=force) return self._refresh_entry_impl(entry, force=force) def _refresh_entry_impl( @@ -1538,8 +1593,8 @@ class CredentialPool: self._persist(removed_ids=entries_to_prune) return available - def _select_unlocked(self) -> Optional[PooledCredential]: - available = self._available_entries(clear_expired=True, refresh=True) + def _select_unlocked(self, *, refresh: bool = True) -> Optional[PooledCredential]: + available = self._available_entries(clear_expired=True, refresh=refresh) if not available: self._current_id = None logger.info("credential pool: no available entries (all exhausted or empty)") @@ -1668,6 +1723,35 @@ class CredentialPool: with self._lock: return self._try_refresh_current_unlocked() + def try_refresh_matching( + self, api_key_hint: Optional[str] = None + ) -> Optional[PooledCredential]: + """Force-refresh the entry that supplied ``api_key_hint``. + + Direct provider integrations may reload the pool after a request has + already failed, so they cannot rely on ``current_id`` identifying the + issuing credential. With no hint, select an entry without first doing + the normal proactive refresh; the forced refresh below must consume a + rotating refresh token exactly once. + """ + with self._lock: + entry = None + if api_key_hint: + entry = next( + ( + candidate + for candidate in self._entries + if candidate.runtime_api_key == api_key_hint + ), + None, + ) + else: + entry = self.current() or self._select_unlocked(refresh=False) + if entry is None: + return None + self._current_id = entry.id + return self._try_refresh_current_unlocked() + def _try_refresh_current_unlocked(self) -> Optional[PooledCredential]: entry = self.current() if entry is None: diff --git a/agent/error_classifier.py b/agent/error_classifier.py index 6afdf4e076e0a..b374d4c32d610 100644 --- a/agent/error_classifier.py +++ b/agent/error_classifier.py @@ -123,6 +123,25 @@ _BILLING_PATTERNS = [ "not available on the free tier", ] +# xAI's explicit Grok credit-exhaustion code. Keep the HTTP 403 special case +# provider-scoped: other providers' generic billing codes historically remain +# auth failures when they arrive as 403. +_XAI_SPENDING_LIMIT_ERROR_CODE = "personal-team-blocked:spending-limit" + +# Structured provider codes that mean the account cannot serve paid traffic +# until credits/subscription capacity is restored. xAI returns its explicit +# Grok spending-limit signal as HTTP 403 rather than 402. +_BILLING_ERROR_CODES = frozenset({ + "insufficient_quota", + "billing_not_active", + "payment_required", + "insufficient_credits", + "no_usable_credits", + "balance_depleted", + "model_not_supported_on_free_tier", + _XAI_SPENDING_LIMIT_ERROR_CODE, +}) + # Patterns that indicate rate limiting (transient, will resolve) _RATE_LIMIT_PATTERNS = [ "rate limit", @@ -906,7 +925,11 @@ def _classify_by_status( # OpenRouter 403 "key limit exceeded" is actually billing. Other # providers also use 403 for account-plan or credit exhaustion. if ( - "key limit exceeded" in error_msg + ( + provider == "xai-oauth" + and error_code.lower() == _XAI_SPENDING_LIMIT_ERROR_CODE + ) + or "key limit exceeded" in error_msg or "spending limit" in error_msg or any(p in error_msg for p in _BILLING_PATTERNS) ): @@ -1292,15 +1315,7 @@ def _classify_by_error_code( should_rotate_credential=True, ) - if code_lower in { - "insufficient_quota", - "billing_not_active", - "payment_required", - "insufficient_credits", - "no_usable_credits", - "balance_depleted", - "model_not_supported_on_free_tier", - }: + if code_lower in _BILLING_ERROR_CODES: return result_fn( FailoverReason.billing, retryable=False, diff --git a/hermes_cli/auth_commands.py b/hermes_cli/auth_commands.py index 15d5fab61c0ac..a5b3e7210b8ed 100644 --- a/hermes_cli/auth_commands.py +++ b/hermes_cli/auth_commands.py @@ -314,8 +314,9 @@ def auth_add_command(args) -> None: _oauth_default_label(provider, len(pool.entries()) + 1), ) # Add a distinct, self-contained pool entry per account (matching the - # xai-oauth / qwen-oauth patterns) instead of - # routing through the singleton ``_save_codex_tokens`` save path. + # qwen-oauth / minimax-oauth multi-account patterns, and the + # xai-oauth path below) instead of routing through the singleton + # ``_save_codex_tokens`` save path. # The singleton round-trip collapsed every added account into the # latest login: a second ``hermes auth add openai-codex`` overwrote # the first account's singleton-mirrored ``device_code`` entry rather @@ -349,19 +350,40 @@ def auth_add_command(args) -> None: timeout_seconds=getattr(args, "timeout", None) or 20.0, open_browser=not getattr(args, "no_browser", False), ) - auth_mod._save_xai_oauth_tokens( - creds["tokens"], - discovery=creds.get("discovery"), - redirect_uri=creds.get("redirect_uri", ""), + label = (getattr(args, "label", None) or "").strip() or label_from_token( + creds["tokens"]["access_token"], + _oauth_default_label(provider, len(pool.entries()) + 1), + ) + # Add a distinct, self-contained pool entry per account (matching the + # openai-codex / qwen-oauth / minimax-oauth patterns) instead of + # routing through the singleton ``_save_xai_oauth_tokens`` save path. + # The singleton round-trip collapsed every added account into the + # latest login: a second ``hermes auth add xai-oauth`` overwrote the + # first account's singleton-mirrored ``device_code`` entry rather than + # creating an independent one. ``manual:device_code`` entries refresh + # from their own token pair (``_sync_xai_oauth_entry_from_auth_store`` + # only adopts the singleton for ``source=="device_code"``), so they + # need no singleton shadow. + entry = PooledCredential( + provider=provider, + id=uuid.uuid4().hex[:6], + label=label, + auth_type=AUTH_TYPE_OAUTH, + priority=0, + source=SOURCE_MANUAL_DEVICE_CODE, + access_token=creds["tokens"]["access_token"], + refresh_token=creds["tokens"].get("refresh_token"), + base_url=creds.get("base_url") or auth_mod.DEFAULT_XAI_OAUTH_BASE_URL, last_refresh=creds.get("last_refresh"), - auth_mode="oauth_device_code", ) - pool = load_pool(provider) - entry = next((e for e in pool.entries() if getattr(e, "source", "") == "device_code"), None) - shown_label = entry.label if entry is not None else label_from_token( - creds["tokens"]["access_token"], _oauth_default_label(provider, 1) - ) - print(f'Saved {provider} OAuth credentials: "{shown_label}"') + first_credential = not pool.entries() + pool.add_entry(entry) + # Adding the first xAI credential should make it the active provider + # (the old singleton save path did this implicitly via + # _save_provider_state). Subsequent adds leave the active provider as-is. + if first_credential: + auth_mod.mark_provider_active_if_unset(provider) + print(f'Added {provider} OAuth credential #{len(pool.entries())}: "{entry.label}"') return if provider == "qwen-oauth": diff --git a/plugins/web/xai/provider.py b/plugins/web/xai/provider.py index 9e163a228d6f3..77d80a4398156 100644 --- a/plugins/web/xai/provider.py +++ b/plugins/web/xai/provider.py @@ -270,7 +270,10 @@ class XAIWebSearchProvider(WebSearchProvider): "refresh and retrying once.", ) try: - refreshed = resolve_xai_http_credentials(force_refresh=True) + refreshed = resolve_xai_http_credentials( + force_refresh=True, + api_key_hint=api_key, + ) refreshed_key = str(refreshed.get("api_key") or "").strip() if refreshed_key and refreshed_key != api_key: api_key = refreshed_key diff --git a/tests/agent/test_credential_pool.py b/tests/agent/test_credential_pool.py index 68a923355e597..3e76470fbace6 100644 --- a/tests/agent/test_credential_pool.py +++ b/tests/agent/test_credential_pool.py @@ -2970,6 +2970,78 @@ def test_xai_oauth_nonterminal_refresh_does_not_quarantine(tmp_path, monkeypatch assert tokens.get("refresh_token") == "old-refresh-token" +def test_xai_oauth_concurrent_pool_instances_refresh_single_use_token_once( + tmp_path, monkeypatch +): + import threading + import time + + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) + monkeypatch.delenv("XAI_API_KEY", raising=False) + monkeypatch.delenv("XAI_OAUTH_ACCESS_TOKEN", raising=False) + + _write_auth_store(tmp_path, { + "version": 1, + "providers": {}, + "credential_pool": { + "xai-oauth": [{ + "id": "manual-xai", + "label": "manual-xai", + "auth_type": "oauth", + "priority": 0, + "source": "manual:xai_pkce", + "access_token": "old-access-token", + "refresh_token": "one-time-refresh-token", + "base_url": "https://api.x.ai/v1", + }], + }, + }) + + from agent.credential_pool import load_pool + import hermes_cli.auth as auth_mod + + pools = [load_pool("xai-oauth"), load_pool("xai-oauth")] + start = threading.Barrier(2) + refresh_calls: list[tuple[str, str]] = [] + + def _refresh(access_token, refresh_token, **_kwargs): + refresh_calls.append((access_token, refresh_token)) + time.sleep(0.1) + return { + "access_token": "fresh-access-token", + "refresh_token": "fresh-refresh-token", + "last_refresh": "2026-07-12T00:00:00+00:00", + } + + monkeypatch.setattr(auth_mod, "refresh_xai_oauth_pure", _refresh) + results = [] + errors = [] + + def _worker(pool): + try: + start.wait() + results.append(pool.try_refresh_matching("old-access-token")) + except Exception as exc: + errors.append(exc) + + threads = [threading.Thread(target=_worker, args=(pool,)) for pool in pools] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert not errors + assert refresh_calls == [("old-access-token", "one-time-refresh-token")] + assert sorted(entry.access_token for entry in results) == [ + "fresh-access-token", + "fresh-access-token", + ] + persisted = json.loads((tmp_path / "hermes" / "auth.json").read_text()) + stored = persisted["credential_pool"]["xai-oauth"][0] + assert stored["access_token"] == "fresh-access-token" + assert stored["refresh_token"] == "fresh-refresh-token" + + # --------------------------------------------------------------------------- # Codex OAuth terminal error quarantine # --------------------------------------------------------------------------- diff --git a/tests/agent/test_credential_pool_routing.py b/tests/agent/test_credential_pool_routing.py index 9c7c047305151..a23c2280096b6 100644 --- a/tests/agent/test_credential_pool_routing.py +++ b/tests/agent/test_credential_pool_routing.py @@ -162,7 +162,7 @@ class TestPoolRotationCycle: # mark_exhausted_and_rotate returns next entry until exhausted self._rotation_index = 0 - def rotate(status_code=None, error_context=None): + def rotate(status_code=None, error_context=None, api_key_hint=None): self._rotation_index += 1 if self._rotation_index < pool_entries: return entries[self._rotation_index] @@ -220,7 +220,11 @@ class TestPoolRotationCycle: ) assert recovered is True assert has_retried is False - pool.mark_exhausted_and_rotate.assert_called_once_with(status_code=402, error_context=None) + pool.mark_exhausted_and_rotate.assert_called_once_with( + status_code=402, + error_context=None, + api_key_hint=None, + ) def test_no_pool_returns_false(self): """No pool should return (False, unchanged).""" diff --git a/tests/agent/test_error_classifier.py b/tests/agent/test_error_classifier.py index 616628967ecb4..6259cd825cc9b 100644 --- a/tests/agent/test_error_classifier.py +++ b/tests/agent/test_error_classifier.py @@ -261,6 +261,40 @@ class TestClassifyApiError: result = classify_api_error(e, provider="openrouter") assert result.reason == FailoverReason.billing + def test_xai_403_structured_spending_limit_code_classified_as_billing(self): + """xAI reports exhausted Grok credits as a provider-specific 403 code.""" + e = MockAPIError( + "Error code: 403", + status_code=403, + body={ + "code": "personal-team-blocked:spending-limit", + "error": ( + "You have run out of credits or need a Grok subscription. " + "Add credits at Grok or upgrade at Grok." + ), + }, + ) + + result = classify_api_error(e, provider="xai-oauth") + + assert result.reason == FailoverReason.billing + assert result.retryable is False + assert result.should_rotate_credential is True + assert result.should_fallback is True + + def test_non_xai_403_generic_billing_code_remains_auth(self): + """Do not broaden generic providers' historical structured-403 behavior.""" + e = MockAPIError( + "Error code: 403", + status_code=403, + body={"code": "insufficient_quota", "error": "Forbidden"}, + ) + + result = classify_api_error(e, provider="openrouter") + + assert result.reason == FailoverReason.auth + assert result.should_rotate_credential is False + # ── Billing ── def test_402_plain_billing(self): diff --git a/tests/hermes_cli/test_auth_commands.py b/tests/hermes_cli/test_auth_commands.py index f2e65dd6cd45b..02c2ac0f40c7b 100644 --- a/tests/hermes_cli/test_auth_commands.py +++ b/tests/hermes_cli/test_auth_commands.py @@ -509,12 +509,15 @@ def test_codex_runtime_pool_only_rate_limit_is_not_missing_auth(tmp_path, monkey def test_auth_add_xai_oauth_sets_active_provider(tmp_path, monkeypatch): - """hermes auth add xai-oauth must write providers singleton and set active_provider. + """hermes auth add xai-oauth must set active_provider and write a pool entry. - Previously pool.add_entry() was called directly, which wrote only the - credential-pool entry without setting active_provider. _model_section_has_credentials() - checks get_active_provider() first; with it unset, the setup wizard would - report "No inference provider configured" after a successful OAuth login. + Regression history: + - Early path called ``pool.add_entry()`` without ``active_provider``, so + the setup wizard reported "No inference provider configured". + - Intermediate path fixed that by routing through ``_save_xai_oauth_tokens`` + (singleton), which set active_provider but collapsed multi-account adds. + - Current path mirrors openai-codex: pool-only ``manual:device_code`` entry + plus ``mark_provider_active_if_unset`` on first add. """ monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) _write_auth_store(tmp_path, {"version": 1, "providers": {}}) @@ -549,15 +552,104 @@ def test_auth_add_xai_oauth_sets_active_provider(tmp_path, monkeypatch): auth_add_command(_Args()) payload = json.loads((tmp_path / "hermes" / "auth.json").read_text()) - # active_provider must be set — the core of this regression + # active_provider must be set — the core of the original regression assert payload["active_provider"] == "xai-oauth" - # providers singleton written by _save_xai_oauth_tokens - assert payload["providers"]["xai-oauth"]["tokens"]["access_token"] == access_token - assert payload["providers"]["xai-oauth"]["auth_mode"] == "oauth_device_code" - # pool seeded from singleton by _seed_from_singletons("xai-oauth") + # Pool-only multi-account path: no providers.xai-oauth singleton write + assert "xai-oauth" not in payload.get("providers", {}) entries = payload["credential_pool"]["xai-oauth"] - entry = next(item for item in entries if item["source"] == "device_code") + entry = next(item for item in entries if item["source"] == "manual:device_code") + assert entry["access_token"] == access_token assert entry["refresh_token"] == "xai-refresh-token" + assert entry["base_url"] == "https://api.x.ai/v1" + + +def test_auth_add_xai_oauth_keeps_distinct_pool_accounts(tmp_path, monkeypatch): + """Two ``hermes auth add xai-oauth`` runs must produce independent pool entries. + + Regression for the same collapse class as #39236 / #42316 for Codex: the + add path used to route through the singleton ``_save_xai_oauth_tokens`` + save, so the second login overwrote the first account's singleton-mirrored + ``device_code`` entry instead of adding a second independent one. + """ + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) + _write_auth_store(tmp_path, {"version": 1, "providers": {}}) + first_token = "xai-access-token-account-a" + second_token = "xai-access-token-account-b" + logins = iter( + [ + { + "tokens": { + "access_token": first_token, + "refresh_token": "first-xai-refresh", + "id_token": "", + "token_type": "Bearer", + }, + "discovery": {"token_endpoint": "https://auth.x.ai/token"}, + "redirect_uri": "", + "base_url": "https://api.x.ai/v1", + "last_refresh": "2026-07-10T10:00:00Z", + "source": "oauth-device-code", + }, + { + "tokens": { + "access_token": second_token, + "refresh_token": "second-xai-refresh", + "id_token": "", + "token_type": "Bearer", + }, + "discovery": {"token_endpoint": "https://auth.x.ai/token"}, + "redirect_uri": "", + "base_url": "https://api.x.ai/v1", + "last_refresh": "2026-07-10T10:05:00Z", + "source": "oauth-device-code", + }, + ] + ) + monkeypatch.setattr( + "hermes_cli.auth._xai_oauth_device_code_login", + lambda **kwargs: next(logins), + ) + + from hermes_cli.auth_commands import auth_add_command + from agent.credential_pool import load_pool + + class _Args: + provider = "xai-oauth" + auth_type = "oauth" + api_key = None + label = None + timeout = None + no_browser = False + + # Distinct labels so order is unambiguous even without JWT email claims. + class _ArgsA(_Args): + label = "xai-heavy" + + class _ArgsB(_Args): + label = "xai-premium" + + auth_add_command(_ArgsA()) + auth_add_command(_ArgsB()) + + pool = load_pool("xai-oauth") + entries = pool.entries() + + assert [entry.source for entry in entries] == [ + "manual:device_code", + "manual:device_code", + ] + assert [entry.label for entry in entries] == ["xai-heavy", "xai-premium"] + assert [entry.access_token for entry in entries] == [first_token, second_token] + assert [entry.refresh_token for entry in entries] == [ + "first-xai-refresh", + "second-xai-refresh", + ] + + payload = json.loads((tmp_path / "hermes" / "auth.json").read_text()) + # No singleton block — the add path is now pool-only. + assert "xai-oauth" not in payload.get("providers", {}) + # First add activated the provider; second add left it as-is. + assert payload["active_provider"] == "xai-oauth" def test_auth_remove_reindexes_priorities(tmp_path, monkeypatch): diff --git a/tests/run_agent/test_codex_xai_oauth_recovery.py b/tests/run_agent/test_codex_xai_oauth_recovery.py index 2bc31686e75e3..33d287c2b2fa5 100644 --- a/tests/run_agent/test_codex_xai_oauth_recovery.py +++ b/tests/run_agent/test_codex_xai_oauth_recovery.py @@ -603,6 +603,69 @@ def test_recover_with_credential_pool_skips_refresh_on_entitlement_403(): assert refresh_calls["n"] == 0, "try_refresh_current must NOT be called on entitlement 403" +def test_recover_with_credential_pool_rotates_on_xai_spending_limit_403(): + """xAI's explicit spending-limit 403 must rotate, not hit the entitlement guard.""" + from agent.error_classifier import FailoverReason, classify_api_error + + agent = _make_codex_agent() + next_entry = MagicMock(id="healthy-account") + refresh_calls = {"n": 0} + + class _SpendingLimitError(Exception): + status_code = 403 + body = { + "code": "personal-team-blocked:spending-limit", + "error": ( + "You have run out of credits or need a Grok subscription. " + "Add credits at Grok or upgrade at Grok." + ), + } + + class _FakePool: + provider = "xai-oauth" + + def try_refresh_current(self): + refresh_calls["n"] += 1 + return MagicMock(id="should_not_be_called") + + def mark_exhausted_and_rotate( + self, + *, + status_code, + error_context=None, + api_key_hint=None, + ): + assert status_code == 403 + assert api_key_hint == "test-key" + assert error_context == { + "reason": "personal-team-blocked:spending-limit", + "message": ( + "You have run out of credits or need a Grok subscription. " + "Add credits at Grok or upgrade at Grok." + ), + } + return next_entry + + error = _SpendingLimitError("Error code: 403") + classified = classify_api_error(error, provider="xai-oauth", model="grok-4.5") + error_context = agent._extract_api_error_context(error) + setattr(agent, "_credential_pool", _FakePool()) + agent._swap_credential = MagicMock() + + recovered, retried_429 = agent._recover_with_credential_pool( + status_code=error.status_code, + has_retried_429=False, + classified_reason=classified.reason, + error_context=error_context, + ) + + assert classified.reason == FailoverReason.billing + assert recovered is True + assert retried_429 is False + assert refresh_calls["n"] == 0 + agent._swap_credential.assert_called_once_with(next_entry) + + def test_recover_with_credential_pool_skips_refresh_on_bare_403_for_xai_oauth(): """A bare HTTP 403 from ``xai-oauth`` (no keyword match) must NOT loop refresh. diff --git a/tests/run_agent/test_run_agent.py b/tests/run_agent/test_run_agent.py index c218e75564ee6..2b0db24f6b733 100644 --- a/tests/run_agent/test_run_agent.py +++ b/tests/run_agent/test_run_agent.py @@ -6185,9 +6185,16 @@ class TestCredentialPoolRecovery: def current(self): return current - def mark_exhausted_and_rotate(self, *, status_code, error_context=None): + def mark_exhausted_and_rotate( + self, + *, + status_code, + error_context=None, + api_key_hint=None, + ): assert status_code == 402 assert error_context is None + assert api_key_hint == agent.api_key return next_entry agent._credential_pool = _Pool() @@ -6206,9 +6213,16 @@ class TestCredentialPoolRecovery: next_entry = SimpleNamespace(label="secondary") class _Pool: - def mark_exhausted_and_rotate(self, *, status_code, error_context=None): + def mark_exhausted_and_rotate( + self, + *, + status_code, + error_context=None, + api_key_hint=None, + ): assert status_code == 400 assert error_context == {"reason": "out_of_extra_usage"} + assert api_key_hint == agent.api_key return next_entry agent._credential_pool = _Pool() diff --git a/tests/tools/test_web_providers_xai.py b/tests/tools/test_web_providers_xai.py index 2a6f0c63b81c4..9a5b00fe9b27e 100644 --- a/tests/tools/test_web_providers_xai.py +++ b/tests/tools/test_web_providers_xai.py @@ -515,9 +515,10 @@ class TestXAIProviderSearchErrors: raise unauthorized return _mock_resp(_responses_payload(json.dumps({"results": []}))) - def fake_resolve(*, force_refresh=False): + def fake_resolve(*, force_refresh=False, api_key_hint=None): if force_refresh: calls["refresh_count"] += 1 + assert api_key_hint == "stale-token" return { "provider": "xai-oauth", "api_key": "fresh-after-refresh", @@ -554,11 +555,11 @@ class TestXAIProviderSearchErrors: calls["posts"] += 1 raise unauthorized - def fake_resolve(*, force_refresh=False): + def fake_resolve(*, force_refresh=False, api_key_hint=None): if force_refresh: calls["refreshed"] = True # provider=="xai" signals env-var path; retry must be skipped. - return {"provider": "xai", "api_key": "sk-env-var-key", "base_url": "https://api.x.ai/v1"} + return {"provider": "xai", "api_key": "«redacted:sk-…»", "base_url": "https://api.x.ai/v1"} with patch.object(xai_provider, "resolve_xai_http_credentials", side_effect=fake_resolve), \ patch.object(xai_provider, "_load_xai_web_config", return_value={}), \ @@ -587,9 +588,10 @@ class TestXAIProviderSearchErrors: calls["posts"] += 1 raise unauthorized - def fake_resolve(*, force_refresh=False): + def fake_resolve(*, force_refresh=False, api_key_hint=None): if force_refresh: calls["refresh_count"] += 1 + assert api_key_hint == "same-dead-token" return { "provider": "xai-oauth", "api_key": "same-dead-token", @@ -624,7 +626,7 @@ class TestXAIProviderSearchErrors: calls["posts"] += 1 raise err - def fake_resolve(*, force_refresh=False): + def fake_resolve(*, force_refresh=False, api_key_hint=None): if force_refresh: calls["refreshed"] = True return {"provider": "xai-oauth", "api_key": "tok", "base_url": "https://api.x.ai/v1"} @@ -727,25 +729,30 @@ class TestXAIBackendWiring: class TestXAIProviderOAuthPath: """Verifies the provider works when credentials come from the OAuth - runtime resolver (``hermes auth`` sign-in) rather than an env-var key. - Patches at the ``hermes_cli.runtime_provider.resolve_runtime_provider`` - boundary so the full ``tools.xai_http.resolve_xai_http_credentials`` - chain is exercised end-to-end. + credential pool (``hermes auth`` sign-in) rather than an env-var key. + The full ``tools.xai_http.resolve_xai_http_credentials`` chain is exercised + against a temporary auth store. """ - def test_search_uses_oauth_bearer_token_and_base_url(self, monkeypatch): + def test_search_uses_oauth_bearer_token_and_base_url(self, monkeypatch, tmp_path): from plugins.web.xai import provider as xai_provider # Force the env-var fallback to fail so resolution must go via OAuth. monkeypatch.delenv("XAI_API_KEY", raising=False) - - oauth_runtime = { - "provider": "xai-oauth", - "api_mode": "codex_responses", - "base_url": "https://api.x.ai/v1", - "api_key": "ya29.fake-oauth-access-token", - "source": "hermes-auth-store", - } + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + monkeypatch.setenv("HERMES_XAI_BASE_URL", "https://proxy.x.ai/v1/") + (tmp_path / "auth.json").write_text(json.dumps({ + "version": 1, + "active_provider": "xai-oauth", + "providers": { + "xai-oauth": { + "tokens": { + "access_token": "ya29.fake-oauth-access-token", + "refresh_token": "fake-oauth-refresh-token", + }, + }, + }, + })) captured: dict = {} @@ -754,13 +761,95 @@ class TestXAIProviderOAuthPath: captured["headers"] = kwargs.get("headers", {}) return _mock_resp(_responses_payload(json.dumps({"results": []}))) - with patch( - "hermes_cli.runtime_provider.resolve_runtime_provider", - return_value=oauth_runtime, - ), patch.object(xai_provider, "_load_xai_web_config", return_value={}), \ + with patch.object(xai_provider, "_load_xai_web_config", return_value={}), \ patch("httpx.post", side_effect=fake_post): result = xai_provider.XAIWebSearchProvider().search("q", limit=3) assert result["success"] is True - assert captured["url"] == "https://api.x.ai/v1/responses" + assert captured["url"] == "https://proxy.x.ai/v1/responses" assert captured["headers"].get("Authorization") == "Bearer ya29.fake-oauth-access-token" + + def test_pool_only_direct_refresh_updates_main_runtime(self, monkeypatch, tmp_path): + """A direct 401 refresh must rotate the exact manual pool row. + + xAI refresh tokens are one-time credentials. Persisting the refreshed + pair only to ``providers.xai-oauth`` leaves the manual row stale and + breaks the next main-runtime load. + """ + from hermes_cli.runtime_provider import resolve_runtime_provider + from tools.xai_http import resolve_xai_http_credentials + + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + monkeypatch.delenv("XAI_API_KEY", raising=False) + monkeypatch.delenv("XAI_OAUTH_ACCESS_TOKEN", raising=False) + (tmp_path / "config.yaml").write_text( + "credential_pool_strategies:\n xai-oauth: round_robin\n" + ) + auth_path = tmp_path / "auth.json" + auth_path.write_text(json.dumps({ + "version": 1, + "active_provider": "xai-oauth", + "providers": {}, + "credential_pool": { + "xai-oauth": [{ + "id": "manual-xai", + "label": "pool-only", + "auth_type": "oauth", + "priority": 0, + "source": "manual:device_code", + "access_token": "rejected-access", + "refresh_token": "one-time-refresh", + "base_url": "https://api.x.ai/v1", + }, { + "id": "backup-xai", + "label": "backup", + "auth_type": "oauth", + "priority": 1, + "source": "manual:device_code", + "access_token": "backup-access", + "refresh_token": "backup-refresh", + "base_url": "https://api.x.ai/v1", + }], + }, + })) + + refresh_calls = [] + + def fake_refresh(access_token, refresh_token, **_kwargs): + refresh_calls.append((access_token, refresh_token)) + return { + "access_token": "fresh-access", + "refresh_token": "rotated-refresh", + "last_refresh": "2026-07-12T12:00:00Z", + } + + monkeypatch.setattr( + "hermes_cli.auth.refresh_xai_oauth_pure", + fake_refresh, + ) + + initial = resolve_xai_http_credentials() + assert initial["api_key"] == "rejected-access" + + refreshed = resolve_xai_http_credentials( + force_refresh=True, + api_key_hint=initial["api_key"], + ) + assert refreshed["api_key"] == "fresh-access" + assert refresh_calls == [("rejected-access", "one-time-refresh")] + + stored = json.loads(auth_path.read_text()) + entries = { + item["id"]: item + for item in stored["credential_pool"]["xai-oauth"] + } + entry = entries["manual-xai"] + assert entry["access_token"] == "fresh-access" + assert entry["refresh_token"] == "rotated-refresh" + assert entries["backup-xai"]["access_token"] == "backup-access" + assert entries["backup-xai"]["refresh_token"] == "backup-refresh" + assert "xai-oauth" not in stored.get("providers", {}) + + runtime = resolve_runtime_provider(requested="xai-oauth") + assert runtime["api_key"] == "backup-access" + assert refresh_calls == [("rejected-access", "one-time-refresh")] diff --git a/tools/xai_http.py b/tools/xai_http.py index fb1f523175f0f..8d80ae1bbd47f 100644 --- a/tools/xai_http.py +++ b/tools/xai_http.py @@ -28,6 +28,9 @@ def has_xai_credentials() -> bool: 1. ``XAI_API_KEY`` env var (cheapest; covers explicit-key users). 2. ``~/.hermes/auth.json`` has a non-empty ``providers.xai-oauth.tokens.access_token`` (single file read, no expiry check, no refresh). + 3. ``credential_pool.xai-oauth`` has any entry with a non-empty + ``access_token`` (covers multi-account ``hermes auth add xai-oauth`` + grants that are pool-only / ``manual:device_code``). Returns False on any exception so a corrupted auth store can't block other availability scans. Truthful refresh + expiry handling happens @@ -46,7 +49,23 @@ def has_xai_credentials() -> bool: xai_state = providers.get("xai-oauth") if isinstance(providers, dict) else None tokens = xai_state.get("tokens") if isinstance(xai_state, dict) else None access_token = tokens.get("access_token") if isinstance(tokens, dict) else None - return bool(str(access_token or "").strip()) + if str(access_token or "").strip(): + return True + # Pool-only grants (multi-account ``auth add``) never write the + # providers singleton; still count as present credentials. + credential_pool = store.get("credential_pool") if isinstance(store, dict) else None + entries = ( + credential_pool.get("xai-oauth") + if isinstance(credential_pool, dict) + else None + ) + if isinstance(entries, list): + for entry in entries: + if not isinstance(entry, dict): + continue + if str(entry.get("access_token", "") or "").strip(): + return True + return False except Exception: return False @@ -221,7 +240,11 @@ def maybe_mark_xai_storage_notice_seen(section_name: str) -> Optional[str]: return notice -def resolve_xai_http_credentials(*, force_refresh: bool = False) -> Dict[str, str]: +def resolve_xai_http_credentials( + *, + force_refresh: bool = False, + api_key_hint: Optional[str] = None, +) -> Dict[str, str]: """Resolve bearer credentials for direct xAI HTTP endpoints. Prefers Hermes-managed xAI OAuth credentials when available, then falls back @@ -231,43 +254,53 @@ def resolve_xai_http_credentials(*, force_refresh: bool = False) -> Dict[str, st endpoints (images, TTS, STT, etc.) aligned with the main runtime auth model and preserves the regression contract from PR #17140 / #17163. - Set ``force_refresh=True`` to bypass the resolver's JWT-exp shortcut and - perform an unconditional OAuth refresh. Callers should use this only as a - reactive remediation after a server 401 (mid-window revocation, opaque - tokens where the proactive JWT check is a no-op, etc.), not as a default — - the auth-store lock is held for the duration of the refresh. + Set ``force_refresh=True`` to perform an unconditional OAuth refresh. + Reactive callers should also pass the rejected bearer as ``api_key_hint`` + so a freshly loaded multi-account pool refreshes the exact issuing entry, + not whichever entry its strategy would otherwise select first. """ try: - from hermes_cli.auth import resolve_xai_oauth_runtime_credentials + from agent.credential_pool import load_pool + import hermes_cli.auth as auth_mod - creds = resolve_xai_oauth_runtime_credentials(force_refresh=force_refresh) - access_token = str(creds.get("api_key") or "").strip() - base_url = str(creds.get("base_url") or "").strip().rstrip("/") + pool = load_pool("xai-oauth") + entry = ( + pool.try_refresh_matching(api_key_hint) + if force_refresh + else pool.select() + ) + if force_refresh and entry is None: + # A rejected refresh may quarantine the issuing entry. Continue + # with the next healthy account instead of falling back to the raw + # singleton resolver and resurrecting the stale pool row. + entry = pool.select() + access_token = str( + getattr(entry, "runtime_api_key", None) + or getattr(entry, "access_token", "") + ).strip() + fallback_base_url = str( + getattr(entry, "runtime_base_url", None) + or getattr(entry, "base_url", "") + or auth_mod.DEFAULT_XAI_OAUTH_BASE_URL + ).strip().rstrip("/") + override_base_url = str( + get_env_value("HERMES_XAI_BASE_URL") + or get_env_value("XAI_BASE_URL") + or "" + ).strip().rstrip("/") + base_url = auth_mod._xai_validate_inference_base_url( + override_base_url, + fallback=fallback_base_url, + ) if access_token: return { "provider": "xai-oauth", "api_key": access_token, - "base_url": base_url or "https://api.x.ai/v1", + "base_url": base_url, } except Exception: pass - if not force_refresh: - try: - from hermes_cli.runtime_provider import resolve_runtime_provider - - runtime = resolve_runtime_provider(requested="xai-oauth") - access_token = str(runtime.get("api_key") or "").strip() - base_url = str(runtime.get("base_url") or "").strip().rstrip("/") - if access_token: - return { - "provider": "xai-oauth", - "api_key": access_token, - "base_url": base_url or "https://api.x.ai/v1", - } - except Exception: - pass - api_key = str(get_env_value("XAI_API_KEY") or "").strip() base_url = str(get_env_value("XAI_BASE_URL") or "https://api.x.ai/v1").strip().rstrip("/") return {