fix(credential-pool): bench a billing 403 fully, even as the sole key

The sole-credential cooldown sized the bench from the raw HTTP status, but
403 is overloaded: error_classifier maps OpenRouter's "key limit exceeded"
and xAI's spending-limit block to FailoverReason.billing, while an edge
throttle with the same status is transient. Only 402 was excluded from the
short cooldown, so a spent account on a single key retried every 60 seconds
and re-failed forever.

Thread the classified reason from recover_with_credential_pool through
mark_exhausted_and_rotate to _exhausted_ttl. Billing keeps the full bench
regardless of status; everything else transient still recovers in 60s. The
verdict is stored on the entry (_EXTRA_KEYS, so it persists to auth.json) —
without that a restart would re-read a bare 403 and downgrade the bench.

Tests: sole billing-403 stays benched, survives reload, unclassified 403
still recovers; call-site coverage that the reason actually reaches the pool.
Three existing kwargs assertions updated for the new argument.
This commit is contained in:
Brooklyn Nicholson 2026-08-04 11:43:55 -06:00 committed by kshitij
parent d1eb08fcf3
commit 9cd0338688
4 changed files with 166 additions and 14 deletions

View File

@ -1020,6 +1020,13 @@ def recover_with_credential_pool(
}
if _credential_id:
kwargs["credential_id"] = _credential_id
# Hand the pool the classified semantics, not just the status. A
# billing 403 (OpenRouter "key limit exceeded", xAI spending limit)
# and an edge-throttle 403 are the same number but need opposite
# cooldowns — the pool can only tell them apart if we say which.
# ``effective_reason`` is resolved below; this closure runs after.
if effective_reason is not None:
kwargs["failure_reason"] = effective_reason.value
return pool.mark_exhausted_and_rotate(**kwargs)
effective_reason = classified_reason

View File

@ -131,6 +131,11 @@ EXHAUSTED_TTL_DEFAULT_SECONDS = 60 * 60 # 1 hour
# the short 401 cooldown above. Provider-supplied reset_at still overrides.
EXHAUSTED_TTL_SOLE_CREDENTIAL_SECONDS = 60 # 1 minute
# ``FailoverReason.billing`` as a bare string. The pool stores classified
# failure semantics as plain text (it persists to JSON and must not import
# the classifier), so the value is duplicated here rather than referenced.
FAILURE_REASON_BILLING = "billing"
# Throttle window for the "no available entries" INFO line. Credential
# selection runs on a hot path (every model call, plus auxiliary tasks like
# compression/moa/titles), so when a pool is empty or fully exhausted the
@ -156,6 +161,13 @@ _EXTRA_KEYS = frozenset({
"token_type", "scope", "client_id", "portal_base_url", "obtained_at",
"expires_in", "agent_key_id", "agent_key_expires_in", "agent_key_reused",
"agent_key_obtained_at", "tls", "secret_source", "secret_fingerprint",
# Classified failure semantics for the last exhaustion, as decided by
# agent/error_classifier.py. The raw HTTP status is not enough to size a
# cooldown: providers return 403 for both an edge throttle (transient,
# seconds) and a spending/key limit (billing, needs a real fix). Persisted
# with the entry so a restart doesn't downgrade a billing bench back to a
# 60s transient cooldown.
"failure_reason",
})
@ -295,23 +307,37 @@ def _is_manual_source(source: str) -> bool:
return normalized == SOURCE_MANUAL or normalized.startswith(f"{SOURCE_MANUAL}:")
def _exhausted_ttl(error_code: Optional[int], *, sole_credential: bool = False) -> int:
def _exhausted_ttl(
error_code: Optional[int],
*,
sole_credential: bool = False,
failure_reason: Optional[str] = None,
) -> int:
"""Return cooldown seconds based on the HTTP status that caused exhaustion.
When *sole_credential* is True the pool has no other entry to rotate to, so
a long bench just blocks the only key. Transient throttles (429 and the
catch-all default, which covers 403/5xx/unknown) are capped to a brief
cooldown so the sole key can recover mirroring the short 401 path. 401
keeps its own (already short) TTL; 402 (billing/quota) keeps the full bench
since a quick retry can't help.
keeps its own (already short) TTL.
*failure_reason* is the classified semantics from
``agent/error_classifier.py``. The raw status alone can't size the
cooldown: an OpenRouter ``key limit exceeded`` and an xAI spending-limit
block both arrive as **403** but classify as ``billing``, and a 60s retry
on a spent account just re-fails every minute. Billing keeps the full
bench regardless of status; 402 does too, since it is billing by
definition even when nothing classified it.
"""
if error_code == 401:
return EXHAUSTED_TTL_401_SECONDS
base = EXHAUSTED_TTL_429_SECONDS if error_code == 429 else EXHAUSTED_TTL_DEFAULT_SECONDS
# Sole credential: shorten only TRANSIENT throttles (429 rate-limit, 403
# edge-throttle, 5xx server, or unknown). 402 (billing/quota) is a genuine
# exhaustion where a quick retry can't help, so it keeps the full bench.
if sole_credential and error_code != 402:
# edge-throttle, 5xx server, or unknown). Billing exhaustion — whether
# classified as such or self-evident from a 402 — is a genuine depletion
# where a quick retry can't help, so it keeps the full bench.
is_billing = error_code == 402 or failure_reason == FAILURE_REASON_BILLING
if sole_credential and not is_billing:
return min(base, EXHAUSTED_TTL_SOLE_CREDENTIAL_SECONDS)
return base
@ -402,7 +428,9 @@ def _exhausted_until(entry: PooledCredential, *, sole_credential: bool = False)
return reset_at
if entry.last_status_at:
return entry.last_status_at + _exhausted_ttl(
entry.last_error_code, sole_credential=sole_credential
entry.last_error_code,
sole_credential=sole_credential,
failure_reason=getattr(entry, "failure_reason", None),
)
return None
@ -769,6 +797,7 @@ class CredentialPool:
error_context: Optional[Dict[str, Any]] = None,
*,
persist: bool = True,
failure_reason: Optional[str] = None,
) -> PooledCredential:
normalized_error = _normalize_error_context(error_context)
# Permanent OAuth failures (token_invalidated, token_revoked, etc.)
@ -782,6 +811,15 @@ class CredentialPool:
terminal_status = STATUS_DEAD
else:
terminal_status = STATUS_EXHAUSTED
# Carry the classifier's verdict onto the entry so the cooldown can be
# sized by what actually failed, not just the HTTP status (a billing
# 403 must not get the sole-credential transient cooldown). Absent a
# classification, clear any stale verdict from a previous failure.
updated_extra = dict(entry.extra)
if failure_reason:
updated_extra["failure_reason"] = failure_reason
else:
updated_extra.pop("failure_reason", None)
updated = replace(
entry,
last_status=terminal_status,
@ -790,6 +828,7 @@ class CredentialPool:
last_error_reason=normalized_error.get("reason"),
last_error_message=normalized_error.get("message"),
last_error_reset_at=normalized_error.get("reset_at"),
extra=updated_extra,
)
self._replace_entry(entry, updated)
if persist:
@ -1996,6 +2035,7 @@ class CredentialPool:
error_context: Optional[Dict[str, Any]] = None,
api_key_hint: Optional[str] = None,
credential_id: Optional[str] = None,
failure_reason: Optional[str] = None,
) -> Optional[PooledCredential]:
with self._lock:
entry = None
@ -2074,7 +2114,9 @@ class CredentialPool:
if entry is None:
return None
_label = entry.label or entry.id[:8]
self._mark_exhausted(entry, status_code, error_context)
self._mark_exhausted(
entry, status_code, error_context, failure_reason=failure_reason
)
# A 402/429/401 is an API-keylevel failure: the account is out of
# balance, rate-limited, or its key is rejected. The same key can
# back more than one pool entry (e.g. an explicit pool entry plus a
@ -2094,7 +2136,11 @@ class CredentialPool:
continue
if sibling.runtime_api_key == failed_runtime_key:
self._mark_exhausted(
sibling, status_code, error_context, persist=False
sibling,
status_code,
error_context,
persist=False,
failure_reason=failure_reason,
)
siblings_marked = True
if siblings_marked:

View File

@ -210,7 +210,7 @@ class TestPoolRotationCycle:
)
assert recovered is True
assert has_retried is False # reset after rotation
pool.mark_exhausted_and_rotate.assert_called_once_with(status_code=429, error_context=None, api_key_hint="test-api-key")
pool.mark_exhausted_and_rotate.assert_called_once_with(status_code=429, error_context=None, api_key_hint="test-api-key", failure_reason="rate_limit")
agent._swap_credential.assert_called_once_with(entries[1])
def test_pool_exhaustion_returns_false(self):
@ -236,7 +236,7 @@ class TestPoolRotationCycle:
)
assert recovered is True
assert has_retried is False
pool.mark_exhausted_and_rotate.assert_called_once_with(status_code=402, error_context=None, api_key_hint="test-api-key")
pool.mark_exhausted_and_rotate.assert_called_once_with(status_code=402, error_context=None, api_key_hint="test-api-key", failure_reason="billing")
def test_api_key_hint_from_pool_current_when_agent_key_missing(self):
@ -273,7 +273,8 @@ class TestPoolRotationCycle:
)
assert recovered is True
pool.mark_exhausted_and_rotate.assert_called_once_with(
status_code=402, error_context=None, api_key_hint="pool-current-key"
status_code=402, error_context=None, api_key_hint="pool-current-key",
failure_reason="billing",
)
@ -492,3 +493,53 @@ class TestFailureAttribution:
assert self._statuses(pool)["cred-0"] != "exhausted"
agent._swap_credential.assert_not_called()
def test_classified_billing_403_recorded_on_entry(self, tmp_path, monkeypatch):
"""A billing-classified 403 must reach the pool as `billing`, not a bare 403.
`error_classifier` maps OpenRouter's `key limit exceeded` 403 (and xAI
spending-limit blocks) to FailoverReason.billing, but the pool only
ever saw the raw status so a sole-credential pool gave a spent
account the 60s transient cooldown and re-failed every minute. The
recovery path now forwards the classified reason so the pool can size
the bench correctly.
"""
from agent.error_classifier import FailoverReason
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")
agent._is_entitlement_failure = MagicMock(return_value=False)
from agent.agent_runtime_helpers import recover_with_credential_pool
recover_with_credential_pool(
agent,
status_code=403,
has_retried_429=False,
classified_reason=FailoverReason.billing,
)
failed = {e.id: e for e in pool.entries()}["cred-1"]
assert failed.last_status == "exhausted"
assert failed.failure_reason == "billing"
def test_unclassified_403_records_no_billing_reason(self, tmp_path, monkeypatch):
"""An unclassified 403 stays transient — no billing verdict is invented."""
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")
agent._is_entitlement_failure = MagicMock(return_value=False)
from agent.agent_runtime_helpers import recover_with_credential_pool
recover_with_credential_pool(
agent, status_code=403, has_retried_429=False
)
failed = {e.id: e for e in pool.entries()}["cred-1"]
assert failed.failure_reason != "billing"

View File

@ -19,8 +19,15 @@ def _write_auth_store(tmp_path, payload: dict) -> None:
(hermes_home / "auth.json").write_text(json.dumps(payload, indent=2))
def _entry(error_code: int, *, age_seconds: float, cred_id: str = "cred-1", priority: int = 0) -> dict:
return {
def _entry(
error_code: int,
*,
age_seconds: float,
cred_id: str = "cred-1",
priority: int = 0,
failure_reason: str | None = None,
) -> dict:
entry = {
"id": cred_id,
"label": cred_id,
"auth_type": "api_key",
@ -32,6 +39,9 @@ def _entry(error_code: int, *, age_seconds: float, cred_id: str = "cred-1", prio
"last_status_at": time.time() - age_seconds,
"last_error_code": error_code,
}
if failure_reason is not None:
entry["failure_reason"] = failure_reason
return entry
def _load(tmp_path, monkeypatch, entries: list[dict]):
@ -68,6 +78,44 @@ def test_sole_credential_403_recovers_after_short_cooldown(tmp_path, monkeypatch
assert entry.last_status == "ok"
def test_sole_credential_billing_403_keeps_full_bench(tmp_path, monkeypatch):
"""A 403 classified as BILLING must keep the full bench, not the 60s cooldown.
Providers overload 403: OpenRouter returns it for `key limit exceeded` and
xAI for a spending-limit block, both of which `error_classifier` maps to
FailoverReason.billing. Status alone can't tell those from an edge
throttle, so retrying a spent account every 60s just re-fails forever.
The classified reason rides along on the entry and wins over the status.
"""
pool = _load(
tmp_path,
monkeypatch,
[_entry(403, age_seconds=90, failure_reason="billing")],
)
assert pool.has_available() is False
assert pool.select() is None
def test_sole_credential_billing_403_survives_reload(tmp_path, monkeypatch):
"""The classified reason persists, so a restart can't downgrade the bench.
`failure_reason` is written to auth.json with the entry; without that, a
process restart would re-read a bare 403 and hand the spent key back after
60 seconds.
"""
from agent.credential_pool import _exhausted_ttl
pool = _load(
tmp_path,
monkeypatch,
[_entry(403, age_seconds=90, failure_reason="billing")],
)
entry = pool.entries()[0]
assert entry.failure_reason == "billing"
assert _exhausted_ttl(403, sole_credential=True, failure_reason="billing") == 60 * 60
assert _exhausted_ttl(403, sole_credential=True) == 60
def test_sole_credential_402_keeps_full_bench(tmp_path, monkeypatch):
"""402 (billing/quota) is genuine exhaustion — a quick retry can't help, so
the sole-credential short cooldown must NOT apply."""