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", agent_mode_revision=1, include_mode=True, ): context = { "agent_depth": 0, "root_trace_id": trace_id, } if include_mode: context.update({ "agent_mode": agent_mode, "agent_mode_revision": agent_mode_revision, }) 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"]))) async def test_recursive_revision_2_remote_agent_fails_before_http(self): root_id = await self._create_root( "recursive-v2-remote", agent_mode_revision=2, ) remote = AsyncMock(return_value={"status": "completed"}) with patch( "cyber_agent.tools.builtin.subagent._run_remote_agent", new=remote, ): result = await agent( task="remote task", agent_type="remote_research", context=self._context(root_id), ) self.assertEqual("failed", result["status"]) self.assertEqual( "Writable Recursive protocol does not support remote agents; " "use a local task_brief delegation", result["error"], ) remote.assert_not_awaited() self.assertEqual([], await self._local_children(root_id)) if __name__ == "__main__": unittest.main()