fix: break unbounded 401 retry loop in credential pool OAuth path

When api_key_hint from a 401 response doesn't match any pool entry
(common with OAuth tokens where runtime_api_key rotates), the pool
rotated without marking anything exhausted and handed back a fresh
selection. Because nothing was ever marked, the pool could never reach
the "no available entries" state — the caller retried the same dead
token forever (~6 attempts/sec), starving the event loop so /stop was
never processed; only killing the gateway ended it.

Rebased onto the identity-tracking rework that landed on main
(73c4b5a045): the single-entry escape from that commit already stops
the most common OAuth case, so this fix bounds the REMAINING gap —
multi-entry pools ping-ponging A->B->A with an unmatched hint. Cap
consecutive no-mark rotations at one full lap of the available
entries, then return None so the error surfaces / fallback activates.

Deliberately does NOT mark innocent entries exhausted (the original
PR's approach): that would quarantine a healthy key for the full
cooldown TTL on a hint that provably matches nothing. No cooldown is
written by the escape, so healthy keys stay available next turn --
bounded without hammering.

The streak resets when a rotation identifies a real entry and on any
successful normal select(), so only genuinely consecutive unmatched
rotations trip the bound.

Fixes #70401
This commit is contained in:
JonthanaHanh 2026-07-24 08:08:10 +07:00 committed by Teknium
parent 8b45a8b0d4
commit 431d2a628c
2 changed files with 234 additions and 1 deletions

View file

@ -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:

View file

@ -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"