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:
parent
d1eb08fcf3
commit
9cd0338688
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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-key–level 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:
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
Loading…
Reference in New Issue