Agent/tests/test_subagent_recursion.py

import asyncio
import os
import tempfile
import unittest
from unittest.mock import AsyncMock, patch

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 SubAgentRecursionTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        self.store = FileSystemTraceStore(self.temp_dir.name)
        self.runner = FakeRunner(self.store)
        self.previous_recursion_setting = os.environ.get("AGENT_RECURSION_ENABLED")
        os.environ["AGENT_RECURSION_ENABLED"] = "true"
        self.root_id = await self._create_root("root")

    async def asyncTearDown(self):
        if self.previous_recursion_setting is None:
            os.environ.pop("AGENT_RECURSION_ENABLED", None)
        else:
            os.environ["AGENT_RECURSION_ENABLED"] = self.previous_recursion_setting
        self.temp_dir.cleanup()

    async def _create_root(self, trace_id, uid="user-1"):
        trace = Trace(
            trace_id=trace_id,
            mode="agent",
            task=f"task-{trace_id}",
            uid=uid,
            model="fake-model",
            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,
        }

    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_recursion_disabled_preserves_single_level_behavior(self):
        os.environ["AGENT_RECURSION_ENABLED"] = "false"
        first = await self._spawn_one(self.root_id)
        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("recursion is disabled", blocked["error"])

    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_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_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_legacy_trace_depth_is_recovered_from_parent_links(self):
        legacy_child_id = "legacy-child"
        await self.store.create_trace(
            Trace(
                trace_id=legacy_child_id,
                mode="agent",
                task="legacy child",
                parent_trace_id=self.root_id,
                uid="user-1",
                model="fake-model",
                context={"created_by_tool": "agent"},
            )
        )
        await self.store.update_goal_tree(
            legacy_child_id,
            GoalTree(mission="legacy child"),
        )

        result = await self._spawn_one(legacy_child_id, "grandchild")
        self.assertEqual("completed", result["status"])
        grandchild = await self.store.get_trace(result["sub_trace_id"])
        self.assertEqual(2, grandchild.context["agent_depth"])
        self.assertEqual(self.root_id, grandchild.context["root_trace_id"])

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