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:
Mateus Scheuer Macedo
2026-04-06 22:59:14 -07:00
committed by Teknium
parent 8dee82ea1e
commit f2c11ff30c
4 changed files with 308 additions and 3 deletions

View File

@@ -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()