mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-24 16:54:43 +00:00
fix(credential-pool): stop lost-update cooldown erasure and wrong-key quarantine
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 <noreply@anthropic.com>
This commit is contained in:
Blade
Teknium
parent
fdd3943cb1
commit
3d67f00fe1
4 changed files with 276 additions and 1 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue