diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index 1bdce7a4989d4..8ddf1b95fd80a 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -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 diff --git a/agent/credential_pool.py b/agent/credential_pool.py index 3466464424394..85a6e79a75f7a 100644 --- a/agent/credential_pool.py +++ b/agent/credential_pool.py @@ -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: diff --git a/tests/agent/test_credential_pool_routing.py b/tests/agent/test_credential_pool_routing.py index be05424059d9c..d7e68e840cae6 100644 --- a/tests/agent/test_credential_pool_routing.py +++ b/tests/agent/test_credential_pool_routing.py @@ -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" + diff --git a/tests/agent/test_credential_pool_sole_cooldown.py b/tests/agent/test_credential_pool_sole_cooldown.py index ea933ba93a5cd..26a2ebe4e7806 100644 --- a/tests/agent/test_credential_pool_sole_cooldown.py +++ b/tests/agent/test_credential_pool_sole_cooldown.py @@ -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."""