diff --git a/agent/credential_pool.py b/agent/credential_pool.py index 516e7f61db61..08b0c0ea6b97 100644 --- a/agent/credential_pool.py +++ b/agent/credential_pool.py @@ -594,6 +594,14 @@ class CredentialPool: # Re-armed to None on every successful selection so a recover→re-exhaust # transition logs promptly instead of being swallowed by a stale window. self._last_no_entries_log_at: Optional[float] = None + # #70401: consecutive mark_exhausted_and_rotate() calls whose supplied + # credential identity matched no pool entry (OAuth wrappers whose + # runtime key rotates, entries pruned by another process, ...). These + # rotations mark nothing exhausted, so without a cap the pool can + # never converge to "no available entries" and the caller's 401 retry + # loop runs unbounded and non-interruptible. Reset whenever a real + # entry is identified or an escape path returns None. + self._unmatched_rotation_streak: int = 0 def has_credentials(self) -> bool: with self._lock: @@ -1584,7 +1592,13 @@ class CredentialPool: def select(self) -> Optional[PooledCredential]: with self._lock: - return self._select_unlocked() + entry = self._select_unlocked() + if entry is not None: + # A normal (non-recovery) selection starts a fresh episode — + # don't let a leftover unmatched-rotation streak from an old + # failure trip the #70401 bound early next time. + self._unmatched_rotation_streak = 0 + return entry def _available_entries(self, *, clear_expired: bool = False, refresh: bool = False) -> List[PooledCredential]: """Return entries not currently in exhaustion cooldown. @@ -1809,6 +1823,35 @@ class CredentialPool: # (rotated away, or a wrapper whose runtime key differs). # Falling through to current()/_select_unlocked() would mark an # innocent healthy key exhausted for the full cooldown TTL. + # + # #70401: this branch must still be BOUNDED. With OAuth-token + # auth the upstream 401's key hint never matches any entry's + # ``runtime_api_key``, so every retry lands here, nothing is + # ever marked exhausted, and the pool can never reach the + # "no available entries" state — the caller retries the same + # dead token forever (~6/sec, starving the event loop so chat + # interrupts are never processed). The single-entry case + # below already escapes; multi-entry pools could still + # ping-pong A→B→A indefinitely without marking anything. + # Cap consecutive no-mark rotations at one full lap of the + # available entries: past that, every candidate has been + # handed back at least once without recovery, so stop + # guessing and surface the error (no cooldown is written for + # anybody — healthy keys stay available for the next turn). + self._unmatched_rotation_streak += 1 + available_count = len(self._available_entries()) + if self._unmatched_rotation_streak > max(available_count, 1): + logger.warning( + "credential pool: failed credential identity matched no " + "%s entry for %d consecutive rotations (pool size %d) — " + "surfacing the error instead of rotating again", + self.provider, + self._unmatched_rotation_streak, + available_count, + ) + self._unmatched_rotation_streak = 0 + self._current_id = None + return None logger.info( "credential pool: failed credential identity matched no %s " "entry; rotating without marking any credential exhausted", @@ -1821,9 +1864,13 @@ class CredentialPool: # entry reports a successful recovery without changing # the credential, so the caller retries the same 401 # indefinitely. Let fallback/error propagation proceed. + self._unmatched_rotation_streak = 0 self._current_id = None return None return next_entry + # A real entry was identified — any prior unmatched-rotation + # streak is stale (this mark WILL advance pool state). + self._unmatched_rotation_streak = 0 if entry is None: entry = self._current_unlocked() or self._select_unlocked() if entry is None: diff --git a/tests/agent/test_credential_pool_unmatched_rotation_bound.py b/tests/agent/test_credential_pool_unmatched_rotation_bound.py new file mode 100644 index 000000000000..ded5bbd0302f --- /dev/null +++ b/tests/agent/test_credential_pool_unmatched_rotation_bound.py @@ -0,0 +1,186 @@ +"""#70401: the unmatched-identity rotation branch in +``mark_exhausted_and_rotate()`` must be bounded and must not write cooldowns +onto innocent healthy keys. + +With OAuth-token auth (provider ``nous``), the upstream 401's ``api_key_hint`` +never matches any pool entry's ``runtime_api_key`` — the wrapper's runtime key +rotates. The no-match branch deliberately marks nothing exhausted (marking +would quarantine an innocent healthy key for the full cooldown TTL) and hands +back a fresh selection. But because nothing is ever marked, the pool can never +converge to the "no available entries" state: with the old code the caller +retried the same dead token forever (~6/sec), starving the event loop so chat +``/stop`` interrupts were never processed; only killing the gateway ended it. + +The fix keeps the don't-mark-innocent-keys semantics (see the breaker/cooldown +design notes in ``mark_exhausted_and_rotate`` — the pool only trips on +confirmed-empty state, and no cooldown is invented here) but BOUNDS the +branch: after one full lap of the available entries with no recovery, the +rotation returns None so the caller surfaces the error / activates fallback. +Healthy keys carry no cooldown and are immediately available next turn — this +does not reintroduce hammering, it stops it. +""" +import json + +import pytest + + +def _seed_pool(tmp_path, monkeypatch, entries, provider="openrouter"): + 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": {provider: entries}}) + ) + monkeypatch.setenv("HERMES_HOME", str(hermes_home)) + from agent.credential_pool import load_pool + + return load_pool(provider) + + +def _entry(idx, key): + return { + "id": f"cred-{idx}", + "label": f"key-{idx}", + "auth_type": "api_key", + "priority": idx, + "source": "manual", + "access_token": key, + } + + +class TestUnmatchedHintRotationIsBounded: + def test_single_entry_pool_returns_none_immediately( + self, tmp_path, monkeypatch + ): + """Single-entry OAuth pool + unmatched hint → None right away (a + no-op rotation must not be reported as recovery).""" + pool = _seed_pool(tmp_path, monkeypatch, [_entry(0, "pool-key")]) + assert pool.select() is not None + + result = pool.mark_exhausted_and_rotate( + status_code=401, + error_context={"reason": "unauthorized"}, + api_key_hint="oauth-runtime-token-that-matches-nothing", + ) + + assert result is None + # No innocent key was quarantined. + statuses = {e.id: e.last_status for e in pool._entries} + assert statuses["cred-0"] != "exhausted" + + def test_multi_entry_pool_unmatched_hint_loop_terminates( + self, tmp_path, monkeypatch + ): + """Multi-entry pool: consecutive unmatched-hint rotations must reach + None within one lap of the pool instead of ping-ponging forever.""" + pool = _seed_pool( + tmp_path, monkeypatch, + [_entry(0, "key-a"), _entry(1, "key-b"), _entry(2, "key-c")], + ) + assert pool.select() is not None + + results = [] + for _ in range(10): # caller's retry loop + nxt = pool.mark_exhausted_and_rotate( + status_code=401, + error_context={"reason": "unauthorized"}, + api_key_hint="oauth-runtime-token-that-matches-nothing", + ) + results.append(nxt) + if nxt is None: + break + else: + pytest.fail( + "unbounded 401 retry loop: 10 unmatched-hint rotations never " + "returned None (#70401)" + ) + + # Bounded within one lap (3 available entries → at most 3 rotations + # before the streak trips). + assert len(results) <= 4 + assert results[-1] is None + # The escape must NOT have invented cooldowns for healthy keys. + statuses = {e.id: e.last_status for e in pool._entries} + assert all(status != "exhausted" for status in statuses.values()), ( + f"innocent keys were quarantined: {statuses}" + ) + + def test_streak_resets_when_identity_matches(self, tmp_path, monkeypatch): + """A rotation that identifies a real entry resets the streak — the + bound only fires on CONSECUTIVE unmatched rotations.""" + pool = _seed_pool( + tmp_path, monkeypatch, + [_entry(0, "key-a"), _entry(1, "key-b"), _entry(2, "key-c")], + ) + assert pool.select() is not None + + # Two unmatched rotations (streak = 2, below the 3-entry bound). + for _ in range(2): + assert pool.mark_exhausted_and_rotate( + status_code=401, + error_context={"reason": "unauthorized"}, + api_key_hint="no-match", + ) is not None + + # A matched rotation (key-a) marks a real entry → streak resets. + assert pool.mark_exhausted_and_rotate( + status_code=401, + error_context={"reason": "unauthorized"}, + api_key_hint="key-a", + ) is not None + + # A fresh unmatched episode still gets its full lap (2 available + # entries remain, so two rotations before the bound trips). + assert pool.mark_exhausted_and_rotate( + status_code=401, + error_context={"reason": "unauthorized"}, + api_key_hint="no-match", + ) is not None + + def test_streak_resets_on_normal_selection(self, tmp_path, monkeypatch): + """select() (a fresh episode) clears any leftover streak so the next + failure gets its full rotation budget.""" + pool = _seed_pool( + tmp_path, monkeypatch, + [_entry(0, "key-a"), _entry(1, "key-b")], + ) + assert pool.select() is not None + + # Burn the streak up to (but not past) the bound. + for _ in range(2): + pool.mark_exhausted_and_rotate( + status_code=401, + error_context={"reason": "unauthorized"}, + api_key_hint="no-match", + ) + + # New turn: a successful normal selection resets the streak. + assert pool.select() is not None + assert pool._unmatched_rotation_streak == 0 + + # The next unmatched rotation is attempt 1 of a new episode. + assert pool.mark_exhausted_and_rotate( + status_code=401, + error_context={"reason": "unauthorized"}, + api_key_hint="no-match", + ) is not None + + def test_matched_hint_path_unaffected(self, tmp_path, monkeypatch): + """Regression guard: the normal matched-hint path still marks the + failing entry and rotates to the healthy one.""" + pool = _seed_pool( + tmp_path, monkeypatch, + [_entry(0, "key-healthy"), _entry(1, "key-failed")], + ) + assert pool.select().access_token == "key-healthy" + + nxt = pool.mark_exhausted_and_rotate( + status_code=401, + error_context={"reason": "unauthorized"}, + api_key_hint="key-failed", + ) + + statuses = {e.id: e.last_status for e in pool._entries} + assert statuses["cred-1"] == "exhausted" + assert statuses["cred-0"] != "exhausted" + assert nxt is not None + assert nxt.access_token == "key-healthy"