Agent/tests/test_subagent_recursion.py

import asyncio
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
from cyber_agent.trace.store import FileSystemTraceStore


class FakeToolRegistry:
    TOOL_NAMES = [
        "agent",
        "evaluate",
        "bash_command",
        "read_file",
        "grep_content",
        "glob_files",
        "goal",
    ]

    def get_tool_names(self):
        return list(self.TOOL_NAMES)


class FakeRunner:
    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": f"completed: {config.trace_id}",
            "saved_knowledge_ids": [],
            "stats": {
                "total_messages": 1,
                "total_tokens": 1,
                "total_cost": 0.0,
            },
        }


class SubAgentModeTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        self.store = FileSystemTraceStore(self.temp_dir.name)
        self.runner = FakeRunner(self.store)
        self.root_id = await self._create_root("root")

    async def asyncTearDown(self):
        self.temp_dir.cleanup()

    async def _create_root(
        self,
        trace_id,
        uid="user-1",
        *,
        agent_mode="recursive",
        include_mode=True,
    ):
        context = {
            "agent_depth": 0,
            "root_trace_id": trace_id,
        }
        if include_mode:
            context.update({
                "agent_mode": agent_mode,
                "agent_mode_revision": 1,
            })
        trace = Trace(
            trace_id=trace_id,
            mode="agent",
            task=f"task-{trace_id}",
            uid=uid,
            model="fake-model",
            context=context,
        )
        await self.store.create_trace(trace)
        await self.store.update_goal_tree(trace_id, GoalTree(mission=trace.task))
        return trace_id

    def _context(self, trace_id, goal_id=None):
        return {
            "store": self.store,
            "trace_id": trace_id,
            "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"):
        return await agent(task=task, context=self._context(parent_id))

    async def _local_children(self, parent_id):
        return await self.store.list_traces(
            parent_trace_id=parent_id,
            created_by_tool="agent",
            limit=20,
        )

    async def test_legacy_preserves_single_level_behavior(self):
        legacy_root = await self._create_root("legacy-root", agent_mode="legacy")
        first = await self._spawn_one(legacy_root)
        self.assertEqual("completed", first["status"])

        child_id = first["sub_trace_id"]
        child_config = self.runner.configs[-1]
        self.assertNotIn("agent", child_config.tools)

        blocked = await self._spawn_one(child_id, "grandchild")
        self.assertEqual("failed", blocked["status"])
        self.assertIn("depth limit reached", blocked["error"])
        self.assertIn("mode=legacy", blocked["error"])

    async def test_legacy_has_no_direct_child_quota(self):
        legacy_root = await self._create_root("legacy-unlimited", agent_mode="legacy")
        result = await agent(
            task=[f"legacy-child-{index}" for index in range(8)],
            context=self._context(legacy_root),
        )
        self.assertEqual("completed", result["status"])
        self.assertEqual(8, len(await self._local_children(legacy_root)))

    async def test_five_subagent_levels_and_sixth_level_rejected(self):
        parent_id = self.root_id
        for expected_depth in range(1, 6):
            result = await self._spawn_one(parent_id, f"depth-{expected_depth}")
            self.assertEqual("completed", result["status"])
            child_id = result["sub_trace_id"]
            child = await self.store.get_trace(child_id)
            self.assertEqual(parent_id, child.parent_trace_id)
            self.assertEqual(expected_depth, child.context["agent_depth"])
            self.assertEqual(self.root_id, child.context["root_trace_id"])

            tools = self.runner.configs[-1].tools
            if expected_depth < 5:
                self.assertIn("agent", tools)
            else:
                self.assertNotIn("agent", tools)
            self.assertNotIn("evaluate", tools)
            self.assertNotIn("bash_command", tools)
            parent_id = child_id

        blocked = await self._spawn_one(parent_id, "depth-6")
        self.assertEqual("failed", blocked["status"])
        self.assertIn("depth limit reached", blocked["error"])
        self.assertEqual([], await self._local_children(parent_id))

    async def test_six_children_allowed_and_seventh_rejected(self):
        tasks = [f"child-{index}" for index in range(6)]
        result = await agent(task=tasks, context=self._context(self.root_id))
        self.assertEqual("completed", result["status"])
        self.assertEqual(6, len(await self._local_children(self.root_id)))

        blocked = await self._spawn_one(self.root_id, "child-7")
        self.assertEqual("failed", blocked["status"])
        self.assertIn("child Agent limit exceeded", blocked["error"])
        self.assertEqual(6, len(await self._local_children(self.root_id)))

    async def test_child_batch_is_rejected_without_partial_creation(self):
        initial = [f"child-{index}" for index in range(5)]
        await agent(task=initial, context=self._context(self.root_id))

        blocked = await agent(
            task=["overflow-1", "overflow-2"],
            context=self._context(self.root_id),
        )
        self.assertEqual("failed", blocked["status"])
        self.assertEqual(5, len(await self._local_children(self.root_id)))

    async def test_mismatched_per_agent_messages_are_rejected_before_spawn(self):
        result = await agent(
            task=["first", "second"],
            messages=[[{"role": "user", "content": "only one message list"}]],
            context=self._context(self.root_id),
        )
        self.assertEqual("failed", result["status"])
        self.assertIn("exactly one message list per task", result["error"])
        self.assertEqual([], await self._local_children(self.root_id))

    async def test_goal_keeps_children_created_across_multiple_calls(self):
        tree = await self.store.get_goal_tree(self.root_id)
        tree.goals = [Goal(id="1", description="delegate twice")]
        await self.store.update_goal_tree(self.root_id, tree)
        context = self._context(self.root_id, goal_id="1")

        await agent(task="first", context=context)
        await agent(task="second", context=context)

        updated_tree = await self.store.get_goal_tree(self.root_id)
        sub_trace_ids = updated_tree.find("1").sub_trace_ids
        self.assertEqual(2, len(sub_trace_ids))
        self.assertEqual(2, len({entry["trace_id"] for entry in sub_trace_ids}))

    async def test_legacy_goal_overwrites_sub_traces_with_latest_call(self):
        legacy_root = await self._create_root("legacy-goal", agent_mode="legacy")
        tree = await self.store.get_goal_tree(legacy_root)
        tree.goals = [Goal(id="1", description="delegate twice")]
        await self.store.update_goal_tree(legacy_root, tree)
        context = self._context(legacy_root, goal_id="1")

        first = await agent(task="first", context=context)
        second = await agent(task="second", context=context)

        updated_tree = await self.store.get_goal_tree(legacy_root)
        sub_trace_ids = updated_tree.find("1").sub_trace_ids
        self.assertEqual(
            [{"trace_id": second["sub_trace_id"], "mission": "second"}],
            sub_trace_ids,
        )
        self.assertNotEqual(first["sub_trace_id"], second["sub_trace_id"])

    async def test_concurrent_batches_never_exceed_limit(self):
        results = await asyncio.gather(
            agent(
                task=[f"left-{index}" for index in range(4)],
                context=self._context(self.root_id),
            ),
            agent(
                task=[f"right-{index}" for index in range(4)],
                context=self._context(self.root_id),
            ),
        )
        self.assertEqual({"completed", "failed"}, {item["status"] for item in results})
        self.assertLessEqual(len(await self._local_children(self.root_id)), 6)

    async def test_continue_from_does_not_consume_child_slot(self):
        tasks = [f"child-{index}" for index in range(6)]
        created = await agent(task=tasks, context=self._context(self.root_id))
        child_id = created["sub_trace_ids"][0]["trace_id"]

        continued = await agent(
            task="continue existing child",
            continue_from=child_id,
            context=self._context(self.root_id),
        )
        self.assertEqual("completed", continued["status"])
        self.assertTrue(continued["continue_from"])
        self.assertEqual(6, len(await self._local_children(self.root_id)))

    async def test_continue_from_requires_direct_child_and_matching_owner(self):
        created = await self._spawn_one(self.root_id)
        child_id = created["sub_trace_id"]
        other_root = await self._create_root("other-root")

        wrong_parent = await agent(
            task="continue",
            continue_from=child_id,
            context=self._context(other_root),
        )
        self.assertEqual("failed", wrong_parent["status"])
        self.assertIn("direct child", wrong_parent["error"])

        await self.store.update_trace(child_id, uid="other-user")
        wrong_owner = await agent(
            task="continue",
            continue_from=child_id,
            context=self._context(self.root_id),
        )
        self.assertEqual("failed", wrong_owner["status"])
        self.assertIn("owner", wrong_owner["error"])

    async def test_continue_from_requires_matching_persisted_mode(self):
        created = await self._spawn_one(self.root_id)
        child_id = created["sub_trace_id"]
        child = await self.store.get_trace(child_id)
        child.context["agent_mode"] = "legacy"
        await self.store.update_trace(child_id, context=child.context)

        result = await agent(
            task="continue",
            continue_from=child_id,
            context=self._context(self.root_id),
        )
        self.assertEqual("failed", result["status"])
        self.assertIn("mode does not match", result["error"])

    async def test_continue_from_rejects_cross_tool_traces(self):
        evaluated = await evaluate(messages=[], context=self._context(self.root_id))
        rejected_agent = await agent(
            task="wrong tool",
            continue_from=evaluated["sub_trace_id"],
            context=self._context(self.root_id),
        )
        self.assertEqual("failed", rejected_agent["status"])
        self.assertIn("created by the agent tool", rejected_agent["error"])

        created = await self._spawn_one(self.root_id)
        rejected_evaluator = await evaluate(
            messages=[],
            continue_from=created["sub_trace_id"],
            context=self._context(self.root_id),
        )
        self.assertEqual("failed", rejected_evaluator["status"])
        self.assertIn("created by evaluate", rejected_evaluator["error"])

    async def test_continue_from_rejects_forged_depth_or_root_lineage(self):
        created = await self._spawn_one(self.root_id)
        child_id = created["sub_trace_id"]
        child = await self.store.get_trace(child_id)
        child.context["agent_depth"] = 0
        child.context["root_trace_id"] = "wrong-root"
        await self.store.update_trace(child_id, context=child.context)

        result = await agent(
            task="continue forged child",
            continue_from=child_id,
            context=self._context(self.root_id),
        )
        self.assertEqual("failed", result["status"])
        self.assertIn("lineage does not match", result["error"])

    async def test_evaluator_never_receives_recursive_or_shell_tools(self):
        await evaluate(messages=[], context=self._context(self.root_id))
        config = self.runner.configs[-1]
        self.assertEqual([], config.tool_groups)
        self.assertEqual(
            ["agent", "evaluate", "bash_command"],
            config.exclude_tools,
        )

    async def test_evaluator_traces_do_not_consume_local_child_slots(self):
        for _ in range(2):
            result = await evaluate(messages=[], context=self._context(self.root_id))
            self.assertEqual("completed", result["status"])

        result = await agent(
            task=[f"child-{index}" for index in range(6)],
            context=self._context(self.root_id),
        )
        self.assertEqual("completed", result["status"])
        self.assertEqual(6, len(await self._local_children(self.root_id)))

    async def test_trace_without_mode_is_treated_as_legacy(self):
        historical_root = await self._create_root(
            "historical-root",
            include_mode=False,
        )
        child = await self._spawn_one(historical_root, "legacy child")
        self.assertEqual("completed", child["status"])
        child_trace = await self.store.get_trace(child["sub_trace_id"])
        self.assertEqual("legacy", child_trace.context["agent_mode"])
        self.assertNotIn("agent", self.runner.configs[-1].tools)

    async def test_remote_agent_is_not_counted_as_local_child(self):
        child = await self._spawn_one(self.root_id)
        before = len(await self._local_children(child["sub_trace_id"]))
        remote_result = {
            "mode": "remote",
            "agent_type": "remote_research",
            "status": "completed",
            "summary": "remote completed",
            "stats": {},
        }
        with patch(
            "cyber_agent.tools.builtin.subagent._run_remote_agent",
            new=AsyncMock(return_value=remote_result),
        ):
            result = await agent(
                task="remote task",
                agent_type="remote_research",
                context=self._context(child["sub_trace_id"]),
            )
        self.assertEqual("completed", result["status"])
        self.assertEqual(before, len(await self._local_children(child["sub_trace_id"])))


if __name__ == "__main__":
    unittest.main()