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.core.task_protocol import ensure_task_protocol
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_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"])
                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()