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:
@@ -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:
|
||||
|
||||
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user