Przeglądaj źródła

收紧 recursive 子 Agent 工具权限边界

SamLee 1 dzień temu
rodzic
commit
b3407ceb14

+ 4 - 0
cyber_agent/core/agent_mode.py

@@ -19,6 +19,10 @@ REMOVED_RECURSION_ENV = "AGENT_RECURSION_ENABLED"
 AGENT_MODE_CONTEXT_KEY = "agent_mode"
 AGENT_MODE_REVISION_CONTEXT_KEY = "agent_mode_revision"
 
+# Runner 注入给本地 agent 工具的内部上下文。它们不来自 LLM 参数,也不持久化
+# 到 Trace;Recursive 子级只能继承调用方在本次运行中实际拥有的能力与调度策略。
+RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY = "_recursive_capability_tool_names"
+
 
 class AgentMode(str, Enum):
     LEGACY = "legacy"

+ 150 - 69
cyber_agent/core/runner.py

@@ -41,6 +41,8 @@ from cyber_agent.tools.builtin.knowledge import KnowledgeConfig
 from cyber_agent.core.memory import MemoryConfig
 from cyber_agent.core.agent_mode import (
     AGENT_MODE_CONTEXT_KEY,
+    AgentMode,
+    RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY,
     apply_policy_to_context,
     assert_removed_config_absent,
     policy_from_context,
@@ -1083,7 +1085,12 @@ class AgentRunner:
     ) -> AsyncIterator[Union[Trace, Message]]:
         """ReAct 循环"""
         trace_id = trace.trace_id
-        tool_schemas = self._get_runtime_tool_schemas(config, trace)
+        runtime_tool_names = self._get_runtime_tool_names(config, trace)
+        tool_schemas = self._get_runtime_tool_schemas(
+            config,
+            trace,
+            runtime_tool_names=runtime_tool_names,
+        )
         completion_status = "completed"
 
         # 当前主路径头节点的 sequence(用于设置 parent_sequence)
@@ -1284,10 +1291,22 @@ class AgentRunner:
                 fresh_trace = await self.trace_store.get_trace(trace_id)
                 if fresh_trace:
                     trace = fresh_trace
+            runtime_policy = policy_from_context(trace.context)
+            runtime_tool_names = self._get_runtime_tool_names(
+                config,
+                trace,
+                in_side_branch=side_branch_ctx is not None,
+            )
             tool_schemas = self._get_runtime_tool_schemas(
                 config,
                 trace,
                 in_side_branch=side_branch_ctx is not None,
+                runtime_tool_names=runtime_tool_names,
+            )
+            dispatch_allowlist = (
+                runtime_tool_names
+                if runtime_policy.mode is AgentMode.RECURSIVE
+                else None
             )
 
             # 构建 LLM messages(注入上下文,移除内部字段)
@@ -1327,7 +1346,6 @@ class AgentRunner:
                 1 for tc in (tool_calls or [])
                 if tc.get("function", {}).get("name") in lifecycle_tools
             )
-            runtime_policy = policy_from_context(trace.context)
             protocol_batch_error = None
             if (
                 runtime_policy.requires_task_protocol
@@ -1366,6 +1384,10 @@ class AgentRunner:
                     not has_context_call
                     and not lifecycle_call_count
                     and not protocol_lifecycle_required
+                    and (
+                        runtime_policy.mode is not AgentMode.RECURSIVE
+                        or "get_current_context" in runtime_tool_names
+                    )
                 ):
                     # 手动添加 get_current_context 工具调用
                     context_call_id = f"call_context_{uuid.uuid4().hex[:8]}"
@@ -1383,6 +1405,10 @@ class AgentRunner:
                 and iteration == 0
                 and not lifecycle_call_count
                 and not protocol_lifecycle_required
+                and (
+                    runtime_policy.mode is not AgentMode.RECURSIVE
+                    or "skill" in runtime_tool_names
+                )
             ):
                 skills_to_inject = self._check_skills_need_injection(
                     trace, inject_skills, history, skill_recency_threshold
@@ -1757,8 +1783,20 @@ class AgentRunner:
                                 pass
                         try:
                             tool_result = await self.tools.execute(
-                                tool_name, tool_args, uid=config.uid or "",
-                                context={"store": self.trace_store, "trace_id": trace_id, "goal_id": current_goal_id, "runner": self, "goal_tree": goal_tree, "knowledge_config": config.knowledge, "sequence": sequence, "side_branch": {"type": side_branch_ctx.type, "branch_id": side_branch_ctx.branch_id, "is_side_branch": True, "max_turns": side_branch_ctx.max_turns, "trigger_event": trigger_event_for_tool} if side_branch_ctx else None, **(config.context or {})}
+                                tool_name,
+                                tool_args,
+                                uid=config.uid or "",
+                                context=self._build_tool_context(
+                                    config=config,
+                                    trace=trace,
+                                    trace_id=trace_id,
+                                    goal_id=current_goal_id,
+                                    goal_tree=goal_tree,
+                                    sequence=sequence,
+                                    side_branch_ctx=side_branch_ctx,
+                                    trigger_event=trigger_event_for_tool,
+                                ),
+                                allowed_tool_names=dispatch_allowlist,
                             )
                             return (tc, tool_args, tool_result)
                         except Exception as e:
@@ -1888,25 +1926,17 @@ class AgentRunner:
                                 tool_name,
                                 tool_args,
                                 uid=config.uid or "",
-                                context={
-                                "store": self.trace_store,
-                                "trace_id": trace_id,
-                                "goal_id": current_goal_id,
-                                "runner": self,
-                                "goal_tree": goal_tree,
-                                "knowledge_config": config.knowledge,
-                                "sequence": sequence,  # 添加sequence用于知识注入记录
-                                # 新增:侧分支信息
-                                "side_branch": {
-                                    "type": side_branch_ctx.type,
-                                    "branch_id": side_branch_ctx.branch_id,
-                                    "is_side_branch": True,
-                                    "max_turns": side_branch_ctx.max_turns,
-                                    "trigger_event": trigger_event_for_tool,
-                                } if side_branch_ctx else None,
-                                # 合并用户自定义 context(RunConfig.context)
-                                **(config.context or {}),
-                                },
+                                context=self._build_tool_context(
+                                    config=config,
+                                    trace=trace,
+                                    trace_id=trace_id,
+                                    goal_id=current_goal_id,
+                                    goal_tree=goal_tree,
+                                    sequence=sequence,
+                                    side_branch_ctx=side_branch_ctx,
+                                    trigger_event=trigger_event_for_tool,
+                                ),
+                                allowed_tool_names=dispatch_allowlist,
                             )
     
                         # 如果是 goal 工具,记录执行后的状态
@@ -3212,22 +3242,13 @@ class AgentRunner:
         )
         return messages
 
-    def _get_tool_schemas(
+    def _get_configured_tool_names(
         self,
         tools: Optional[List[str]] = None,
         tool_groups: Optional[List[str]] = None,
         exclude_tools: Optional[List[str]] = None,
-    ) -> List[Dict]:
-        """
-        获取工具 Schema
-
-        合并策略(取并集):
-        - tool_groups 非空: 按分组白名单过滤得到基础工具集
-        - tools 非空: 追加指定的工具名(与 tool_groups 结果取并集)
-        - 两者都为 None: 返回所有已注册工具
-
-        最后再用 exclude_tools 减去禁用的工具(如远程 agent 禁止 agent/evaluate)。
-        """
+    ) -> set[str]:
+        """解析 RunConfig 的基础工具能力集合。"""
         if tool_groups is not None:
             tool_names = set(self.tools.get_tool_names(groups=tool_groups))
         else:
@@ -3236,54 +3257,72 @@ class AgentRunner:
             tool_names |= set(tools)
         if exclude_tools:
             tool_names -= set(exclude_tools)
+        return tool_names
+
+    def _get_tool_schemas(
+        self,
+        tools: Optional[List[str]] = None,
+        tool_groups: Optional[List[str]] = None,
+        exclude_tools: Optional[List[str]] = None,
+    ) -> List[Dict]:
+        """获取 RunConfig 基础能力对应的工具 Schema。"""
+        tool_names = self._get_configured_tool_names(
+            tools,
+            tool_groups,
+            exclude_tools,
+        )
         return self.tools.get_schemas(list(tool_names))
 
-    def _get_runtime_tool_schemas(
+    def _get_runtime_tool_names(
         self,
         config: RunConfig,
         trace: Trace,
         *,
         in_side_branch: bool = False,
-    ) -> List[Dict]:
-        """按 Trace 持久化模式和当前协议状态过滤工具 Schema。"""
-        schemas = deepcopy(self._get_tool_schemas(
+    ) -> set[str]:
+        """在基础能力上应用 Recursive 协议状态门禁。"""
+        tool_names = self._get_configured_tool_names(
             config.tools,
             config.tool_groups,
             config.exclude_tools,
-        ))
+        )
         policy = policy_from_context(trace.context)
         protocol_tools = {"submit_task_report", "review_task_result"}
 
         if not policy.requires_task_protocol or in_side_branch:
-            schemas = [
-                schema for schema in schemas
-                if schema.get("function", {}).get("name") not in protocol_tools
-            ]
-        else:
-            state = ensure_task_protocol(trace.context)
-            if state["pending_reviews"]:
-                schemas = [
-                    schema for schema in schemas
-                    if schema.get("function", {}).get("name") == "review_task_result"
-                ]
-            elif state["next_actions"]:
-                schemas = [
-                    schema for schema in schemas
-                    if schema.get("function", {}).get("name") == "agent"
-                ]
-            else:
-                is_child = bool(trace.parent_trace_id)
-                has_report = state.get("task_report") is not None
-                schemas = [
-                    schema for schema in schemas
-                    if not (
-                        schema.get("function", {}).get("name") == "review_task_result"
-                        or (
-                            schema.get("function", {}).get("name") == "submit_task_report"
-                            and (not is_child or has_report)
-                        )
-                    )
-                ]
+            return tool_names - protocol_tools
+
+        state = ensure_task_protocol(trace.context)
+        if state["pending_reviews"]:
+            return tool_names & {"review_task_result"}
+        if state["next_actions"]:
+            return tool_names & {"agent"}
+
+        tool_names.discard("review_task_result")
+        if not trace.parent_trace_id or state.get("task_report") is not None:
+            tool_names.discard("submit_task_report")
+        return tool_names
+
+    def _get_runtime_tool_schemas(
+        self,
+        config: RunConfig,
+        trace: Trace,
+        *,
+        in_side_branch: bool = False,
+        runtime_tool_names: Optional[set[str]] = None,
+    ) -> List[Dict]:
+        """按 Trace 持久化模式和当前协议状态过滤工具 Schema。"""
+        tool_names = (
+            runtime_tool_names
+            if runtime_tool_names is not None
+            else self._get_runtime_tool_names(
+                config,
+                trace,
+                in_side_branch=in_side_branch,
+            )
+        )
+        schemas = deepcopy(self.tools.get_schemas(list(tool_names)))
+        policy = policy_from_context(trace.context)
 
         for schema in schemas:
             function = schema.get("function", {})
@@ -3309,6 +3348,48 @@ class AgentRunner:
                     required.append("task")
         return schemas
 
+    def _build_tool_context(
+        self,
+        *,
+        config: RunConfig,
+        trace: Trace,
+        trace_id: str,
+        goal_id: Optional[str],
+        goal_tree: Optional[GoalTree],
+        sequence: int,
+        side_branch_ctx: Optional[SideBranchContext],
+        trigger_event: Optional[str],
+    ) -> Dict[str, Any]:
+        """构建工具上下文;Recursive 内部字段最后覆盖,不能由调用方伪造。"""
+        context = {
+            "store": self.trace_store,
+            "trace_id": trace_id,
+            "goal_id": goal_id,
+            "runner": self,
+            "goal_tree": goal_tree,
+            "knowledge_config": config.knowledge,
+            "sequence": sequence,
+            "side_branch": {
+                "type": side_branch_ctx.type,
+                "branch_id": side_branch_ctx.branch_id,
+                "is_side_branch": True,
+                "max_turns": side_branch_ctx.max_turns,
+                "trigger_event": trigger_event,
+            } if side_branch_ctx else None,
+            **(config.context or {}),
+        }
+        if policy_from_context(trace.context).mode is AgentMode.RECURSIVE:
+            context.update({
+                RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY: sorted(
+                    self._get_configured_tool_names(
+                        config.tools,
+                        config.tool_groups,
+                        config.exclude_tools,
+                    )
+                ),
+            })
+        return context
+
     # 默认 system prompt 前缀(当 config.system_prompt 和前端都未提供 system message 时使用)
     # 注意:此常量已迁移到 cyber_agent.core.prompts,这里保留引用以保持向后兼容
 

+ 25 - 4
cyber_agent/tools/builtin/subagent.py

@@ -19,6 +19,7 @@ from cyber_agent.tools import tool
 from cyber_agent.core.agent_mode import (
     AgentMode,
     AgentPolicy,
+    RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY,
     apply_policy_to_context,
     assert_removed_config_absent,
     policy_from_context,
@@ -322,15 +323,26 @@ def _get_allowed_tools(
     """按子 Agent 深度生成唯一的工具权限列表。"""
     runner = context.get("runner")
     if runner and hasattr(runner, "tools") and hasattr(runner.tools, "get_tool_names"):
+        registered_tools = set(runner.tools.get_tool_names())
+        if policy.mode is AgentMode.RECURSIVE:
+            parent_capabilities = context.get(
+                RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY
+            )
+            if not isinstance(parent_capabilities, list) or any(
+                not isinstance(name, str) for name in parent_capabilities
+            ):
+                return None
+            allowed_tools = registered_tools & set(parent_capabilities)
+        else:
+            # Legacy 保留原行为:从全局 Registry 取工具全集。
+            allowed_tools = registered_tools
+
         blocked_tools = {"evaluate", "bash_command"}
         if not policy.requires_task_protocol:
             blocked_tools.update({"submit_task_report", "review_task_result"})
         if policy.mode is AgentMode.LEGACY or agent_depth >= policy.max_depth:
             blocked_tools.add("agent")
-        return [
-            name for name in runner.tools.get_tool_names()
-            if name not in blocked_tools
-        ]
+        return sorted(allowed_tools - blocked_tools)
     return None
 
 
@@ -598,6 +610,15 @@ async def _run_agents(
     except ValueError as exc:
         return {"status": "failed", "error": str(exc)}
 
+    if (
+        policy.mode is AgentMode.RECURSIVE
+        and _get_allowed_tools(context, 0, policy) is None
+    ):
+        return {
+            "status": "failed",
+            "error": "Recursive agent capability context is missing or invalid",
+        }
+
     protocol_state = None
     approved_action = None
     if policy.requires_task_protocol:

+ 7 - 1
cyber_agent/tools/registry.py

@@ -228,7 +228,8 @@ class ToolRegistry:
 		uid: str = "",
 		context: Optional[Dict[str, Any]] = None,
 		sensitive_data: Optional[Dict[str, Any]] = None,
-		inject_values: Optional[Dict[str, Any]] = None
+		inject_values: Optional[Dict[str, Any]] = None,
+		allowed_tool_names: Optional[set[str]] = None,
 	) -> str:
 		"""
 		执行工具调用
@@ -243,6 +244,11 @@ class ToolRegistry:
 		Returns:
 			JSON 字符串格式的结果
 		"""
+		if allowed_tool_names is not None and name not in allowed_tool_names:
+			error_msg = f"Tool not allowed in current run: {name}"
+			logger.warning(f"[ToolRegistry] {error_msg}")
+			return json.dumps({"error": error_msg}, ensure_ascii=False)
+
 		if name not in self._tools:
 			error_msg = f"Unknown tool: {name}"
 			logger.error(f"[ToolRegistry] {error_msg}")

+ 197 - 0
tests/test_recursive_tool_capabilities.py

@@ -0,0 +1,197 @@
+import json
+import os
+import tempfile
+import unittest
+from unittest.mock import patch
+
+from cyber_agent.core.agent_mode import RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY
+from cyber_agent.core.runner import AgentRunner, RunConfig
+from cyber_agent.tools.builtin.subagent import agent
+from cyber_agent.tools.registry import ToolRegistry
+from cyber_agent.trace.goal_models import GoalTree
+from cyber_agent.trace.models import Trace
+from cyber_agent.trace.store import FileSystemTraceStore
+
+
+class FakeToolRegistry:
+    TOOL_NAMES = [
+        "agent",
+        "evaluate",
+        "bash_command",
+        "read_file",
+        "write_file",
+        "goal",
+    ]
+
+    def get_tool_names(self):
+        return list(self.TOOL_NAMES)
+
+
+class RecordingRunner:
+    def __init__(self, store):
+        self.store = store
+        self.tools = FakeToolRegistry()
+        self.debug = False
+        self.stdin_check = None
+        self._cancel_events = {}
+        self.configs = []
+
+    async def run_result(self, messages, config, on_event=None):
+        self.configs.append(config)
+        await self.store.update_trace(config.trace_id, status="completed")
+        return {
+            "status": "completed",
+            "summary": "done",
+            "saved_knowledge_ids": [],
+            "stats": {"total_messages": 1, "total_tokens": 1, "total_cost": 0.0},
+        }
+
+
+class RecursiveCapabilityInheritanceTest(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.store = FileSystemTraceStore(self.temp_dir.name)
+        self.runner = RecordingRunner(self.store)
+        self.root_id = "capability-root"
+        await self.store.create_trace(Trace(
+            trace_id=self.root_id,
+            mode="agent",
+            task="root",
+            uid="user-1",
+            model="fake",
+            context={
+                "agent_mode": "recursive",
+                "agent_mode_revision": 1,
+                "agent_depth": 0,
+                "root_trace_id": self.root_id,
+            },
+        ))
+        await self.store.update_goal_tree(self.root_id, GoalTree(mission="root"))
+
+    async def asyncTearDown(self):
+        self.temp_dir.cleanup()
+
+    def context(self, trace_id, capabilities=None):
+        context = {
+            "store": self.store,
+            "trace_id": trace_id,
+            "runner": self.runner,
+        }
+        if capabilities is not None:
+            context[RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY] = capabilities
+        return context
+
+    async def test_child_and_grandchild_only_inherit_parent_capabilities(self):
+        parent_capabilities = ["agent", "read_file"]
+        child_result = await agent(
+            task="child",
+            context=self.context(self.root_id, parent_capabilities),
+        )
+        child_config = self.runner.configs[-1]
+        self.assertEqual({"agent", "read_file"}, set(child_config.tools))
+        self.assertNotIn("write_file", child_config.tools)
+
+        grandchild_result = await agent(
+            task="grandchild",
+            context=self.context(child_result["sub_trace_id"], child_config.tools),
+        )
+        grandchild_config = self.runner.configs[-1]
+        self.assertEqual({"agent", "read_file"}, set(grandchild_config.tools))
+        self.assertNotIn("write_file", grandchild_config.tools)
+        self.assertEqual(
+            child_result["sub_trace_id"],
+            (await self.store.get_trace(grandchild_result["sub_trace_id"])).parent_trace_id,
+        )
+
+    async def test_missing_recursive_capability_snapshot_fails_before_spawn(self):
+        result = await agent(task="child", context=self.context(self.root_id))
+        self.assertEqual("failed", result["status"])
+        self.assertIn("capability context", result["error"])
+        children = await self.store.list_traces(parent_trace_id=self.root_id, limit=10)
+        self.assertEqual([], children)
+
+
+class RecursiveRuntimeDispatchTest(unittest.IsolatedAsyncioTestCase):
+    async def test_registry_rejects_before_counting_or_calling(self):
+        registry = ToolRegistry()
+        calls = 0
+
+        async def dangerous():
+            nonlocal calls
+            calls += 1
+            return "called"
+
+        registry.register(dangerous)
+        result = json.loads(await registry.execute(
+            "dangerous",
+            {},
+            allowed_tool_names=set(),
+        ))
+        self.assertIn("not allowed", result["error"])
+        self.assertEqual(0, calls)
+        self.assertEqual(
+            0,
+            registry.get_stats("dangerous")["dangerous"]["call_count"],
+        )
+
+    async def test_forged_tool_call_is_blocked_in_serial_and_parallel_dispatch(self):
+        for parallel in (False, True):
+            with self.subTest(parallel=parallel):
+                temp_dir = tempfile.TemporaryDirectory()
+                self.addCleanup(temp_dir.cleanup)
+                store = FileSystemTraceStore(temp_dir.name)
+                registry = ToolRegistry()
+                dangerous_calls = 0
+
+                async def safe():
+                    return "safe"
+
+                async def dangerous():
+                    nonlocal dangerous_calls
+                    dangerous_calls += 1
+                    return "dangerous"
+
+                registry.register(safe, groups=["test"])
+                registry.register(dangerous, groups=["test"])
+                llm_calls = 0
+
+                async def fake_llm(**kwargs):
+                    nonlocal llm_calls
+                    llm_calls += 1
+                    if llm_calls == 1:
+                        return {
+                            "content": "",
+                            "tool_calls": [{
+                                "id": "forged-call",
+                                "type": "function",
+                                "function": {"name": "dangerous", "arguments": "{}"},
+                            }],
+                            "finish_reason": "tool_calls",
+                        }
+                    return {"content": "done", "finish_reason": "stop"}
+
+                runner = AgentRunner(
+                    trace_store=store,
+                    tool_registry=registry,
+                    llm_call=fake_llm,
+                )
+                with patch.dict(os.environ, {"AGENT_MODE": "recursive"}, clear=False):
+                    result = await runner.run_result(
+                        messages=[{"role": "user", "content": "test"}],
+                        config=RunConfig(
+                            tools=["safe"],
+                            tool_groups=[],
+                            parallel_tool_execution=parallel,
+                            enable_research_flow=False,
+                        ),
+                    )
+                self.assertEqual("completed", result["status"])
+                self.assertEqual(0, dangerous_calls)
+                self.assertEqual(
+                    0,
+                    registry.get_stats("dangerous")["dangerous"]["call_count"],
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()

+ 2 - 0
tests/test_subagent_recursion.py

@@ -3,6 +3,7 @@ import tempfile
 import unittest
 from unittest.mock import AsyncMock, patch
 
+from cyber_agent.core.agent_mode import RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY
 from cyber_agent.tools.builtin.subagent import agent, evaluate
 from cyber_agent.trace.goal_models import Goal, GoalTree
 from cyber_agent.trace.models import Trace
@@ -94,6 +95,7 @@ class SubAgentModeTest(unittest.IsolatedAsyncioTestCase):
             "goal_id": goal_id,
             "runner": self.runner,
             "knowledge_config": None,
+            RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY: self.runner.tools.get_tool_names(),
         }
 
     async def _spawn_one(self, parent_id, task="child"):

+ 2 - 0
tests/test_task_protocol.py

@@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, patch
 
 from pydantic import ValidationError
 
+from cyber_agent.core.agent_mode import RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY
 from cyber_agent.core.runner import AgentRunner, RunConfig
 from cyber_agent.core.task_protocol import (
     TaskBrief,
@@ -169,6 +170,7 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
             "runner": self.runner,
             "goal_tree": tree or self.tree,
             "sequence": 10,
+            RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY: self.runner.tools.get_tool_names(),
         }
 
     async def children(self):