From 769381fb3e01e1bfbc7ff980b5704a479281f66f Mon Sep 17 00:00:00 2001 From: solyanviktor-star <233359899+solyanviktor-star@users.noreply.github.com> Date: Sun, 12 Jul 2026 15:41:07 +0300 Subject: [PATCH] fix(credential_pool): complete the locking boundary across the public pool surface Follow-up to review feedback: - Acquire self._lock in the remaining public pool-state methods: has_credentials, reset_statuses, remove_index, resolve_target, and add_entry. All of them read or rebind self._entries (and the mutating ones persist auth.json), so they now hold the same lock as select() and the query methods. None are called from within the lock, so no unlocked helpers are needed. - Make the blocking test deterministic: an instrumented lock records the acquire attempt, and the test first waits for the worker to actually reach self._lock before asserting it blocks. Previously an unlocked method could pass if the worker thread was scheduled late. - Extend the lock test matrix to all nine public methods; the five newly locked ones fail the test without this fix. Co-Authored-By: Claude Fable 5 --- agent/credential_pool.py | 123 +++++++++++++++------------- tests/agent/test_credential_pool.py | 84 ++++++++++++++++--- 2 files changed, 135 insertions(+), 72 deletions(-) diff --git a/agent/credential_pool.py b/agent/credential_pool.py index d710ddde9fb8..0bc113746c79 100644 --- a/agent/credential_pool.py +++ b/agent/credential_pool.py @@ -596,7 +596,8 @@ class CredentialPool: self._last_no_entries_log_at: Optional[float] = None def has_credentials(self) -> bool: - return bool(self._entries) + with self._lock: + return bool(self._entries) def has_available(self) -> bool: """True if at least one entry is not currently in exhaustion cooldown.""" @@ -1922,76 +1923,80 @@ class CredentialPool: return refreshed def reset_statuses(self) -> int: - count = 0 - new_entries = [] - for entry in self._entries: - if entry.last_status or entry.last_status_at or entry.last_error_code: - new_entries.append( - replace( - entry, - last_status=None, - last_status_at=None, - last_error_code=None, - last_error_reason=None, - last_error_message=None, - last_error_reset_at=None, + with self._lock: + count = 0 + new_entries = [] + for entry in self._entries: + if entry.last_status or entry.last_status_at or entry.last_error_code: + new_entries.append( + replace( + entry, + last_status=None, + last_status_at=None, + last_error_code=None, + last_error_reason=None, + last_error_message=None, + last_error_reset_at=None, + ) ) - ) - count += 1 - else: - new_entries.append(entry) - if count: - self._entries = new_entries - self._persist() - return count + count += 1 + else: + new_entries.append(entry) + if count: + self._entries = new_entries + self._persist() + return count def remove_index(self, index: int) -> Optional[PooledCredential]: - if index < 1 or index > len(self._entries): - return None - removed = self._entries.pop(index - 1) - self._entries = [ - replace(entry, priority=new_priority) - for new_priority, entry in enumerate(self._entries) - ] - write_credential_pool( - self.provider, - [entry.to_dict() for entry in self._entries], - removed_ids=[removed.id], - ) - if self._current_id == removed.id: - self._current_id = None - return removed + with self._lock: + if index < 1 or index > len(self._entries): + return None + removed = self._entries.pop(index - 1) + self._entries = [ + replace(entry, priority=new_priority) + for new_priority, entry in enumerate(self._entries) + ] + write_credential_pool( + self.provider, + [entry.to_dict() for entry in self._entries], + removed_ids=[removed.id], + ) + if self._current_id == removed.id: + self._current_id = None + return removed def resolve_target(self, target: Any) -> Tuple[Optional[int], Optional[PooledCredential], Optional[str]]: raw = str(target or "").strip() if not raw: return None, None, "No credential target provided." - for idx, entry in enumerate(self._entries, start=1): - if entry.id == raw: - return idx, entry, None + with self._lock: + for idx, entry in enumerate(self._entries, start=1): + if entry.id == raw: + return idx, entry, None - label_matches = [ - (idx, entry) - for idx, entry in enumerate(self._entries, start=1) - if entry.label.strip().lower() == raw.lower() - ] - if len(label_matches) == 1: - return label_matches[0][0], label_matches[0][1], None - if len(label_matches) > 1: - return None, None, f'Ambiguous credential label "{raw}". Use the numeric index or entry id instead.' - if raw.isdigit(): - index = int(raw) - if 1 <= index <= len(self._entries): - return index, self._entries[index - 1], None - return None, None, f"No credential #{index}." - return None, None, f'No credential matching "{raw}".' + label_matches = [ + (idx, entry) + for idx, entry in enumerate(self._entries, start=1) + if entry.label.strip().lower() == raw.lower() + ] + if len(label_matches) == 1: + return label_matches[0][0], label_matches[0][1], None + if len(label_matches) > 1: + return None, None, f'Ambiguous credential label "{raw}". Use the numeric index or entry id instead.' + if raw.isdigit(): + index = int(raw) + if 1 <= index <= len(self._entries): + return index, self._entries[index - 1], None + return None, None, f"No credential #{index}." + return None, None, f'No credential matching "{raw}".' def add_entry(self, entry: PooledCredential) -> PooledCredential: - entry = replace(entry, priority=_next_priority(self._entries)) - self._entries.append(entry) - self._persist() - return entry + with self._lock: + entry = replace(entry, priority=_next_priority(self._entries)) + self._entries.append(entry) + self._persist() + return entry def _upsert_entry(entries: List[PooledCredential], provider: str, source: str, payload: Dict[str, Any]) -> bool: diff --git a/tests/agent/test_credential_pool.py b/tests/agent/test_credential_pool.py index 98792284018a..4fbe23efe5aa 100644 --- a/tests/agent/test_credential_pool.py +++ b/tests/agent/test_credential_pool.py @@ -3585,15 +3585,25 @@ def _load_two_ok_pool(tmp_path, monkeypatch): return load_pool("anthropic") +def _fresh_entry(pool): + """A copy of the pool's first entry under a new id, for add_entry().""" + from dataclasses import replace as dc_replace + + return dc_replace(pool.entries()[0], id="cred-new") + + class TestCredentialPoolQueryLocking: - """Read/status methods must run under ``self._lock``. + """Public pool-state 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. + and the management surface (``has_credentials``/``reset_statuses``/ + ``remove_index``/``resolve_target``/``add_entry``) reads or rebinds + ``self._entries`` and persists auth.json, so they must all 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): @@ -3605,33 +3615,81 @@ class TestCredentialPoolQueryLocking: assert pool.current() is not None assert pool.peek() is not None assert pool.has_available() is True + assert pool.has_credentials() is True + assert pool.resolve_target("cred-1")[1] is not None # (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): + @pytest.mark.parametrize( + "method,get_args", + [ + ("has_available", lambda pool: ()), + ("peek", lambda pool: ()), + ("current", lambda pool: ()), + ("entries", lambda pool: ()), + ("has_credentials", lambda pool: ()), + ("reset_statuses", lambda pool: ()), + ("resolve_target", lambda pool: ("cred-1",)), + ("remove_index", lambda pool: (1,)), + ("add_entry", lambda pool: (_fresh_entry(pool),)), + ], + ) + def test_query_method_acquires_lock(self, tmp_path, monkeypatch, method, get_args): import threading pool = _load_two_ok_pool(tmp_path, monkeypatch) pool.select() + args = get_args(pool) + + inner = pool._lock + + class _InstrumentedLock: + """Probe that records acquire attempts, so the test can prove the + worker actually reached ``self._lock`` before asserting that it + blocks (a plain timed wait passes spuriously if the worker is + simply never scheduled).""" + + def __init__(self): + self.attempted = threading.Event() + + def acquire(self, *args, **kwargs): + self.attempted.set() + return inner.acquire(*args, **kwargs) + + def release(self): + inner.release() + + def __enter__(self): + self.acquire() + return self + + def __exit__(self, *exc): + self.release() + + probe = _InstrumentedLock() + pool._lock = probe done = threading.Event() def _call(): - getattr(pool, method)() + getattr(pool, method)(*args) 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() + # Hold the real lock (without tripping the probe), then fire the query + # on another thread. If the method acquires self._lock (as it must), + # it blocks until we release. + inner.acquire() try: worker = threading.Thread(target=_call, daemon=True) worker.start() + assert probe.attempted.wait(timeout=2.0), ( + f"{method}() never attempted to acquire self._lock" + ) assert not done.wait(timeout=0.5), ( f"{method}() returned while the pool lock was held — it is not " - f"acquiring self._lock" + f"blocking on self._lock" ) finally: - pool._lock.release() + inner.release() assert done.wait(timeout=2.0), f"{method}() did not complete after lock release"