From 34496edd1e408dc2dcb8003bf26ee3fc1301dbd1 Mon Sep 17 00:00:00 2001 From: Roger Perez Garcia Date: Tue, 18 Aug 2026 09:39:34 -0400 Subject: [PATCH] feat(delegation): enforce per-child tool allowlists --- agent/tool_executor.py | 18 ++- run_agent.py | 1 + .../test_tool_call_guardrail_runtime.py | 47 ++++++ tests/tools/test_delegate.py | 153 ++++++++++++++++++ tests/tools/test_refresh_agent_mcp_tools.py | 22 +++ tools/delegate_tool.py | 95 +++++++++++ tools/mcp_tool.py | 14 ++ 7 files changed, 349 insertions(+), 1 deletion(-) diff --git a/agent/tool_executor.py b/agent/tool_executor.py index 381f1000e9bbc..5eedec0cf5c5c 100644 --- a/agent/tool_executor.py +++ b/agent/tool_executor.py @@ -596,6 +596,17 @@ def _advance_start_order(callback=None) -> None: block_message = scope_block block_error_type = "tool_scope_block" + tool_allowlist = getattr(agent, "_tool_allowlist", None) + if ( + block_message is None + and tool_allowlist is not None + and function_name not in tool_allowlist + ): + block_message = ( + f"Tool '{function_name}' is not permitted by this delegated " + "agent's runtime tool_allowlist." + ) + block_error_type = "tool_allowlist_block" if block_message is None: block_error_type = "plugin_block" @@ -640,7 +651,12 @@ def _resolve_pre_tool_block(): _advance_start_order() state["blocked"] = True if block_message is not None: - result = json.dumps({"error": block_message}, ensure_ascii=False) + error_payload = {"error": block_message} + if block_error_type == "tool_allowlist_block": + error_payload.update( + {"error_type": block_error_type, "tool": function_name} + ) + result = json.dumps(error_payload, ensure_ascii=False) error_type = block_error_type error_message = block_message else: diff --git a/run_agent.py b/run_agent.py index b09262eae4448..086407139878e 100644 --- a/run_agent.py +++ b/run_agent.py @@ -8238,6 +8238,7 @@ def _dispatch_delegate_task(self, function_args: dict) -> str: max_iterations=function_args.get("max_iterations"), role=function_args.get("role"), background=(not _is_subagent), + tool_allowlist=function_args.get("tool_allowlist"), action=function_args.get("action"), subagent_id=function_args.get("subagent_id"), message=function_args.get("message"), diff --git a/tests/run_agent/test_tool_call_guardrail_runtime.py b/tests/run_agent/test_tool_call_guardrail_runtime.py index ca6e80aac4c4f..de3ed16d72066 100644 --- a/tests/run_agent/test_tool_call_guardrail_runtime.py +++ b/tests/run_agent/test_tool_call_guardrail_runtime.py @@ -136,6 +136,53 @@ def test_config_enabled_hard_stop_blocks_repeated_exact_failure_before_execution assert "repeated_exact_failure_block" in messages[0]["content"] +def test_tool_allowlist_blocks_disallowed_sequential_call_before_dispatch(): + agent = _make_agent("read_file", "write_file", "patch", "terminal") + agent._tool_allowlist = frozenset({"read_file", "search_files", "grep"}) + tc = _mock_tool_call( + "write_file", + json.dumps({"path": "/tmp/must-not-exist", "content": "blocked"}), + "c-allowlist-sequential", + ) + msg = SimpleNamespace(content="", tool_calls=[tc]) + messages = [] + + with patch("run_agent.handle_function_call", return_value="SHOULD_NOT_RUN") as mock_hfc: + agent._execute_tool_calls_sequential(msg, messages, "task-allowlist") + + mock_hfc.assert_not_called() + result = json.loads(messages[0]["content"]) + assert result["error_type"] == "tool_allowlist_block" + assert result["tool"] == "write_file" + + +def test_tool_allowlist_blocks_disallowed_concurrent_calls_before_dispatch(): + agent = _make_agent( + "read_file", "write_file", "patch", "terminal", "mcp__roshhome__update_request" + ) + agent._tool_allowlist = frozenset({"read_file", "search_files", "grep"}) + calls = [ + _mock_tool_call("patch", "{}", "c-allowlist-patch"), + _mock_tool_call("terminal", "{}", "c-allowlist-terminal"), + _mock_tool_call("mcp__roshhome__update_request", "{}", "c-allowlist-mcp"), + ] + msg = SimpleNamespace(content="", tool_calls=calls) + messages = [] + + with patch("run_agent.handle_function_call", return_value="SHOULD_NOT_RUN") as mock_hfc: + agent._execute_tool_calls_concurrent(msg, messages, "task-allowlist") + + mock_hfc.assert_not_called() + contents = [message["content"] for message in messages] + assert all('"error_type": "tool_allowlist_block"' in item for item in contents) + assert all( + f'"tool": "{tool_name}"' in content + for tool_name, content in zip( + ["patch", "terminal", "mcp__roshhome__update_request"], contents + ) + ) + + def test_sequential_after_call_appends_guidance_to_tool_result_without_extra_messages(): agent = _make_agent("web_search") args = {"query": "same"} diff --git a/tests/tools/test_delegate.py b/tests/tools/test_delegate.py index 0d7030ddb0c32..b9aaed2a05877 100644 --- a/tests/tools/test_delegate.py +++ b/tests/tools/test_delegate.py @@ -70,6 +70,10 @@ def test_schema_valid(self): # capability-selection surface the model should not control. self.assertNotIn("toolsets", props) self.assertNotIn("toolsets", props["tasks"]["items"]["properties"]) + self.assertIn("tool_allowlist", props) + self.assertIn( + "tool_allowlist", props["tasks"]["items"]["properties"] + ) # max_iterations is intentionally NOT exposed to the model — it's # config-authoritative via delegation.max_iterations so users get # predictable budgets. @@ -251,6 +255,80 @@ def test_orchestrator_composite_regains_only_delegate_task(self): ) +class TestChildToolAllowlist(unittest.TestCase): + @staticmethod + def _tool(name): + return { + "type": "function", + "function": {"name": name, "description": "", "parameters": {}}, + } + + def _build(self, tool_allowlist): + parent = _make_mock_parent() + parent.enabled_toolsets = ["file", "terminal", "mcp-roshhome"] + child = MagicMock() + child.tools = [ + self._tool("read_file"), + self._tool("write_file"), + self._tool("terminal"), + self._tool("mcp__roshhome__update_request"), + ] + child.valid_tool_names = { + "read_file", + "write_file", + "terminal", + "mcp__roshhome__update_request", + } + child._context_engine_tool_names = set() + + with patch("run_agent.AIAgent", return_value=child): + result = _build_child_agent( + task_index=0, + goal="Inspect safely", + context=None, + toolsets=None, + model=None, + max_iterations=10, + task_count=1, + parent_agent=parent, + tool_allowlist=tool_allowlist, + ) + return result + + def test_allowlist_intersects_final_builtin_and_mcp_snapshot(self): + child = self._build(["read_file", "search_files", "grep"]) + + self.assertEqual(child.valid_tool_names, {"read_file"}) + self.assertEqual( + [tool["function"]["name"] for tool in child.tools], ["read_file"] + ) + self.assertEqual( + child._tool_allowlist, + frozenset({"read_file", "search_files", "grep"}), + ) + + def test_empty_allowlist_is_deny_all(self): + child = self._build([]) + + self.assertEqual(child.tools, []) + self.assertEqual(child.valid_tool_names, set()) + self.assertEqual(child._tool_allowlist, frozenset()) + + def test_absent_allowlist_preserves_current_snapshot(self): + child = self._build(None) + + self.assertEqual( + child.valid_tool_names, + { + "read_file", + "write_file", + "terminal", + "mcp__roshhome__update_request", + }, + ) + self.assertIsNone(child._tool_allowlist) + + class TestDelegateTask(unittest.TestCase): def test_no_parent_agent(self): result = json.loads(delegate_task(goal="test")) @@ -1366,11 +1444,13 @@ def fake_delegate_task(**kwargs): parent, { "goal": "test", + "tool_allowlist": ["read_file", "search_files"], "acp_command": "claude", "acp_args": ["--acp", "--stdio"], "tasks": [ { "goal": "nested", + "tool_allowlist": [], "acp_command": "codex", "acp_args": ["--acp"], }, @@ -1381,6 +1461,8 @@ def fake_delegate_task(**kwargs): self.assertNotIn("acp_command", captured) self.assertNotIn("acp_args", captured) self.assertEqual(captured["goal"], "test") + self.assertEqual(captured["tool_allowlist"], ["read_file", "search_files"]) + self.assertEqual(captured["tasks"][0]["tool_allowlist"], []) self.assertNotIn("acp_command", captured["tasks"][0]) self.assertNotIn("acp_args", captured["tasks"][0]) @@ -1620,6 +1702,77 @@ def test_orchestrator_role_keeps_delegation_at_depth_1( self.assertIn("delegation", kwargs["enabled_toolsets"]) self.assertEqual(mock_child._delegate_role, "orchestrator") + @patch("tools.delegate_tool._resolve_delegation_credentials") + @patch("tools.delegate_tool._load_config", return_value={"max_spawn_depth": 2}) + def test_top_level_tool_allowlist_reaches_child_runtime( + self, mock_cfg, mock_creds + ): + mock_creds.return_value = { + "provider": None, "base_url": None, + "api_key": None, "api_mode": None, "model": None, + } + parent = _make_mock_parent(depth=0) + parent.enabled_toolsets = ["terminal", "file"] + with patch("run_agent.AIAgent") as MockAgent: + mock_child = _make_role_mock_child() + mock_child.tools = [] + mock_child.valid_tool_names = set() + MockAgent.return_value = mock_child + + delegate_task( + goal="Inspect safely", + tool_allowlist=["read_file", "search_files", "grep"], + parent_agent=parent, + ) + + self.assertEqual( + mock_child._tool_allowlist, + frozenset({"read_file", "search_files", "grep"}), + ) + + @patch("tools.delegate_tool._resolve_delegation_credentials") + @patch("tools.delegate_tool._load_config", return_value={"max_spawn_depth": 2}) + def test_per_task_tool_allowlist_overrides_top_level( + self, mock_cfg, mock_creds + ): + mock_creds.return_value = { + "provider": None, "base_url": None, + "api_key": None, "api_mode": None, "model": None, + } + parent = _make_mock_parent(depth=0) + parent.enabled_toolsets = ["terminal", "file"] + children = [_make_role_mock_child(), _make_role_mock_child()] + for child in children: + child.tools = [] + child.valid_tool_names = set() + + with patch("run_agent.AIAgent", side_effect=children): + delegate_task( + tasks=[ + {"goal": "Inspect the first component thoroughly"}, + { + "goal": "Inspect the second component thoroughly", + "tool_allowlist": [], + }, + ], + tool_allowlist=["read_file"], + parent_agent=parent, + ) + + self.assertEqual(children[0]._tool_allowlist, frozenset({"read_file"})) + self.assertEqual(children[1]._tool_allowlist, frozenset()) + + def test_invalid_tool_allowlist_fails_before_spawn(self): + parent = _make_mock_parent(depth=0) + result = json.loads( + delegate_task( + goal="Inspect safely", + tool_allowlist="read_file", + parent_agent=parent, + ) + ) + self.assertIn("tool_allowlist must be an array", result["error"]) + @patch("tools.delegate_tool._resolve_delegation_credentials") @patch("tools.delegate_tool._load_config", return_value={"max_spawn_depth": 2}) diff --git a/tests/tools/test_refresh_agent_mcp_tools.py b/tests/tools/test_refresh_agent_mcp_tools.py index b1aa95e12f560..394d5c045e5b5 100644 --- a/tests/tools/test_refresh_agent_mcp_tools.py +++ b/tests/tools/test_refresh_agent_mcp_tools.py @@ -44,6 +44,28 @@ def test_refresh_adds_late_landing_tools(monkeypatch): assert len(agent.tools) == 3 +def test_refresh_intersects_late_mcp_tools_with_agent_allowlist(monkeypatch): + """A late MCP refresh must not reopen a delegated child's hard boundary.""" + agent = _agent(["read_file"]) + agent._tool_allowlist = frozenset({"read_file", "search_files", "grep"}) + + import model_tools + monkeypatch.setattr( + model_tools, + "get_tool_definitions", + lambda **kw: [ + _tool("read_file"), + _tool("mcp__roshhome__update_request"), + ], + ) + + added = mcp_tool.refresh_agent_mcp_tools(agent) + + assert added == set() + assert agent.valid_tool_names == {"read_file"} + assert [tool["function"]["name"] for tool in agent.tools] == ["read_file"] + + def test_refresh_preserves_memory_provider_and_context_engine_tools(monkeypatch): """B1 regression: a rebuild must NOT drop post-build-injected tools. diff --git a/tools/delegate_tool.py b/tools/delegate_tool.py index 8699edb984707..cf55fa70b9631 100644 --- a/tools/delegate_tool.py +++ b/tools/delegate_tool.py @@ -1296,6 +1296,49 @@ def _strip_blocked_tools(toolsets: List[str]) -> List[str]: return [t for t in toolsets if t not in blocked_toolset_names] +def _normalize_tool_allowlist( + value: Any, *, field: str = "tool_allowlist" +) -> tuple[Optional[List[str]], Optional[str]]: + """Validate and de-duplicate an optional per-child tool-name allowlist.""" + if value is None: + return None, None + if not isinstance(value, list): + return None, f"{field} must be an array of non-empty tool names." + + normalized: List[str] = [] + seen = set() + for index, item in enumerate(value): + if not isinstance(item, str) or not item.strip(): + return None, f"{field}[{index}] must be a non-empty string tool name." + name = item.strip() + if name not in seen: + normalized.append(name) + seen.add(name) + return normalized, None + + +def _apply_child_tool_allowlist(child: Any, tool_allowlist: Optional[List[str]]) -> None: + """Publish a deny-all, exact-name tool boundary on a delegated child.""" + allowset = None if tool_allowlist is None else frozenset(tool_allowlist) + child._tool_allowlist = allowset + if allowset is None: + return + + child.tools = [ + tool + for tool in (getattr(child, "tools", None) or []) + if tool.get("function", {}).get("name") in allowset + ] + child.valid_tool_names = { + tool["function"]["name"] for tool in child.tools + } + engine_names = getattr(child, "_context_engine_tool_names", None) + if isinstance(engine_names, set): + engine_names.intersection_update(allowset) + # A precomputed deferred-tool scope must never outlive the narrowed surface. + child._tool_search_scope_cache = None + + def _blocked_toolsets_for_role(role: str) -> List[str]: """Return one-tool deny toolsets for a delegated child role. @@ -1598,6 +1641,9 @@ def _build_child_agent( # 'leaf' (default) cannot; 'orchestrator' retains the delegation # toolset subject to depth/kill-switch bounds applied below. role: str = "leaf", + # Exact tool-name capability boundary. None preserves legacy behavior; + # an empty list is intentional deny-all. + tool_allowlist: Optional[List[str]] = None, ): """ Build a child AIAgent on the main thread (thread-safe construction). @@ -1985,6 +2031,10 @@ def _child_thinking(text: str) -> None: except Exception: pass raise + # Apply the exact-name boundary only after AIAgent has materialized its + # complete builtin + MCP + post-build tool snapshot. This is a runtime + # capability restriction, not a prompt instruction. + _apply_child_tool_allowlist(child, tool_allowlist) child._print_fn = getattr(parent_agent, "_print_fn", None) # Ownership transfer for the dedicated handle: the child's close() must # release it (nothing else holds a reference), and no parent teardown can @@ -3602,6 +3652,7 @@ def delegate_task( role: Optional[str] = None, background: Optional[bool] = None, output_schema: Optional[Dict[str, Any]] = None, + tool_allowlist: Optional[List[str]] = None, action: Optional[str] = None, subagent_id: Optional[str] = None, message: Optional[str] = None, @@ -3733,6 +3784,8 @@ def delegate_task( single_task: Dict[str, Any] = {"goal": goal, "context": context, "role": top_role} if output_schema is not None: single_task["output_schema"] = output_schema + if tool_allowlist is not None: + single_task["tool_allowlist"] = tool_allowlist task_list = [single_task] else: return tool_error("Provide either 'goal' (single task) or 'tasks' (batch).") @@ -3749,6 +3802,25 @@ def delegate_task( if not task.get("goal", "").strip(): return tool_error(f"Task {i} is missing a 'goal'.") + normalized_top_allowlist, allowlist_error = _normalize_tool_allowlist( + tool_allowlist + ) + if allowlist_error: + return tool_error(allowlist_error) + task_tool_allowlists: List[Optional[List[str]]] = [] + for i, task in enumerate(task_list): + raw_allowlist = ( + task["tool_allowlist"] + if "tool_allowlist" in task + else normalized_top_allowlist + ) + normalized_allowlist, allowlist_error = _normalize_tool_allowlist( + raw_allowlist, field=f"tasks[{i}].tool_allowlist" + ) + if allowlist_error: + return tool_error(allowlist_error) + task_tool_allowlists.append(normalized_allowlist) + # Batch-only quality gate: catch malformed fan-outs (placeholder goals, # unexpanded multi-word template markers, 1-task batches) before any # child is spawned. The single-`goal` form is deliberately exempt — @@ -3858,6 +3930,7 @@ def delegate_task( override_acp_command=creds.get("command"), override_acp_args=creds.get("args"), role=effective_role, + tool_allowlist=task_tool_allowlists[i], ) except ValueError as exc: # Explicit-pin preflight failures (e.g. pinned delegation.command @@ -4776,6 +4849,17 @@ def _build_dynamic_schema_overrides() -> dict: "enum": ["leaf", "orchestrator"], "description": "Per-task role override. See top-level 'role' for semantics.", }, + "tool_allowlist": { + "type": "array", + "items": {"type": "string"}, + "description": ( + "Optional exact tool-name allowlist for this child. " + "When present, the runtime exposes and executes only " + "the intersection of the inherited final builtin/MCP " + "surface with these names. An empty list denies all " + "tools. Per-task value overrides the top-level value." + ), + }, "output_schema": { "type": "object", "description": ( @@ -4802,6 +4886,16 @@ def _build_dynamic_schema_overrides() -> dict: "enum": ["leaf", "orchestrator"], "description": "(rebuilt at get_definitions() time)", }, + "tool_allowlist": { + "type": "array", + "items": {"type": "string"}, + "description": ( + "Optional exact tool-name allowlist for spawned children. " + "When present, deny all tools except names in this list after " + "inheritance, blocked-tool removal, and MCP resolution. " + "An empty list denies every tool." + ), + }, "output_schema": { "type": "object", "description": ( @@ -4915,6 +5009,7 @@ def _strip_model_hidden_task_fields(tasks: Any) -> Any: role=args.get("role"), background=_model_background_value(args, kw.get("parent_agent")), output_schema=args.get("output_schema"), + tool_allowlist=args.get("tool_allowlist"), action=args.get("action"), subagent_id=args.get("subagent_id"), message=args.get("message"), diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index 93b4f2585cff8..d5312c506d7c0 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -7758,6 +7758,20 @@ def refresh_agent_mcp_tools( # this rebuild actually appended (matching agent_init's dedup-aware add). staged_engine_names = _reinject_post_build_tools(agent, new_defs, new_names) + # Delegated children may carry an exact-name runtime capability boundary. + # Reapply it after registry + MCP + post-build injection so late MCP + # connections cannot reopen tools that were absent from the allowlist. + tool_allowlist = getattr(agent, "_tool_allowlist", None) + if tool_allowlist is not None: + allowed_names = frozenset(tool_allowlist) + new_defs = [ + tool + for tool in new_defs + if tool.get("function", {}).get("name") in allowed_names + ] + new_names = {tool["function"]["name"] for tool in new_defs} + staged_engine_names.intersection_update(allowed_names) + # Single atomic read-diff-publish so the returned ``added`` is consistent # with what was actually published, even under concurrent callers, and a # stale (older-generation) rebuild can't overwrite a newer published one.