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.core.task_protocol import ensure_task_protocol from cyber_agent.tools.builtin.knowledge import KnowledgeConfig from cyber_agent.tools.builtin.subagent import agent from cyber_agent.tools.builtin.task_protocol import update_task_progress 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_recursive_framework_context_fields_cannot_be_overridden(self): with tempfile.TemporaryDirectory() as directory: store = FileSystemTraceStore(directory) runner = AgentRunner(trace_store=store, llm_call=lambda **kwargs: None) trace = Trace( trace_id="real-trace", mode="agent", task="test", context={"agent_mode": "recursive", "agent_mode_revision": 2}, ) goal_tree = GoalTree(mission="test") config = RunConfig(context={ "store": "spoof-store", "trace_id": "spoof-trace", "goal_id": "spoof-goal", "runner": "spoof-runner", "goal_tree": "spoof-tree", "sequence": -1, RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY: ["spoof-tool"], }) context = runner._build_tool_context( config=config, trace=trace, trace_id=trace.trace_id, goal_id="real-goal", goal_tree=goal_tree, sequence=7, side_branch_ctx=None, trigger_event=None, ) self.assertIs(store, context["store"]) self.assertEqual("real-trace", context["trace_id"]) self.assertEqual("real-goal", context["goal_id"]) self.assertIs(runner, context["runner"]) self.assertIs(goal_tree, context["goal_tree"]) self.assertEqual(7, context["sequence"]) self.assertNotEqual( ["spoof-tool"], context[RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY], ) async def test_pending_review_runtime_allowlist_rejects_forged_agent(self): registry = ToolRegistry() agent_calls = 0 async def agent(): nonlocal agent_calls agent_calls += 1 return "called" async def review_task_result(): return "reviewed" registry.register(agent) registry.register(review_task_result) runner = AgentRunner(tool_registry=registry, llm_call=lambda **kwargs: None) trace = Trace( trace_id="pending-root", mode="agent", task="pending", context={"agent_mode": "recursive", "agent_mode_revision": 2}, ) state = ensure_task_protocol(trace.context) state["pending_reviews"]["child"] = {"task_report": {}} config = RunConfig( tools=["agent", "review_task_result"], tool_groups=[], ) allowed = runner._get_runtime_tool_names(config, trace) self.assertEqual({"review_task_result"}, allowed) result = json.loads(await registry.execute( "agent", {}, allowed_tool_names=allowed, )) self.assertIn("not allowed", result["error"]) self.assertEqual(0, agent_calls) self.assertEqual(0, registry.get_stats("agent")["agent"]["call_count"]) 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"]) registry.register( update_task_progress, hidden_params=["context"], groups=["core"], ) llm_calls = 0 async def fake_llm(**kwargs): nonlocal llm_calls llm_calls += 1 messages = kwargs.get("messages") or [] if messages and "recursive_validation_protocol" in messages[0]["content"]: packet = json.loads(messages[1]["content"]) scope = packet["validation_scope"] return { "content": json.dumps({ "outcome": "passed", "scope": scope, "checks": [ { "check_id": item["check_id"], "status": "passed", "evidence_refs": [], "issue": None, } for item in packet["validation_plan"]["checks"] if item["scope"] == scope ], "reason": "unauthorized tools were not executed", "retry_from": None, }), "tool_calls": [], "finish_reason": "stop", } if llm_calls == 1: return { "content": "", "tool_calls": [{ "id": "forged-call", "type": "function", "function": {"name": "dangerous", "arguments": "{}"}, }], "finish_reason": "tool_calls", } if not any( message.get("role") == "tool" and message.get("name") == "update_task_progress" for message in messages ): return { "content": "", "tool_calls": [{ "id": "progress-ready", "type": "function", "function": { "name": "update_task_progress", "arguments": json.dumps({ "expected_revision": 1, "progress": { "phase": "ready_to_submit", "questions": [], "blockers": [], "findings": [], "hypotheses": [], "work_items": [], "decision_rationale": ( "The forged tool was rejected." ), }, }), }, }], "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", "update_task_progress"], tool_groups=[], parallel_tool_execution=parallel, enable_research_flow=False, root_task_anchor={ "objective": "reject forged tool calls", "completion_criteria": [ "forged tools are never executed" ], }, knowledge=KnowledgeConfig( enable_extraction=False, enable_completion_extraction=False, enable_injection=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()