diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index 5cc10079d09f..606edbad7836 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -922,12 +922,17 @@ def recover_with_credential_pool( ) return False, has_retried_429 - # Capture the current API key before any rotation — needed to - # identify which credential actually failed when - # mark_exhausted_and_rotate is called. Without this hint the - # pool falls back to current() or _select_unlocked(), which may - # return the NEXT (healthy) entry after a prior rotation, marking - # the wrong credential as exhausted (#43747). + # Attribute the failure to the API key the agent actually dispatched the + # request with, not to pool.current(). The current() pointer is shared, + # mutable state — round-robin select() advances it on every call, and + # concurrent turns or a second process (gateway/dashboard) reloading the + # pool reset it to None — so by the time recovery runs it routinely points + # at a DIFFERENT, healthy entry. Marking that entry exhausted copies this + # request's error/reset time onto it and can take the whole pool offline + # from a single rate-limited key (#43747). ``_swap_credential`` keeps + # ``agent.api_key`` in sync with the entry in use, so it identifies the + # failing entry exactly; fall back to current()'s key only when the agent + # carries no key at all. _api_key_hint = getattr(agent, "api_key", None) or None if not _api_key_hint: _cur = pool.current() @@ -987,7 +992,16 @@ def recover_with_credential_pool( # rotate immediately. This prevents the "cancel-between-429s" trap # where has_retried_429 (a local var) gets reset on each new prompt, # causing the pool to retry the same exhausted credential forever. - current_entry = pool.current() + # Prefer the entry matching the failing key over the shared current() + # pointer, for the same attribution reason as above. + current_entry = None + if _api_key_hint: + current_entry = next( + (e for e in pool.entries() if e.runtime_api_key == _api_key_hint), + None, + ) + if current_entry is None: + current_entry = pool.current() current_last_status = getattr(current_entry, "last_status", None) if current_entry else None if current_last_status == STATUS_EXHAUSTED: _ra().logger.info( @@ -995,7 +1009,11 @@ def recover_with_credential_pool( current_last_status, ) rotate_status = status_code if status_code is not None else 429 - next_entry = pool.mark_exhausted_and_rotate(status_code=rotate_status, error_context=error_context, api_key_hint=_api_key_hint) + next_entry = pool.mark_exhausted_and_rotate( + status_code=rotate_status, + error_context=error_context, + api_key_hint=_api_key_hint, + ) if next_entry is not None: _ra().logger.info( "Credential %s (rate limit, pre-exhausted) — rotated to pool entry %s", @@ -1019,7 +1037,11 @@ def recover_with_credential_pool( if not has_retried_429 and not usage_limit_reached: return False, True rotate_status = status_code if status_code is not None else 429 - next_entry = pool.mark_exhausted_and_rotate(status_code=rotate_status, error_context=error_context, api_key_hint=_api_key_hint) + next_entry = pool.mark_exhausted_and_rotate( + status_code=rotate_status, + error_context=error_context, + api_key_hint=_api_key_hint, + ) if next_entry is not None: _ra().logger.info( "Credential %s (rate limit) — rotated to pool entry %s", @@ -1119,7 +1141,11 @@ def recover_with_credential_pool( # Refresh failed — rotate to next credential instead of giving up. # The failed entry is already marked exhausted by try_refresh_current(). rotate_status = status_code if status_code is not None else 401 - next_entry = pool.mark_exhausted_and_rotate(status_code=rotate_status, error_context=error_context, api_key_hint=_api_key_hint) + next_entry = pool.mark_exhausted_and_rotate( + status_code=rotate_status, + error_context=error_context, + api_key_hint=_api_key_hint, + ) if next_entry is not None: _ra().logger.info( "Credential %s (auth refresh failed) — rotated to pool entry %s", diff --git a/tests/agent/test_credential_pool_routing.py b/tests/agent/test_credential_pool_routing.py index 810c18eea280..54b63f90c008 100644 --- a/tests/agent/test_credential_pool_routing.py +++ b/tests/agent/test_credential_pool_routing.py @@ -6,8 +6,12 @@ Covers: 3. Eager fallback deferred when credential pool has credentials 4. Eager fallback fires when no credential pool exists 5. Full 429 rotation cycle: retry-same → rotate → exhaust → fallback +6. Failure attribution: the entry matching the failing API key is marked + exhausted, not whatever pool.current() happens to point at """ +import json +import time from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -361,3 +365,127 @@ class TestApiKeyHintRealPool: statuses = {e.id: e.last_status for e in pool._entries} assert statuses["cred-healthy"] == "exhausted" assert statuses["cred-failed"] in (None, "ok") + + +# --------------------------------------------------------------------------- +# 7. Failure attribution — mark the key that failed, not pool.current() +# --------------------------------------------------------------------------- + +class TestFailureAttribution: + """Regression: recover_with_credential_pool must mark the entry whose API + key actually produced the failure. + + pool.current() is shared mutable state: round-robin select() advances it, + concurrent turns move it, and a freshly loaded pool (second process) has + current() == None — in which case the old code fell through to + _select_unlocked() and exhausted the NEXT (healthy) entry, copying the + failing key's error/reset time onto it until the whole pool went offline. + """ + + def _make_pool(self, tmp_path, monkeypatch, entries): + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) + hermes_home = tmp_path / "hermes" + hermes_home.mkdir(parents=True, exist_ok=True) + (hermes_home / "auth.json").write_text( + json.dumps({"version": 1, "credential_pool": {"anthropic": entries}}) + ) + from agent.credential_pool import load_pool + + return load_pool("anthropic") + + def _entry(self, idx, key, **overrides): + entry = { + "id": f"cred-{idx}", + "label": f"key-{idx}", + "auth_type": "api_key", + "priority": idx, + "source": "manual", + "access_token": key, + } + entry.update(overrides) + return entry + + def _agent(self, pool, failing_key): + return SimpleNamespace( + provider="anthropic", + api_key=failing_key, + _credential_pool=pool, + _swap_credential=MagicMock(), + ) + + def _statuses(self, pool): + return {e.id: e.last_status for e in pool.entries()} + + def test_billing_marks_failing_key_not_pointer(self, tmp_path, monkeypatch): + """Freshly loaded pool (current() is None): a 402 on key B must mark + entry B exhausted, not entry A (which _select_unlocked would return).""" + pool = self._make_pool( + tmp_path, monkeypatch, + [self._entry(0, "key-a"), self._entry(1, "key-b")], + ) + assert pool.current() is None + agent = self._agent(pool, failing_key="key-b") + + from agent.agent_runtime_helpers import recover_with_credential_pool + + recovered, _ = recover_with_credential_pool( + agent, status_code=402, has_retried_429=False + ) + + assert recovered is True + statuses = self._statuses(pool) + assert statuses["cred-1"] == "exhausted" + assert statuses["cred-0"] != "exhausted" + swapped = agent._swap_credential.call_args[0][0] + assert swapped.id == "cred-0" + + def test_rate_limit_marks_failing_key_not_pointer(self, tmp_path, monkeypatch): + """Same attribution for the 429 rotation path (second consecutive 429).""" + pool = self._make_pool( + tmp_path, monkeypatch, + [self._entry(0, "key-a"), self._entry(1, "key-b")], + ) + agent = self._agent(pool, failing_key="key-b") + + from agent.agent_runtime_helpers import recover_with_credential_pool + + recovered, has_retried = recover_with_credential_pool( + agent, status_code=429, has_retried_429=True + ) + + assert recovered is True + assert has_retried is False + statuses = self._statuses(pool) + assert statuses["cred-1"] == "exhausted" + assert statuses["cred-0"] != "exhausted" + + def test_pre_exhausted_check_uses_failing_key(self, tmp_path, monkeypatch): + """The 'already exhausted → rotate immediately' check must inspect the + failing entry, not pool.current(): first 429 on an already-exhausted + key rotates without burning a retry.""" + pool = self._make_pool( + tmp_path, monkeypatch, + [ + self._entry(0, "key-a"), + self._entry( + 1, "key-b", + last_status="exhausted", + last_status_at=time.time(), + last_error_code=429, + ), + ], + ) + agent = self._agent(pool, failing_key="key-b") + + from agent.agent_runtime_helpers import recover_with_credential_pool + + recovered, has_retried = recover_with_credential_pool( + agent, status_code=429, has_retried_429=False + ) + + assert recovered is True + assert has_retried is False + statuses = self._statuses(pool) + assert statuses["cred-0"] != "exhausted" + swapped = agent._swap_credential.call_args[0][0] + assert swapped.id == "cred-0" diff --git a/tests/run_agent/test_run_agent.py b/tests/run_agent/test_run_agent.py index 3e719882de04..f4edc65896bb 100644 --- a/tests/run_agent/test_run_agent.py +++ b/tests/run_agent/test_run_agent.py @@ -6625,6 +6625,9 @@ class TestCredentialPoolRecovery: def current(self): return SimpleNamespace(label="primary") + def entries(self): + return [] + def mark_exhausted_and_rotate( self, *, status_code, error_context=None, api_key_hint=None ): @@ -6839,6 +6842,9 @@ class TestCredentialPoolRecovery: def current(self): return SimpleNamespace(label="primary") + def entries(self): + return [] + def mark_exhausted_and_rotate( self, *, status_code, error_context=None, api_key_hint=None ):