fix(delegate): share credential pools with subagents + per-task leasing
Cherry-picked from PR #5580 by MestreY0d4-Uninter. - Share parent's credential pool with child agents for key rotation - Leasing layer spreads parallel children across keys (least-loaded) - Thread-safe acquire_lease/release_lease in CredentialPool - Reverted sneaked-in tool-name restoration change (kept original getattr + isinstance guard pattern)
This commit is contained in:
committed by
Teknium
parent
8dee82ea1e
commit
f2c11ff30c
@@ -26,6 +26,7 @@ from tools.delegate_tool import (
|
||||
_build_child_agent,
|
||||
_build_child_system_prompt,
|
||||
_strip_blocked_tools,
|
||||
_resolve_child_credential_pool,
|
||||
_resolve_delegation_credentials,
|
||||
)
|
||||
|
||||
@@ -930,5 +931,126 @@ class TestDelegationProviderIntegration(unittest.TestCase):
|
||||
self.assertEqual(kwargs["base_url"], parent.base_url)
|
||||
|
||||
|
||||
class TestChildCredentialPoolResolution(unittest.TestCase):
|
||||
def test_same_provider_shares_parent_pool(self):
|
||||
parent = _make_mock_parent()
|
||||
mock_pool = MagicMock()
|
||||
parent._credential_pool = mock_pool
|
||||
|
||||
result = _resolve_child_credential_pool("openrouter", parent)
|
||||
self.assertIs(result, mock_pool)
|
||||
|
||||
def test_no_provider_inherits_parent_pool(self):
|
||||
parent = _make_mock_parent()
|
||||
mock_pool = MagicMock()
|
||||
parent._credential_pool = mock_pool
|
||||
|
||||
result = _resolve_child_credential_pool(None, parent)
|
||||
self.assertIs(result, mock_pool)
|
||||
|
||||
def test_different_provider_loads_own_pool(self):
|
||||
parent = _make_mock_parent()
|
||||
parent._credential_pool = MagicMock()
|
||||
mock_pool = MagicMock()
|
||||
mock_pool.has_credentials.return_value = True
|
||||
|
||||
with patch("agent.credential_pool.load_pool", return_value=mock_pool):
|
||||
result = _resolve_child_credential_pool("anthropic", parent)
|
||||
|
||||
self.assertIs(result, mock_pool)
|
||||
|
||||
def test_different_provider_empty_pool_returns_none(self):
|
||||
parent = _make_mock_parent()
|
||||
parent._credential_pool = MagicMock()
|
||||
mock_pool = MagicMock()
|
||||
mock_pool.has_credentials.return_value = False
|
||||
|
||||
with patch("agent.credential_pool.load_pool", return_value=mock_pool):
|
||||
result = _resolve_child_credential_pool("anthropic", parent)
|
||||
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_different_provider_load_failure_returns_none(self):
|
||||
parent = _make_mock_parent()
|
||||
parent._credential_pool = MagicMock()
|
||||
|
||||
with patch("agent.credential_pool.load_pool", side_effect=Exception("disk error")):
|
||||
result = _resolve_child_credential_pool("anthropic", parent)
|
||||
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_build_child_agent_assigns_parent_pool_when_shared(self):
|
||||
parent = _make_mock_parent()
|
||||
mock_pool = MagicMock()
|
||||
parent._credential_pool = mock_pool
|
||||
|
||||
with patch("run_agent.AIAgent") as MockAgent:
|
||||
mock_child = MagicMock()
|
||||
MockAgent.return_value = mock_child
|
||||
|
||||
_build_child_agent(
|
||||
task_index=0,
|
||||
goal="Test pool assignment",
|
||||
context=None,
|
||||
toolsets=["terminal"],
|
||||
model=None,
|
||||
max_iterations=10,
|
||||
parent_agent=parent,
|
||||
)
|
||||
|
||||
self.assertEqual(mock_child._credential_pool, mock_pool)
|
||||
|
||||
|
||||
class TestChildCredentialLeasing(unittest.TestCase):
|
||||
def test_run_single_child_acquires_and_releases_lease(self):
|
||||
from tools.delegate_tool import _run_single_child
|
||||
|
||||
leased_entry = MagicMock()
|
||||
leased_entry.id = "cred-b"
|
||||
|
||||
child = MagicMock()
|
||||
child._credential_pool = MagicMock()
|
||||
child._credential_pool.acquire_lease.return_value = "cred-b"
|
||||
child._credential_pool.current.return_value = leased_entry
|
||||
child.run_conversation.return_value = {
|
||||
"final_response": "done",
|
||||
"completed": True,
|
||||
"interrupted": False,
|
||||
"api_calls": 1,
|
||||
"messages": [],
|
||||
}
|
||||
|
||||
result = _run_single_child(
|
||||
task_index=0,
|
||||
goal="Investigate rate limits",
|
||||
child=child,
|
||||
parent_agent=_make_mock_parent(),
|
||||
)
|
||||
|
||||
self.assertEqual(result["status"], "completed")
|
||||
child._credential_pool.acquire_lease.assert_called_once_with()
|
||||
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_releases_lease_after_failure(self):
|
||||
from tools.delegate_tool import _run_single_child
|
||||
|
||||
child = MagicMock()
|
||||
child._credential_pool = MagicMock()
|
||||
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")
|
||||
|
||||
result = _run_single_child(
|
||||
task_index=1,
|
||||
goal="Trigger failure",
|
||||
child=child,
|
||||
parent_agent=_make_mock_parent(),
|
||||
)
|
||||
|
||||
self.assertEqual(result["status"], "error")
|
||||
child._credential_pool.release_lease.assert_called_once_with("cred-a")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user