Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 56 additions & 0 deletions tests/tools/test_delegate.py
Original file line number Diff line number Diff line change
Expand Up @@ -1409,14 +1409,31 @@ 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_parent_provider_rejects_mismatched_parent_pool(self):
parent = _make_mock_parent()
parent.provider = "openai-codex"
mismatched_pool = MagicMock()
mismatched_pool.provider = "xiaomi"
parent._credential_pool = mismatched_pool
own_pool = MagicMock()
own_pool.has_credentials.return_value = True

with patch("agent.credential_pool.load_pool", return_value=own_pool) as mock_load:
result = _resolve_child_credential_pool("openai-codex", parent)

mock_load.assert_called_once_with("openai-codex")
self.assertIs(result, own_pool)

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)
Expand Down Expand Up @@ -1537,7 +1554,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 = {
Expand All @@ -1560,6 +1580,42 @@ 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_cross_provider_leased_credential(self):
from tools.delegate_tool import _run_single_child

leased_entry = MagicMock()
leased_entry.id = "xiaomi-cred"
leased_entry.provider = "xiaomi"
leased_entry.runtime_base_url = "https://token-plan-cn.xiaomimimo.com/v1"
leased_entry.base_url = "https://token-plan-cn.xiaomimimo.com/v1"

child = MagicMock()
child.provider = "openai-codex"
child.base_url = "https://chatgpt.com/backend-api/codex"
child._credential_pool = MagicMock()
child._credential_pool.provider = "xiaomi"
child._credential_pool.acquire_lease.return_value = "xiaomi-cred"
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="Review safely",
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_not_called()
child._credential_pool.release_lease.assert_called_once_with("xiaomi-cred")

def test_run_single_child_releases_lease_after_failure(self):
from tools.delegate_tool import _run_single_child

Expand Down
77 changes: 75 additions & 2 deletions tools/delegate_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -1318,6 +1318,67 @@ def _w(line: str = "") -> None:
return None


def _normalize_provider_name(value: Any) -> str:
"""Return a conservative provider identity for runtime safety checks."""
if not isinstance(value, str):
return ""
return value.strip().lower()


def _provider_identities_match(candidate: Any, expected: Any) -> bool:
"""Return True when two provider identifiers can share credentials.

Most pools use the canonical provider name directly. Named custom-provider
pools are keyed as ``custom:<name>`` while the runtime child provider may be
the bare configured name; treat those as equivalent, but never collapse two
unrelated concrete providers.
"""
candidate_name = _normalize_provider_name(candidate)
expected_name = _normalize_provider_name(expected)
if not candidate_name or not expected_name:
return False
if candidate_name == expected_name:
return True
if (
candidate_name.startswith("custom:")
and candidate_name.removeprefix("custom:") == expected_name
):
return True
if (
expected_name.startswith("custom:")
and expected_name.removeprefix("custom:") == candidate_name
):
return True
return False


def _credential_entry_matches_child_runtime(child: Any, entry: Any) -> bool:
"""Guard against binding a leased credential to the wrong child runtime."""
child_provider = _normalize_provider_name(getattr(child, "provider", None))
entry_provider = _normalize_provider_name(getattr(entry, "provider", None))
pool_provider = _normalize_provider_name(
getattr(getattr(child, "_credential_pool", None), "provider", None)
)
credential_provider = entry_provider or pool_provider
if credential_provider and child_provider and not _provider_identities_match(
credential_provider,
child_provider,
):
return False

child_base = getattr(child, "base_url", None)
entry_base = (
getattr(entry, "runtime_base_url", None)
or getattr(entry, "base_url", None)
)
if isinstance(child_base, str) and isinstance(entry_base, str):
normalized_child_base = child_base.strip().rstrip("/")
normalized_entry_base = entry_base.strip().rstrip("/")
if normalized_child_base and normalized_entry_base:
return normalized_child_base == normalized_entry_base
return True


def _run_single_child(
task_index: int,
goal: str,
Expand Down Expand Up @@ -1349,7 +1410,11 @@ def _run_single_child(
if leased_cred_id is not None:
try:
leased_entry = child_pool.current()
if leased_entry is not None and hasattr(child, "_swap_credential"):
if (
leased_entry is not None
and hasattr(child, "_swap_credential")
and _credential_entry_matches_child_runtime(child, leased_entry)
):
child._swap_credential(leased_entry)
except Exception as exc:
logger.debug("Failed to bind child to leased credential: %s", exc)
Expand Down Expand Up @@ -2325,7 +2390,15 @@ def _resolve_child_credential_pool(effective_provider: Optional[str], parent_age
parent_provider = getattr(parent_agent, "provider", None) or ""
parent_pool = getattr(parent_agent, "_credential_pool", None)
if parent_pool is not None and effective_provider == parent_provider:
return parent_pool
pool_provider = _normalize_provider_name(getattr(parent_pool, "provider", None))
if not pool_provider or _provider_identities_match(pool_provider, effective_provider):
return parent_pool
logger.debug(
"Parent credential pool provider '%s' does not match child provider "
"'%s'; loading child pool",
getattr(parent_pool, "provider", None),
effective_provider,
)

try:
from agent.credential_pool import load_pool
Expand Down
Loading