fix(credential_pool): acquire the pool lock in has_available/peek/current/entries

`has_available()`, `peek()`, `current()` and `entries()` read (and, via
`_available_entries()`, mutate and persist) `self._entries` without holding
`self._lock`, while every other entry point — `select()`,
`mark_exhausted_and_rotate()`, `acquire_lease()`, `try_refresh_current()` —
guards the exact same access with the lock.

`_available_entries()` is not read-only: it prunes aged-out DEAD manual
entries (rebinding `self._entries` at the prune step) and calls `_persist()`
(writes auth.json). The gateway runs platform adapters in threads and cron
runs jobs in a ThreadPoolExecutor, so a status probe via `has_available()`
or `peek()` can race a concurrent `select()`/rotation: torn iteration of
`self._entries`, interleaved auth.json writes, or a lost token rotation.

Fix: take `self._lock` in all four query methods. Because the lock is
non-reentrant and `peek()` composes `current()` + `_available_entries()`,
add a lock-free `_current_unlocked()` helper and route the already-locked
internal callers (`_select_unlocked`, `mark_exhausted_and_rotate`,
`_try_refresh_current_unlocked`) through it to avoid self-deadlock.

Added regression tests: a no-deadlock check (peek re-entrancy) and a
lock-held-blocks-the-call check for each of the four methods.
This commit is contained in:
solyanviktor-star 2026-07-11 14:42:28 +03:00 committed by Teknium
parent 65d42e35d4
commit 5b794c984e
2 changed files with 105 additions and 11 deletions

View file

@ -600,16 +600,27 @@ class CredentialPool:
def has_available(self) -> bool:
"""True if at least one entry is not currently in exhaustion cooldown."""
return bool(self._available_entries())
# ``_available_entries`` is not read-only: it prunes aged-out DEAD
# manual entries (rebinding ``self._entries``) and persists. It must
# run under ``self._lock`` like every other caller (``select`` etc.),
# otherwise a status probe here can race a concurrent ``select`` /
# rotation and tear ``self._entries`` or double-write auth.json.
with self._lock:
return bool(self._available_entries())
def entries(self) -> List[PooledCredential]:
return list(self._entries)
with self._lock:
return list(self._entries)
def current(self) -> Optional[PooledCredential]:
def _current_unlocked(self) -> Optional[PooledCredential]:
if not self._current_id:
return None
return next((entry for entry in self._entries if entry.id == self._current_id), None)
def current(self) -> Optional[PooledCredential]:
with self._lock:
return self._current_unlocked()
def _replace_entry(self, old: PooledCredential, new: PooledCredential) -> None:
"""Swap an entry in-place by id, preserving sort order."""
for idx, entry in enumerate(self._entries):
@ -1729,18 +1740,21 @@ class CredentialPool:
self._entries = [replace(candidate, priority=idx) for idx, candidate in enumerate(rotated)]
self._persist()
self._current_id = entry.id
return self.current() or entry
return self._current_unlocked() or entry
entry = available[0]
self._current_id = entry.id
return entry
def peek(self) -> Optional[PooledCredential]:
current = self.current()
if current is not None:
return current
available = self._available_entries()
return available[0] if available else None
# Single lock acquisition for the whole read; call the unlocked
# helpers so we don't re-enter the non-reentrant ``self._lock``.
with self._lock:
current = self._current_unlocked()
if current is not None:
return current
available = self._available_entries()
return available[0] if available else None
def mark_exhausted_and_rotate(
self,
@ -1775,7 +1789,7 @@ class CredentialPool:
self._current_id = None
return self._select_unlocked()
if entry is None:
entry = self.current() or self._select_unlocked()
entry = self._current_unlocked() or self._select_unlocked()
if entry is None:
return None
_label = entry.label or entry.id[:8]
@ -1899,7 +1913,7 @@ class CredentialPool:
return self._try_refresh_current_unlocked()
def _try_refresh_current_unlocked(self) -> Optional[PooledCredential]:
entry = self.current()
entry = self._current_unlocked()
if entry is None:
return None
refreshed = self._refresh_entry(entry, force=True)

View file

@ -3555,3 +3555,83 @@ def test_sync_anthropic_entry_clears_all_error_fields(tmp_path, monkeypatch):
assert synced.last_error_reason is None
assert synced.last_error_message is None
assert synced.last_error_reset_at is None
def _load_two_ok_pool(tmp_path, monkeypatch):
"""A pool with two OK anthropic entries, current = cred-1."""
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": "***",
"last_status": "ok", "last_status_at": None, "last_error_code": None,
},
{
"id": "cred-2", "label": "secondary", "auth_type": "api_key",
"priority": 1, "source": "manual", "access_token": "***",
"last_status": "ok", "last_status_at": None, "last_error_code": None,
},
]
},
},
)
from agent.credential_pool import load_pool
return load_pool("anthropic")
class TestCredentialPoolQueryLocking:
"""Read/status methods must run under ``self._lock``.
``has_available``/``peek``/``current``/``entries`` all touch
``self._entries`` (and ``_available_entries`` even prunes + persists),
so they must hold the same lock every mutating entry point uses. A
naive fix would deadlock because the lock is non-reentrant and ``peek``
calls ``current`` + ``_available_entries``; these tests guard both the
no-deadlock and the actually-locked properties.
"""
def test_query_methods_do_not_deadlock(self, tmp_path, monkeypatch):
pool = _load_two_ok_pool(tmp_path, monkeypatch)
pool.select() # set a current entry
# peek() internally calls current() + _available_entries(); if any of
# these re-acquired the non-reentrant lock we'd hang here forever.
assert pool.current() is not None
assert pool.peek() is not None
assert pool.has_available() is True
# (env may seed extra singleton entries; just assert ours are present)
assert {"cred-1", "cred-2"} <= {e.id for e in pool.entries()}
@pytest.mark.parametrize("method", ["has_available", "peek", "current", "entries"])
def test_query_method_acquires_lock(self, tmp_path, monkeypatch, method):
import threading
pool = _load_two_ok_pool(tmp_path, monkeypatch)
pool.select()
done = threading.Event()
def _call():
getattr(pool, method)()
done.set()
# Hold the pool lock, then fire the query on another thread. If the
# method acquires self._lock (as it must), it blocks until we release.
pool._lock.acquire()
try:
worker = threading.Thread(target=_call, daemon=True)
worker.start()
assert not done.wait(timeout=0.5), (
f"{method}() returned while the pool lock was held — it is not "
f"acquiring self._lock"
)
finally:
pool._lock.release()
assert done.wait(timeout=2.0), f"{method}() did not complete after lock release"