fix: arm the pool revert only when the benched credential outranks the one rotated to

A session already on the fallback (preferred entry benched by another session) that
rotates UP to the preferred entry once its window reopened must not be pulled back
DOWN when the fallback's own cooldown lifts. Compare priorities in _rotate_and_swap
before arming _credential_pool_revert_id; cover the two-session interleaving as the
control case inside the existing positive test.
This commit is contained in:
teknium1
2026-09-18 03:57:53 -07:00
committed by Teknium
parent 92bb5b92b8
commit 4c5a70065c
2 changed files with 34 additions and 6 deletions

View File

@@ -854,16 +854,21 @@ def recover_with_credential_pool(
rotate_status, label, getattr(next_entry, "id", "?"),
)
swapped = agent._swap_credential(next_entry) is not False
benched = next((e for e in pool.entries() if e.id == credential_id), None) if credential_id else None
if (
swapped
and credential_id
and benched is not None
and benched.priority < getattr(next_entry, "priority", benched.priority)
and not getattr(agent, "_credential_pool_revert_id", None)
and effective_reason in (FailoverReason.rate_limit, FailoverReason.billing)
):
# A quota bench (429/402) lifts when the window reopens, and a fresh session's
# select() would go straight back to this entry; arm the per-turn hook so the live
# session does too (#114501). Keep the FIRST benched entry across chained rotations
# — it is the preferred one. Auth benches are not windows; they stay as they are.
# session does too (#114501). Only when the benched entry OUTRANKS the one we rotated
# to: a session that was already on the fallback (preferred benched elsewhere) and
# rotates UP once the preferred window reopened must not be pulled back down when
# the fallback's cooldown lifts. Keep the FIRST benched entry across chained
# rotations — it is the preferred one. Auth benches are not windows; they stay.
agent._credential_pool_revert_id = credential_id
return swapped
if effective_reason == FailoverReason.upstream_rate_limit:

View File

@@ -52,12 +52,13 @@ class _LiveAgent:
return False
def _expire_cooldowns(monkeypatch):
real = time.time
monkeypatch.setattr(cp.time, "time", lambda: real() + EXHAUSTED_TTL_429_SECONDS + 120)
def _expire_cooldowns(monkeypatch, real=None, windows=1):
real = real or time.time
monkeypatch.setattr(cp.time, "time", lambda: real() + windows * (EXHAUSTED_TTL_429_SECONDS + 120))
def test_live_session_reverts_to_quota_benched_credential_once_cooldown_lifts(monkeypatch):
real_time = time.time
pool = CredentialPool(provider="anthropic", entries=[
_entry("pref0000", "subscription-oauth", priority=0, auth_type="oauth", token="sk-ant-oat01-PREF"),
_entry("fall0000", "paid-api-key", priority=1, auth_type="api_key", token="sk-ant-api03-FALL"),
@@ -83,6 +84,28 @@ def test_live_session_reverts_to_quota_benched_credential_once_cooldown_lifts(mo
# The pool agrees the preferred entry is healthy again (cooldown cleared, not merely elapsed).
assert next(e for e in pool.entries() if e.id == "pref0000").last_status != cp.STATUS_EXHAUSTED
# Control (two-session interleaving): a session that started on the FALLBACK because another
# session benched the preferred entry, then rotates UP to the preferred entry once its window
# reopened, must not be pulled back DOWN when the fallback's own cooldown lifts.
pool2 = CredentialPool(provider="anthropic", entries=[
_entry("pref0000", "subscription-oauth", priority=0, auth_type="oauth", token="sk-ant-oat01-PREF"),
_entry("fall0000", "paid-api-key", priority=1, auth_type="api_key", token="sk-ant-api03-FALL"),
])
monkeypatch.setattr(cp.time, "time", real_time)
pool2.mark_exhausted_and_rotate( # the OTHER session benches the preferred entry
status_code=429, credential_id="pref0000", failure_reason="rate_limit", error_context={"message": "Error"},
)
late = _LiveAgent(pool2)
assert late.api_key == "sk-ant-api03-FALL"
_expire_cooldowns(monkeypatch, real_time, 1) # pref's window reopened; fall benched from now
recover_with_credential_pool(late, status_code=429, has_retried_429=False, error_context={"message": "Error"})
recovered, _ = recover_with_credential_pool(late, status_code=429, has_retried_429=True, error_context={"message": "Error"})
assert recovered and late.api_key == "sk-ant-oat01-PREF"
assert getattr(late, "_credential_pool_revert_id", None) is None # rotated UP: nothing to revert to
_expire_cooldowns(monkeypatch, real_time, 2) # fall's cooldown lifts too
assert restore_primary_runtime(late) is False
assert late.api_key == "sk-ant-oat01-PREF" and pool2.select().id == "pref0000"
def test_auth_bench_does_not_arm_a_revert(monkeypatch):
"""A 401 bench is not a quota window: the session keeps the credential it rotated to."""