diff --git a/tests/tools/test_delegate.py b/tests/tools/test_delegate.py index 96d895038260c..b2c4c11521077 100644 --- a/tests/tools/test_delegate.py +++ b/tests/tools/test_delegate.py @@ -1853,14 +1853,28 @@ class TestChildCredentialPoolResolution(unittest.TestCase): def test_same_provider_shares_parent_pool(self): parent = _make_mock_parent() mock_pool = MagicMock() + mock_pool.provider = "openrouter" parent._credential_pool = mock_pool result = _resolve_child_credential_pool("openrouter", parent) self.assertIs(result, mock_pool) + def test_same_provider_rejects_mismatched_parent_pool(self): + parent = _make_mock_parent() + parent.provider = "deepseek" + parent._credential_pool = MagicMock() + parent._credential_pool.provider = "zai" + + with patch("agent.credential_pool.load_pool", return_value=None) as load_mock: + result = _resolve_child_credential_pool("deepseek", parent) + + self.assertIsNone(result) + load_mock.assert_called_once_with("deepseek") + def test_no_provider_inherits_parent_pool(self): parent = _make_mock_parent() mock_pool = MagicMock() + mock_pool.provider = "openrouter" parent._credential_pool = mock_pool result = _resolve_child_credential_pool(None, parent) @@ -1906,6 +1920,7 @@ def test_custom_different_endpoint_does_not_inherit_parent_pool(self): parent.provider = "custom" parent.base_url = "https://endpoint-a.example.com/v1" parent._credential_pool = MagicMock(name="parent_custom_a_pool") + parent._credential_pool.provider = "custom:endpoint-a" child_pool = MagicMock(name="endpoint_b_pool") child_pool.has_credentials.return_value = True @@ -1934,6 +1949,7 @@ def test_custom_same_endpoint_shares_parent_pool(self): parent.provider = "custom" parent.base_url = "https://endpoint-a.example.com/v1" parent._credential_pool = MagicMock(name="parent_custom_a_pool") + parent._credential_pool.provider = "custom:endpoint-a" with patch( "agent.credential_pool.get_custom_provider_pool_key", @@ -1945,6 +1961,24 @@ def test_custom_same_endpoint_shares_parent_pool(self): self.assertIs(result, parent._credential_pool) + def test_custom_endpoint_rejects_parent_pool_for_different_endpoint(self): + parent = _make_mock_parent() + parent.provider = "custom" + parent.base_url = "https://endpoint-a.example.com/v1" + parent._credential_pool = MagicMock() + parent._credential_pool.provider = "custom:endpoint-b" + + with patch( + "agent.credential_pool.get_custom_provider_pool_key", + return_value="custom:endpoint-a", + ), patch("agent.credential_pool.load_pool", return_value=None) as load_mock: + result = _resolve_child_credential_pool( + "custom", parent, "https://endpoint-a.example.com/v1" + ) + + self.assertIsNone(result) + load_mock.assert_called_once_with("custom:endpoint-a") + def test_custom_unregistered_endpoint_returns_none(self): """A raw delegation.base_url with no matching custom_providers entry must NOT inherit the parent's pool — return None so the child keeps its @@ -1967,6 +2001,7 @@ def test_custom_unregistered_endpoint_returns_none(self): def test_build_child_agent_assigns_parent_pool_when_shared(self): parent = _make_mock_parent() mock_pool = MagicMock() + mock_pool.provider = "openrouter" parent._credential_pool = mock_pool with patch("run_agent.AIAgent") as MockAgent: @@ -2048,7 +2083,10 @@ def test_run_single_child_acquires_and_releases_lease(self): leased_entry.id = "cred-b" child = MagicMock() + child.provider = "openrouter" + child.base_url = "https://openrouter.ai/api/v1" child._credential_pool = MagicMock() + child._credential_pool.provider = "openrouter" child._credential_pool.acquire_lease.return_value = "cred-b" child._credential_pool.current.return_value = leased_entry child.run_conversation.return_value = { @@ -2071,11 +2109,75 @@ def test_run_single_child_acquires_and_releases_lease(self): child._swap_credential.assert_called_once_with(leased_entry) child._credential_pool.release_lease.assert_called_once_with("cred-b") + def test_run_single_child_skips_mismatched_credential_pool(self): + from tools.delegate_tool import _run_single_child + + child = MagicMock() + child.provider = "deepseek" + child._credential_pool = MagicMock() + child._credential_pool.provider = "zai" + child._credential_pool.acquire_lease.return_value = "zai-cred" + child.run_conversation.return_value = { + "final_response": "done", + "completed": True, + "interrupted": False, + "api_calls": 1, + "messages": [], + } + + result = _run_single_child( + task_index=0, + goal="Use delegated model", + child=child, + parent_agent=_make_mock_parent(), + ) + + self.assertEqual(result["status"], "completed") + child._credential_pool.acquire_lease.assert_not_called() + child._swap_credential.assert_not_called() + child._credential_pool.release_lease.assert_not_called() + + def test_run_single_child_skips_custom_pool_for_different_endpoint(self): + from tools.delegate_tool import _run_single_child + + child = MagicMock() + child.provider = "custom" + child.base_url = "https://endpoint-a.example.com/v1" + child._credential_pool = MagicMock() + child._credential_pool.provider = "custom:endpoint-b" + child._credential_pool.acquire_lease.return_value = "endpoint-b-cred" + child.run_conversation.return_value = { + "final_response": "done", + "completed": True, + "interrupted": False, + "api_calls": 1, + "messages": [], + } + + with patch( + "agent.credential_pool.get_custom_provider_pool_key", + return_value="custom:endpoint-a", + ): + result = _run_single_child( + task_index=0, + goal="Use endpoint A", + child=child, + parent_agent=_make_mock_parent(), + ) + + self.assertEqual(result["status"], "completed") + child._credential_pool.acquire_lease.assert_not_called() + child._swap_credential.assert_not_called() + child._credential_pool.release_lease.assert_not_called() + def test_run_single_child_releases_lease_after_failure(self): from tools.delegate_tool import _run_single_child child = MagicMock() + child.provider = "openrouter" + child.base_url = "https://openrouter.ai/api/v1" child._credential_pool = MagicMock() + child._credential_pool.provider = "openrouter" child._credential_pool.acquire_lease.return_value = "cred-a" child._credential_pool.current.return_value = MagicMock(id="cred-a") child.run_conversation.side_effect = RuntimeError("boom") diff --git a/tools/delegate_tool.py b/tools/delegate_tool.py index 12f94a180a8ef..5eb947e5a266a 100644 --- a/tools/delegate_tool.py +++ b/tools/delegate_tool.py @@ -31,6 +31,7 @@ ) from typing import Any, Dict, List, Optional +from agent.credential_pool import credential_pool_matches_provider from toolsets import TOOLSETS # Sentinel value used by the runtime provider system for providers that are @@ -1813,6 +1814,17 @@ def _run_single_child( ) child_pool = getattr(child, "_credential_pool", None) + if child_pool is not None and not credential_pool_matches_provider( + child_pool, + getattr(child, "provider", None), + base_url=getattr(child, "base_url", None), + ): + logger.debug( + "Skipping child credential pool for provider mismatch: child=%r pool=%r", + getattr(child, "provider", None), + getattr(child_pool, "provider", None), + ) + child_pool = None leased_cred_id = None if child_pool is not None: leased_cred_id = child_pool.acquire_lease() @@ -3036,11 +3048,18 @@ def _resolve_child_credential_pool( ``custom:`` pool key derived from the base_url) and only share the parent's pool when both resolve to the *same* custom endpoint. """ - if not effective_provider: - return getattr(parent_agent, "_credential_pool", None) - parent_provider = getattr(parent_agent, "provider", None) or "" parent_pool = getattr(parent_agent, "_credential_pool", None) + if not effective_provider: + return ( + parent_pool + if credential_pool_matches_provider( + parent_pool, + parent_provider, + base_url=effective_base_url, + ) + else None + ) # Custom endpoints: distinguish by endpoint identity, not the bare "custom" # provider string. Two custom runtimes are only interchangeable when they @@ -3066,6 +3085,11 @@ def _resolve_child_credential_pool( and parent_provider == "custom" and parent_key is not None and parent_key == child_key + and credential_pool_matches_provider( + parent_pool, + effective_provider, + base_url=effective_base_url, + ) ): return parent_pool @@ -3080,7 +3104,15 @@ def _resolve_child_credential_pool( ) return None - if parent_pool is not None and effective_provider == parent_provider: + if ( + parent_pool is not None + and effective_provider == parent_provider + and credential_pool_matches_provider( + parent_pool, + effective_provider, + base_url=effective_base_url, + ) + ): return parent_pool try: