From dc116202c14ac474d2c32de92fe822d15af37f17 Mon Sep 17 00:00:00 2001 From: Ofer LaOr Date: Sat, 12 Sep 2026 14:03:48 +0000 Subject: [PATCH] fix(delegation): keep credential pools endpoint-coherent --- agent/credential_pool.py | 54 +++++++++++++++-- ..._credential_pool_lease_refresh_reselect.py | 18 +++++- tests/tools/test_delegate.py | 56 ++++++++++++++++++ tools/delegate_tool_child_run.py | 32 +++++++++- tools/delegate_tool_config.py | 59 ++++++++++++++++++- 5 files changed, 209 insertions(+), 10 deletions(-) diff --git a/agent/credential_pool.py b/agent/credential_pool.py index b76e53d866a77..08372dfdf37d8 100644 --- a/agent/credential_pool.py +++ b/agent/credential_pool.py @@ -1010,6 +1010,34 @@ def current(self) -> Optional[PooledCredential]: with self._lock: return self._current_unlocked() + def leased_entry( + self, + credential_id: str, + *, + entry_filter: Optional[Callable[[PooledCredential], bool]] = None, + ) -> Optional[PooledCredential]: + """Return the exact entry covered by an active lease. + + ``current()`` is a shared rotation cursor and may advance after a + child acquires its lease. Revalidate endpoint eligibility while holding + the pool lock so callers bind the leased credential rather than a + concurrent child's selection. + """ + with self._lock: + if self._active_leases.get(credential_id, 0) <= 0: + return None + entry = self._find(lambda candidate: candidate.id == credential_id) + if entry is None: + return None + if entry_filter is not None: + try: + if not entry_filter(entry): + return None + except Exception as exc: + logger.warning("credential pool: leased-entry filter failed: %s", exc) + return None + return entry + def entry_id_for_api_key(self, api_key_hint: Any = None) -> Optional[str]: """Stable id for the runtime credential in use. @@ -2088,7 +2116,12 @@ def mark_exhausted_and_rotate( # ---- leases ------------------------------------------------------------ - def acquire_lease(self, credential_id: Optional[str] = None) -> Optional[str]: + def acquire_lease( + self, + credential_id: Optional[str] = None, + *, + entry_filter: Optional[Callable[[PooledCredential], bool]] = None, + ) -> Optional[str]: """Acquire a soft lease on a credential. With *credential_id*, lease that entry directly. Otherwise prefer the @@ -2096,18 +2129,25 @@ def acquire_lease(self, credential_id: Optional[str] = None) -> Optional[str]: every credential is at the soft cap, still return the least-leased one instead of blocking. """ - chosen_id, pending_refresh = self._acquire_lease_under_lock(credential_id) + chosen_id, pending_refresh = self._acquire_lease_under_lock( + credential_id, entry_filter=entry_filter + ) if pending_refresh: self._refresh_pending_entries(pending_refresh) # Mirror select(): a pool whose entries all needed a deferred # refresh must retry once they are back in rotation, or the caller # sees "no credentials available" after a successful refresh. if chosen_id is None: - chosen_id, _ = self._acquire_lease_under_lock(credential_id) + chosen_id, _ = self._acquire_lease_under_lock( + credential_id, entry_filter=entry_filter + ) return chosen_id def _acquire_lease_under_lock( - self, credential_id: Optional[str], + self, + credential_id: Optional[str], + *, + entry_filter: Optional[Callable[[PooledCredential], bool]] = None, ) -> Tuple[Optional[str], List[PooledCredential]]: with self._lock: if credential_id: @@ -2116,6 +2156,12 @@ def _acquire_lease_under_lock( return credential_id, [] available, pending_refresh = self._available_entries(clear_expired=True, refresh=True) + if entry_filter is not None: + try: + available = [entry for entry in available if entry_filter(entry)] + except Exception as exc: + logger.warning("credential pool: lease entry filter failed: %s", exc) + return None, pending_refresh if not available: return None, pending_refresh diff --git a/tests/agent/test_credential_pool_lease_refresh_reselect.py b/tests/agent/test_credential_pool_lease_refresh_reselect.py index 267acf9ae519a..2438efcc8038b 100644 --- a/tests/agent/test_credential_pool_lease_refresh_reselect.py +++ b/tests/agent/test_credential_pool_lease_refresh_reselect.py @@ -91,9 +91,9 @@ def test_acquire_lease_without_pending_refresh_does_not_double_select(): passes = {"n": 0} original = pool._acquire_lease_under_lock - def counting(credential_id): + def counting(credential_id, *, entry_filter=None): passes["n"] += 1 - return original(credential_id) + return original(credential_id, entry_filter=entry_filter) pool._acquire_lease_under_lock = counting @@ -113,3 +113,17 @@ def test_acquire_lease_still_none_when_refresh_does_not_help(): assert pool.acquire_lease() is None assert state["refresh_calls"] == 1, "retry must not refresh repeatedly" assert pool._active_leases == {} + + +def test_acquire_lease_preserves_filter_after_deferred_refresh(): + """The post-refresh retry must not lease an excluded credential.""" + wrong = _entry("wrong") + right = _entry("right") + pool = _bare_pool([wrong, right]) + state = _wire_deferred_refresh(pool) + + lease = pool.acquire_lease(entry_filter=lambda entry: entry.id == "right") + + assert state["refresh_calls"] == 1 + assert lease == "right" + assert pool.leased_entry("right", entry_filter=lambda entry: entry.id == "right") is right diff --git a/tests/tools/test_delegate.py b/tests/tools/test_delegate.py index dd3fee084132d..21acc9b59bd03 100644 --- a/tests/tools/test_delegate.py +++ b/tests/tools/test_delegate.py @@ -17,6 +17,7 @@ import unittest from unittest.mock import MagicMock, patch +from agent.credential_pool import CredentialPool, PooledCredential from tools.delegate_tool import ( DELEGATE_BLOCKED_TOOLS, DELEGATE_TASK_SCHEMA, @@ -1281,6 +1282,24 @@ def test_same_provider_shares_parent_pool(self): result = _resolve_child_credential_pool("openrouter", parent) self.assertIs(result, mock_pool) + def test_same_provider_different_endpoint_refuses_parent_pool(self): + """A child must not lease a same-provider pool for another endpoint.""" + parent = _make_mock_parent() + parent.provider = "openai-api" + parent.base_url = "https://azure.example.com/openai/v1" + parent_pool = MagicMock(provider="openai-api") + parent_pool.entries.return_value = [ + MagicMock(runtime_base_url="https://api.openai.com/v1") + ] + parent._credential_pool = parent_pool + + with patch("agent.credential_pool.load_pool", return_value=None): + result = _resolve_child_credential_pool( + "openai-api", parent, parent.base_url + ) + + self.assertIsNone(result) + # --- Custom-endpoint identity resolution (issue #7833) --- @@ -1314,6 +1333,43 @@ def test_build_child_agent_strict_intersection_when_opted_out(self, mock_cfg): class TestChildCredentialLeasing(unittest.TestCase): + def test_real_pool_leases_and_binds_only_the_child_endpoint(self): + """The least-leased credential at another endpoint cannot be bound.""" + from tools.delegate_tool_child_run import _lease_child_credential + + wrong = PooledCredential( + provider="openai-api", id="wrong", label="wrong", auth_type="api_key", + priority=0, source="manual", access_token="wrong", + base_url="https://api.openai.com/v1", + ) + right = PooledCredential( + provider="openai-api", id="right", label="right", auth_type="api_key", + priority=1, source="manual", access_token="right", + base_url="https://azure.example.com/openai/v1", + ) + pool = CredentialPool("openai-api", [wrong, right]) + original_leased_entry = pool.leased_entry + + def advance_cursor_before_binding(credential_id, *, entry_filter=None): + pool.acquire_lease("wrong") + return original_leased_entry(credential_id, entry_filter=entry_filter) + + pool.leased_entry = advance_cursor_before_binding + child = MagicMock( + provider="openai-api", + base_url="https://azure.example.com/openai/v1", + ) + child._credential_pool = pool + + leased_pool, lease_id = _lease_child_credential(child) + + self.assertIs(leased_pool, pool) + self.assertEqual(lease_id, "right") + child._swap_credential.assert_called_once_with(right) + self.assertIs(pool.current(), wrong) + pool.release_lease("right") + pool.release_lease("wrong") + def test_run_single_child_acquires_and_releases_lease(self): from tools.delegate_tool import _run_single_child diff --git a/tools/delegate_tool_child_run.py b/tools/delegate_tool_child_run.py index 7d49038a8df4b..7f2a2e7059aba 100644 --- a/tools/delegate_tool_child_run.py +++ b/tools/delegate_tool_child_run.py @@ -365,10 +365,38 @@ def _lease_child_credential(child: Any) -> tuple[Any, Optional[str]]: child_pool = getattr(child, "_credential_pool", None) if child_pool is None: return None, None - leased_cred_id = child_pool.acquire_lease() + from tools.delegate_tool_config import ( + _credential_pool_entry_matches_runtime, + _credential_pool_matches_runtime, + ) + + provider = getattr(child, "provider", None) + base_url = getattr(child, "base_url", None) + if not _credential_pool_matches_runtime(child_pool, provider, base_url): + logger.warning( + "Skipping delegated credential pool lease: pool does not match child runtime %s at %s", + provider, base_url, + ) + child._credential_pool = None + return None, None + + from agent.credential_pool import CredentialPool + + real_pool = isinstance(child_pool, CredentialPool) + entry_filter = ( + (lambda entry: _credential_pool_entry_matches_runtime(entry, base_url)) + if real_pool else None + ) + leased_cred_id = ( + child_pool.acquire_lease(entry_filter=entry_filter) + if real_pool else child_pool.acquire_lease() + ) if leased_cred_id is not None: with _quiet("Failed to bind child to leased credential: %s"): - leased_entry = child_pool.current() + leased_entry = ( + child_pool.leased_entry(leased_cred_id, entry_filter=entry_filter) + if real_pool else child_pool.current() + ) if leased_entry is not None and hasattr(child, "_swap_credential"): child._swap_credential(leased_entry) return child_pool, leased_cred_id diff --git a/tools/delegate_tool_config.py b/tools/delegate_tool_config.py index 6fe745e7f9644..1def35ca1d842 100644 --- a/tools/delegate_tool_config.py +++ b/tools/delegate_tool_config.py @@ -206,6 +206,54 @@ def _loaded_pool(key: Any): pool = load_pool(key) return pool if pool is not None and pool.has_credentials() else None + +def _credential_pool_matches_runtime( + pool: Any, + provider: Optional[str], + base_url: Optional[str], +) -> bool: + """Whether a pool can safely supply credentials to this exact runtime. + + A provider name alone does not distinguish OpenAI-compatible endpoints: + an ``openai-api`` pool for api.openai.com must not bind an Azure child. + Lightweight legacy adapters without provider/entry metadata remain usable. + """ + from agent.credential_pool import credential_pool_matches_provider + + raw_pool_provider = getattr(pool, "provider", None) + if not isinstance(raw_pool_provider, str): + return True + if not credential_pool_matches_provider(pool, provider, base_url=base_url): + return False + expected = str(base_url or "").strip().rstrip("/").lower() + if not expected: + return True + entries_fn = getattr(pool, "entries", None) + if not callable(entries_fn): + return True + try: + entries = entries_fn() + if not isinstance(entries, (list, tuple)): + return False + except Exception: + return False + endpoints = [ + value.strip().rstrip("/").lower() + for entry in entries + for value in [getattr(entry, "runtime_base_url", None) or getattr(entry, "base_url", None)] + if isinstance(value, str) and value.strip() + ] + return bool(endpoints) and expected in endpoints + + +def _credential_pool_entry_matches_runtime(entry: Any, base_url: Optional[str]) -> bool: + """Whether one pooled credential targets the child endpoint.""" + expected = str(base_url or "").strip().rstrip("/").lower() + if not expected: + return True + value = getattr(entry, "runtime_base_url", None) or getattr(entry, "base_url", None) + return isinstance(value, str) and value.strip().rstrip("/").lower() == expected + def _resolve_child_credential_pool( effective_provider: Optional[str], parent_agent, effective_base_url: Optional[str] = None, ): @@ -236,8 +284,15 @@ def _resolve_child_credential_pool( return parent_pool return _loaded_pool(child_key) if parent_pool is not None and effective_provider == parent_provider: - return parent_pool - return _loaded_pool(effective_provider) + if _credential_pool_matches_runtime(parent_pool, effective_provider, effective_base_url): + return parent_pool + logger.debug( + "Parent credential pool does not match child runtime %s at %s; resolving independently", + effective_provider, effective_base_url, + ) + pool = _loaded_pool(effective_provider) + if pool is not None and _credential_pool_matches_runtime(pool, effective_provider, effective_base_url): + return pool except Exception as exc: if effective_provider == "custom": logger.debug("Could not resolve custom credential pool for child endpoint '%s': %s", effective_base_url, exc)