From 3d67f00fe1ef24fba928a14ef2796e17017ab4ce Mon Sep 17 00:00:00 2001 From: Blade Date: Thu, 16 Jul 2026 12:55:33 -0400 Subject: [PATCH] fix(credential-pool): stop lost-update cooldown erasure and wrong-key quarantine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two related races in credential-pool cooldown state: 1. Lost update across processes: write_credential_pool merged only entries missing from the caller's snapshot; for entries present on both sides the caller's in-memory copy won wholesale. A process holding a snapshot taken before another process marked a key exhausted would, on its next persist (e.g. a round-robin rotation), write the key back as healthy — erasing the cooldown so every process resumes hammering a rate-limited key. Merge status fields by last_status_at recency: adopt the on-disk status only when it is strictly newer AND still binding (DEAD, or EXHAUSTED with an unexpired cooldown), and never onto re-authed (token-changed) entries, so legitimate expiry-clears and fresh logins are preserved. 2. Wrong-key quarantine: when mark_exhausted_and_rotate received an api_key_hint that matched no entry, it fell through to current()/_select_unlocked() — on a freshly loaded pool that selects the NEXT healthy key and benches it for the full cooldown TTL, punishing an innocent credential. When a hint is provided but unmatched, rotate without marking anything instead of guessing. Includes regression tests. Co-Authored-By: Claude Fable 5 --- agent/credential_pool.py | 14 +++ hermes_cli/auth.py | 85 ++++++++++++- tests/agent/test_credential_pool.py | 61 +++++++++ .../hermes_cli/test_auth_profile_fallback.py | 117 ++++++++++++++++++ 4 files changed, 276 insertions(+), 1 deletion(-) diff --git a/agent/credential_pool.py b/agent/credential_pool.py index f0fac12b509c..2134a3bc41d2 100644 --- a/agent/credential_pool.py +++ b/agent/credential_pool.py @@ -1760,6 +1760,20 @@ class CredentialPool: (e for e in self._entries if e.runtime_api_key == api_key_hint), None, ) + if entry is None: + # The failed key is identifiable but matches no entry + # (rotated away, or a wrapper whose runtime key differs). + # Falling through to current()/_select_unlocked() would + # mark an INNOCENT healthy key exhausted for the full + # cooldown TTL. Don't guess — just hand back a fresh + # selection so the caller can retry. + logger.info( + "credential pool: failed key hint matched no %s entry; " + "rotating without marking any credential exhausted", + self.provider, + ) + self._current_id = None + return self._select_unlocked() if entry is None: entry = self.current() or self._select_unlocked() if entry is None: diff --git a/hermes_cli/auth.py b/hermes_cli/auth.py index 507dd08a0220..23d0e51ea3e0 100644 --- a/hermes_cli/auth.py +++ b/hermes_cli/auth.py @@ -1447,6 +1447,73 @@ def read_credential_pool(provider_id: Optional[str] = None) -> Dict[str, Any]: return list(global_entries) if isinstance(global_entries, list) else [] +_POOL_STATUS_FIELDS = ( + "last_status", + "last_status_at", + "last_error_code", + "last_error_reason", + "last_error_message", + "last_error_reset_at", +) + + +def _merge_disk_cooldown_state( + entry: Dict[str, Any], + disk_entry: Optional[Dict[str, Any]], + provider_id: str, +) -> Dict[str, Any]: + """Keep a newer on-disk cooldown/quarantine over a stale in-memory one. + + ``write_credential_pool`` callers persist an in-memory snapshot that may + predate another process marking the same credential exhausted or dead + (last-writer-wins lost update). Without this merge, process B's later + rewrite resurrects a rate-limited key as healthy and both processes + resume hammering it. Adopt the on-disk status fields only when they are + strictly more recent (by ``last_status_at``) AND still binding — a DEAD + marker, or an EXHAUSTED cooldown that has not yet expired. Expired + cooldowns are not resurrected, so the pool's own expiry-clear (which + resets ``last_status_at`` to None) is never overridden. + """ + if not isinstance(disk_entry, dict): + return entry + try: + from agent.credential_pool import ( + PooledCredential, + STATUS_DEAD, + STATUS_EXHAUSTED, + _exhausted_until, + _parse_absolute_timestamp, + ) + + disk_status = disk_entry.get("last_status") + if disk_status not in (STATUS_DEAD, STATUS_EXHAUSTED): + return entry + # A token change means the caller re-authed/refreshed this entry and + # intentionally cleared its status (e.g. _sync_codex_entry_from_ + # auth_store after a fresh device-code login) — never resurrect the + # old cooldown onto fresh credentials. + mem_access = entry.get("access_token") or "" + disk_access = disk_entry.get("access_token") or "" + if mem_access and disk_access and mem_access != disk_access: + return entry + disk_ts = _parse_absolute_timestamp(disk_entry.get("last_status_at")) or 0.0 + mem_ts = _parse_absolute_timestamp(entry.get("last_status_at")) or 0.0 + if disk_ts <= mem_ts: + return entry + if disk_status == STATUS_EXHAUSTED: + until = _exhausted_until( + PooledCredential.from_dict(provider_id, disk_entry) + ) + if until is None or until <= time.time(): + return entry + merged_entry = dict(entry) + for status_field in _POOL_STATUS_FIELDS: + merged_entry[status_field] = disk_entry.get(status_field) + return merged_entry + except Exception: # pragma: no cover - best-effort merge + return entry + + def write_credential_pool( provider_id: str, entries: List[Dict[str, Any]], @@ -1464,6 +1531,10 @@ def write_credential_pool( the caller loaded its in-memory snapshot; without this merge a later rotation/exhaustion rewrite drops the concurrent credential. + For entries present on BOTH sides, status fields are merged by + ``last_status_at`` recency via ``_merge_disk_cooldown_state`` so a stale + snapshot cannot erase a cooldown/quarantine another process just wrote. + Pass ``removed_ids`` for entries the caller intentionally removed, so the merge does not resurrect them from the on-disk copy. """ @@ -1481,12 +1552,24 @@ def write_credential_pool( ] existing = pool.get(provider_id) existing_list = existing if isinstance(existing, list) else [] + existing_by_id = { + entry.get("id"): entry + for entry in existing_list + if isinstance(entry, dict) and entry.get("id") + } new_ids = { entry.get("id") for entry in sanitized_entries if isinstance(entry, dict) and entry.get("id") } - merged: List[Dict[str, Any]] = list(sanitized_entries) + merged: List[Dict[str, Any]] = [ + _merge_disk_cooldown_state( + entry, existing_by_id.get(entry.get("id")), provider_id + ) + if isinstance(entry, dict) + else entry + for entry in sanitized_entries + ] for disk_entry in existing_list: if not isinstance(disk_entry, dict): continue diff --git a/tests/agent/test_credential_pool.py b/tests/agent/test_credential_pool.py index 2a21062d66a5..8e3884ca37a6 100644 --- a/tests/agent/test_credential_pool.py +++ b/tests/agent/test_credential_pool.py @@ -443,6 +443,67 @@ def test_billing_rotation_marks_all_entries_sharing_failed_key(tmp_path, monkeyp assert statuses["cred-model-config"] == STATUS_EXHAUSTED +def test_unmatched_api_key_hint_rotates_without_benching_innocent_key(tmp_path, monkeypatch): + """An api_key_hint matching no entry must not quarantine a healthy key. + + Regression: when the hint was unmatched (key rotated away, or a wrapper + whose runtime key differs), mark_exhausted_and_rotate fell through to + current()/_select_unlocked() — on a freshly loaded pool that selects the + NEXT healthy key and benched it for the full cooldown TTL, punishing an + innocent credential. Now it rotates without marking anything. + """ + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) + _write_auth_store( + tmp_path, + { + "version": 1, + "credential_pool": { + "anthropic": [ + { + "id": "cred-1", + "label": "primary", + "auth_type": "api_key", + "priority": 0, + "source": "manual", + "access_token": "sk-ant-api-primary", + }, + { + "id": "cred-2", + "label": "secondary", + "auth_type": "api_key", + "priority": 1, + "source": "manual", + "access_token": "sk-ant-api-secondary", + }, + ] + }, + }, + ) + + from agent.credential_pool import load_pool, STATUS_DEAD, STATUS_EXHAUSTED + + # Freshly loaded pool: current() is None, exactly the shape of the bug. + pool = load_pool("anthropic") + + next_entry = pool.mark_exhausted_and_rotate( + status_code=429, + api_key_hint="sk-ant-api-rotated-away", + ) + + # A fresh selection is still handed back so the caller can retry... + assert next_entry is not None + + # ...but no credential was benched, in memory or on disk. + assert all( + entry.last_status not in (STATUS_EXHAUSTED, STATUS_DEAD) + for entry in pool.entries() + ) + auth_payload = json.loads((tmp_path / "hermes" / "auth.json").read_text()) + for persisted in auth_payload["credential_pool"]["anthropic"]: + assert persisted.get("last_status") not in (STATUS_EXHAUSTED, STATUS_DEAD) + assert persisted.get("last_error_code") is None + + def test_token_invalidated_marks_credential_dead(tmp_path, monkeypatch): """OpenAI Codex token_invalidated must mark the credential DEAD, not exhausted. diff --git a/tests/hermes_cli/test_auth_profile_fallback.py b/tests/hermes_cli/test_auth_profile_fallback.py index dc6a31b1bd48..902c7f881ffa 100644 --- a/tests/hermes_cli/test_auth_profile_fallback.py +++ b/tests/hermes_cli/test_auth_profile_fallback.py @@ -12,6 +12,7 @@ authenticated only at the global root. from __future__ import annotations import json +import time from contextlib import contextmanager from pathlib import Path @@ -519,3 +520,119 @@ def test_auth_lock_reentrancy_is_scoped_after_profile_context_switch(profile_env reset_hermes_home_override(token) assert getattr(holder_a, "depth", 0) == 0 + + +# --------------------------------------------------------------------------- +# write_credential_pool — stale-snapshot cooldown merge +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def classic_env(tmp_path, monkeypatch): + """Classic single-root layout (HERMES_HOME != ~/.hermes, no profiles).""" + fake_home = tmp_path / "home" + fake_home.mkdir() + monkeypatch.setattr(Path, "home", lambda: fake_home) + hermes_home = tmp_path / "classic" + hermes_home.mkdir() + monkeypatch.setenv("HERMES_HOME", str(hermes_home)) + return hermes_home + + +def _pool_entry(**overrides) -> dict: + entry = { + "id": "cred-x", + "label": "key-x", + "auth_type": "api_key", + "priority": 0, + "source": "manual", + "access_token": "sk-x", + } + entry.update(overrides) + return entry + + +@pytest.mark.parametrize( + "disk_status,error_code", + [("exhausted", 429), ("dead", 401)], +) +def test_write_pool_stale_snapshot_keeps_newer_disk_cooldown( + classic_env, disk_status, error_code, +): + """A stale healthy snapshot must not erase a newer binding cooldown. + + Process A benches a key (EXHAUSTED with an unexpired cooldown, or DEAD); + process B persists a snapshot taken *before* that. The on-disk status is + strictly newer and still binding, so it must survive the rewrite instead + of the key being resurrected as healthy. + """ + from hermes_cli.auth import write_credential_pool + + benched_at = time.time() - 60 # newer than the snapshot, cooldown unexpired + _write(classic_env / "auth.json", _make_auth_store(pool={ + "openrouter": [_pool_entry( + last_status=disk_status, + last_status_at=benched_at, + last_error_code=error_code, + )], + })) + + # Stale in-memory snapshot: same entry, still healthy (no status fields). + write_credential_pool("openrouter", [_pool_entry()]) + + data = json.loads((classic_env / "auth.json").read_text()) + persisted = data["credential_pool"]["openrouter"][0] + assert persisted["last_status"] == disk_status + assert persisted["last_status_at"] == benched_at + assert persisted["last_error_code"] == error_code + + +def test_write_pool_expired_disk_cooldown_is_not_resurrected(classic_env): + """An expired on-disk cooldown is NOT re-adopted onto the snapshot. + + The pool's own expiry-clear (and any caller that legitimately observed + the cooldown lapse) must win: only still-binding cooldowns are merged. + """ + from hermes_cli.auth import write_credential_pool + + _write(classic_env / "auth.json", _make_auth_store(pool={ + "openrouter": [_pool_entry( + last_status="exhausted", + last_status_at=time.time() - 90_000, # far past the 1h 429 TTL + last_error_code=429, + )], + })) + + write_credential_pool("openrouter", [_pool_entry()]) + + data = json.loads((classic_env / "auth.json").read_text()) + persisted = data["credential_pool"]["openrouter"][0] + assert persisted.get("last_status") != "exhausted" + assert persisted.get("last_error_code") is None + + +def test_write_pool_never_merges_cooldown_onto_reauthed_entry(classic_env): + """A token change means re-auth: the old cooldown must never carry over. + + A fresh login intentionally clears the entry's status; resurrecting the + stale cooldown onto the new credentials would bench a just-authorized key. + """ + from hermes_cli.auth import write_credential_pool + + _write(classic_env / "auth.json", _make_auth_store(pool={ + "openrouter": [_pool_entry( + access_token="sk-old", + last_status="exhausted", + last_status_at=time.time() - 60, # newer AND unexpired + last_error_code=429, + )], + })) + + # Same entry id, freshly re-authed with a new token and cleared status. + write_credential_pool("openrouter", [_pool_entry(access_token="sk-new")]) + + data = json.loads((classic_env / "auth.json").read_text()) + persisted = data["credential_pool"]["openrouter"][0] + assert persisted["access_token"] == "sk-new" + assert persisted.get("last_status") != "exhausted" + assert persisted.get("last_error_code") is None