Agent/tests/test_recursive_tool_capabilities.py

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