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"