diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index 84beba810c..fab406a3eb 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -36,7 +36,7 @@ from hermes_cli.timeouts import get_provider_request_timeout from agent.prompt_builder import format_steer_marker from agent.tool_dispatch_helpers import _trajectory_normalize_msg, make_tool_result_message from agent.trajectory import convert_scratchpad_to_think -from agent.credential_pool import STATUS_EXHAUSTED +from agent.credential_pool import STATUS_EXHAUSTED, credential_pool_matches_provider from agent.error_classifier import FailoverReason from agent.turn_context import drop_stale_api_content from utils import base_url_host_matches, base_url_hostname, env_var_enabled, atomic_json_write @@ -1464,6 +1464,64 @@ def restore_primary_runtime(agent) -> bool: if getattr(agent, "_rate_limited_until", 0) > time.monotonic(): return False # primary still in rate-limit cooldown, stay on fallback + # ── Reset-aware gate ── + # The 60s ``_rate_limited_until`` cooldown covers transient rate limits, + # but subscription-style providers (Claude Pro/Max 5-hour windows, ChatGPT + # weekly limits) report reset times hours or days away. The credential + # pool already stores those timestamps (``last_error_reset_at``); until + # the earliest one elapses, every restore attempt is a *guaranteed* + # failure that costs two prompt-cache invalidations per turn (switch to + # primary, fail, switch back to fallback) and re-marshals the full + # context each way. Skip the restore while the pool says nobody can + # serve, and come back the moment the reset time passes. + # + # Fail-open by design: any error (unreadable auth store, legacy pool + # adapter without ``next_available_at``) falls through to the existing + # every-turn retry. A pool with no reset info returns ``None`` and also + # falls through — this gate only ever *adds* skips for provably + # limited windows, so recovery can never be later than it is today. + # + # When the attached pool belongs to the fallback provider (cross-provider + # fallback rebinds it), the primary pool is loaded here and handed to the + # pool-rebind block below via ``prefetched_primary_pool`` so the load + # happens at most once per restore. + prefetched_primary_pool = None + try: + primary_provider = str( + (agent._primary_runtime or {}).get("provider") or "" + ).strip().lower() + pool = getattr(agent, "_credential_pool", None) + if not credential_pool_matches_provider( + pool, + primary_provider, + base_url=str((agent._primary_runtime or {}).get("base_url") or ""), + ): + from agent.credential_pool import load_pool + + prefetched_primary_pool = ( + load_pool(primary_provider) if primary_provider else None + ) + pool = prefetched_primary_pool + next_at = getattr(pool, "next_available_at", lambda: None)() + if next_at is not None and next_at > time.time(): + if not getattr(agent, "_restore_wait_logged", False): + agent._restore_wait_logged = True + logger.info( + "Primary %s rate-limited until %s; staying on fallback " + "%s/%s until the reset elapses", + primary_provider or "?", + datetime.fromtimestamp(next_at).isoformat(timespec="seconds"), + agent.provider, + agent.model, + ) + return False + except Exception: + logger.debug( + "Reset-aware restore gate failed; falling back to per-turn retry", + exc_info=True, + ) + agent._restore_wait_logged = False + rt = agent._primary_runtime try: # ── Core runtime state ── @@ -1557,9 +1615,14 @@ def restore_primary_runtime(agent) -> bool: agent._credential_pool = None agent._credential_pool_entry_id = None try: - from agent.credential_pool import load_pool + if prefetched_primary_pool is not None: + # Reuse the pool the reset-aware gate already loaded for + # this restore — avoids a second disk read of auth.json. + agent._credential_pool = prefetched_primary_pool + else: + from agent.credential_pool import load_pool - agent._credential_pool = load_pool(primary_provider) + agent._credential_pool = load_pool(primary_provider) except Exception as exc: logger.warning( "Restore could not reload primary credential pool for %s: %s", diff --git a/agent/credential_pool.py b/agent/credential_pool.py index 4beb3492b0..1ebc2a3b73 100644 --- a/agent/credential_pool.py +++ b/agent/credential_pool.py @@ -620,6 +620,30 @@ class CredentialPool: with self._lock: return bool(self._available_entries()) + def next_available_at(self) -> Optional[float]: + """Earliest epoch time (seconds) any entry re-enters rotation. + + Returns ``None`` when at least one entry is available right now, or + when no exhausted entry carries a usable recovery time (empty pool, + or only ``STATUS_DEAD`` entries, which never re-enter via TTL). + Callers must treat ``None`` as "no wait information", not + "unavailable". + + Like :meth:`has_available`, expired cooldowns are left uncleared + (``clear_expired=False``); the only writes are the same + re-auth/token sync paths ``has_available`` already performs. + """ + if self._available_entries(): + return None + candidates: List[float] = [] + for entry in self._entries: + if entry.last_status != STATUS_EXHAUSTED: + continue + until = _exhausted_until(entry) + if until is not None: + candidates.append(until) + return min(candidates) if candidates else None + def entries(self) -> List[PooledCredential]: with self._lock: return list(self._entries) diff --git a/tests/run_agent/test_reset_aware_primary_restore.py b/tests/run_agent/test_reset_aware_primary_restore.py new file mode 100644 index 0000000000..7fad82e686 --- /dev/null +++ b/tests/run_agent/test_reset_aware_primary_restore.py @@ -0,0 +1,307 @@ +"""Reset-aware primary restore — stay on fallback until the primary's +rate-limit window actually resets. + +``restore_primary_runtime`` retries the primary at the top of every turn +once the 60s ``_rate_limited_until`` cooldown clears. For transient 429s +that is correct, but subscription-window limits (Claude Pro/Max 5-hour +windows, ChatGPT weekly caps) report reset times hours or days away. The +credential pool already knows that timestamp (``last_error_reset_at``), +and until it elapses every restore attempt is a guaranteed failure that +invalidates the prompt cache twice per turn (primary → fail → fallback). + +The gate must FAIL OPEN: no pool, no reset info, or any error in the gate +falls through to the existing per-turn retry, so recovery can never happen +later than it does today. +""" + +import time +from unittest.mock import MagicMock, patch + +from run_agent import AIAgent +from agent.credential_pool import ( + STATUS_DEAD, + STATUS_EXHAUSTED, + STATUS_OK, + CredentialPool, + PooledCredential, +) + + +# ============================================================================= +# Helpers +# ============================================================================= + +def _entry( + provider="openrouter", + id="cred-1", + status=STATUS_OK, + reset_at=None, + status_at=None, + error_code=None, +): + return PooledCredential( + provider=provider, + id=id, + label=f"label-{id}", + auth_type="api_key", + priority=1, + source="manual", + access_token="sk-test-1234567890", + last_status=status, + last_status_at=status_at, + last_error_code=error_code, + last_error_reset_at=reset_at, + ) + + +class _FakePool: + """Minimal stand-in for CredentialPool in agent-level tests.""" + + def __init__(self, provider, next_at=None, available=False, raise_on_next=False): + self.provider = provider + self._next_at = next_at + self._available = available + self._raise = raise_on_next + self.next_available_calls = 0 + + def next_available_at(self): + self.next_available_calls += 1 + if self._raise: + raise RuntimeError("boom") + return self._next_at + + def has_credentials(self): + return True + + def has_available(self): + return self._available + + def select(self): + return None + + +def _make_tool_defs(*names): + return [ + { + "type": "function", + "function": { + "name": n, + "description": f"{n} tool", + "parameters": {"type": "object", "properties": {}}, + }, + } + for n in names + ] + + +def _make_agent(fallback_model=None): + with ( + patch("run_agent.get_tool_definitions", return_value=_make_tool_defs("web_search")), + patch("run_agent.check_toolset_requirements", return_value={}), + patch("run_agent.OpenAI"), + ): + agent = AIAgent( + api_key="test-key-12345678", + base_url="https://my-llm.example.com/v1", + provider="custom", + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + fallback_model=fallback_model, + ) + agent.client = MagicMock() + return agent + + +def _activate_fallback(agent): + mock_client = MagicMock() + mock_client.api_key = "fallback-key-1234" + mock_client.base_url = "https://openrouter.ai/api/v1" + with patch( + "agent.auxiliary_client.resolve_provider_client", + return_value=(mock_client, None), + ): + assert agent._try_activate_fallback() is True + assert agent._fallback_activated is True + + +# ============================================================================= +# CredentialPool.next_available_at() +# ============================================================================= + +class TestNextAvailableAt: + def test_all_exhausted_returns_earliest_reset(self): + now = time.time() + pool = CredentialPool( + "openrouter", + [ + _entry(id="a", status=STATUS_EXHAUSTED, reset_at=now + 7200, error_code=429), + _entry(id="b", status=STATUS_EXHAUSTED, reset_at=now + 3600, error_code=429), + ], + ) + assert pool.next_available_at() == now + 3600 + + def test_available_entry_returns_none(self): + now = time.time() + pool = CredentialPool( + "openrouter", + [ + _entry(id="a", status=STATUS_OK), + _entry(id="b", status=STATUS_EXHAUSTED, reset_at=now + 3600, error_code=429), + ], + ) + assert pool.next_available_at() is None + + def test_elapsed_cooldown_counts_as_available(self): + """An exhausted entry whose reset time has passed re-enters rotation, + so the pool reports available (None) even without clear_expired.""" + now = time.time() + pool = CredentialPool( + "openrouter", + [_entry(id="a", status=STATUS_EXHAUSTED, reset_at=now - 10, error_code=429)], + ) + assert pool.next_available_at() is None + + def test_exhausted_without_timestamps_returns_none(self): + """No reset info at all -> None (fail open), not a guess.""" + pool = CredentialPool( + "openrouter", + [_entry(id="a", status=STATUS_EXHAUSTED, reset_at=None, status_at=None)], + ) + assert pool.next_available_at() is None + + def test_exhausted_with_status_at_uses_ttl(self): + """Without an explicit reset_at, last_status_at + TTL is the estimate.""" + now = time.time() + pool = CredentialPool( + "openrouter", + [_entry(id="a", status=STATUS_EXHAUSTED, status_at=now, error_code=429)], + ) + result = pool.next_available_at() + assert result is not None + assert result > now + + def test_dead_only_returns_none(self): + """DEAD entries never re-enter via TTL; report no wait info.""" + pool = CredentialPool( + "openrouter", + [_entry(id="a", status=STATUS_DEAD, status_at=time.time())], + ) + assert pool.next_available_at() is None + + def test_empty_pool_returns_none(self): + pool = CredentialPool("openrouter", []) + assert pool.next_available_at() is None + + +# ============================================================================= +# restore_primary_runtime() gate +# ============================================================================= + +class TestResetAwareRestoreGate: + FB = {"provider": "openrouter", "model": "anthropic/claude-sonnet-4"} + + def test_stays_on_fallback_until_reset(self): + agent = _make_agent(fallback_model=self.FB) + _activate_fallback(agent) + agent._rate_limited_until = 0 # 60s transient cooldown already cleared + + # Attached pool matches the primary provider and says: nobody can + # serve until an hour from now. + agent._credential_pool = _FakePool("custom", next_at=time.time() + 3600) + + assert agent._restore_primary_runtime() is False + assert agent._fallback_activated is True + assert agent.provider == "openrouter" + assert agent.model == "anthropic/claude-sonnet-4" + + def test_restores_once_reset_elapsed(self): + agent = _make_agent(fallback_model=self.FB) + original_model = agent.model + _activate_fallback(agent) + agent._rate_limited_until = 0 + + agent._credential_pool = _FakePool("custom", next_at=None) + + with patch("run_agent.OpenAI", return_value=MagicMock()): + assert agent._restore_primary_runtime() is True + assert agent._fallback_activated is False + assert agent.model == original_model + assert agent.provider == "custom" + + def test_past_reset_time_does_not_block(self): + agent = _make_agent(fallback_model=self.FB) + _activate_fallback(agent) + agent._rate_limited_until = 0 + + agent._credential_pool = _FakePool("custom", next_at=time.time() - 5) + + with patch("run_agent.OpenAI", return_value=MagicMock()): + assert agent._restore_primary_runtime() is True + assert agent._fallback_activated is False + + def test_fails_open_on_pool_error(self): + """Any exception inside the gate must not break restore.""" + agent = _make_agent(fallback_model=self.FB) + _activate_fallback(agent) + agent._rate_limited_until = 0 + + agent._credential_pool = _FakePool("custom", raise_on_next=True) + + with patch("run_agent.OpenAI", return_value=MagicMock()): + assert agent._restore_primary_runtime() is True + assert agent._fallback_activated is False + + def test_cross_provider_fallback_loads_primary_pool(self): + """After a cross-provider fallback the attached pool belongs to the + fallback provider; the gate must consult the PRIMARY's pool.""" + agent = _make_agent(fallback_model=self.FB) + _activate_fallback(agent) + agent._rate_limited_until = 0 + + # Attached pool is the fallback provider's (mismatch with "custom"). + agent._credential_pool = _FakePool("openrouter", next_at=None) + primary_pool = _FakePool("custom", next_at=time.time() + 3600) + + with patch("agent.credential_pool.load_pool", return_value=primary_pool) as lp: + assert agent._restore_primary_runtime() is False + assert any(c.args == ("custom",) for c in lp.call_args_list) + assert primary_pool.next_available_calls == 1 + assert agent._fallback_activated is True + + def test_no_pool_info_falls_through(self): + """Pool present but no reset info -> existing per-turn retry.""" + agent = _make_agent(fallback_model=self.FB) + _activate_fallback(agent) + agent._rate_limited_until = 0 + + agent._credential_pool = _FakePool("custom", next_at=None) + + with patch("run_agent.OpenAI", return_value=MagicMock()): + assert agent._restore_primary_runtime() is True + + def test_logs_wait_only_once(self, caplog): + import logging + + agent = _make_agent(fallback_model=self.FB) + _activate_fallback(agent) + agent._rate_limited_until = 0 + agent._credential_pool = _FakePool("custom", next_at=time.time() + 3600) + + with caplog.at_level(logging.INFO, logger="agent.agent_runtime_helpers"): + assert agent._restore_primary_runtime() is False + assert agent._restore_primary_runtime() is False + waits = [r for r in caplog.records if "staying on fallback" in r.getMessage()] + assert len(waits) == 1 + + def test_transient_cooldown_still_respected(self): + """The existing 60s monotonic gate fires before the reset-aware one.""" + agent = _make_agent(fallback_model=self.FB) + _activate_fallback(agent) + agent._rate_limited_until = time.monotonic() + 60 + pool = _FakePool("custom", next_at=None) + agent._credential_pool = pool + + assert agent._restore_primary_runtime() is False + # Reset-aware gate never consulted — short-circuited by the 60s gate. + assert pool.next_available_calls == 0 diff --git a/website/docs/user-guide/features/fallback-providers.md b/website/docs/user-guide/features/fallback-providers.md index 3d2dde7820..b44b882f59 100644 --- a/website/docs/user-guide/features/fallback-providers.md +++ b/website/docs/user-guide/features/fallback-providers.md @@ -123,6 +123,8 @@ Prompt caches are keyed to the model (and on most providers, the account) servin :::info Per-Turn, Not Per-Session Fallback is **turn-scoped**: each new user message starts with the primary model restored. If the primary fails mid-turn, fallback activates for that turn only. On the next message, Hermes tries the primary again. Within a single turn, fallback activates at most once — if the fallback also fails, normal error handling takes over (retries, then error message). This prevents cascading failover loops within a turn while giving the primary model a fresh chance every turn. + +The per-turn retry is **reset-aware**: when the primary's credentials report a rate-limit reset time that hasn't elapsed yet (subscription windows like Claude Pro/Max's 5-hour blocks or Codex weekly limits report these as hours or days), Hermes skips the doomed retry and stays on the fallback until the reset passes — avoiding two pointless provider switches (and two prompt-cache invalidations) per turn. The moment the reset time elapses, the next turn goes back to the primary automatically. Transient 429s without a reset time keep the existing behavior: a short cooldown, then retry every turn. ::: ### Examples