import asyncio import json import os import tempfile import unittest from unittest.mock import AsyncMock, patch from cyber_agent.core.runner import AgentRunner, RunConfig from cyber_agent.core.task_protocol import ensure_task_protocol from cyber_agent.tools import get_tool_registry from cyber_agent.tools.builtin.subagent import _load_or_create_task_report from cyber_agent.trace import run_api from cyber_agent.trace.models import Trace from cyber_agent.trace.store import FileSystemTraceStore def brief(index): return { "objective": f"child-{index}", "reason": "exercise subtree cancellation", "completion_criteria": [f"finish child-{index}"], "expected_outputs": [f"result for child-{index}"], } class RecursiveSubtreeCancellationTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.temp_dir = tempfile.TemporaryDirectory() self.store = FileSystemTraceStore(self.temp_dir.name) async def asyncTearDown(self): self.temp_dir.cleanup() async def test_stop_child_only_cancels_its_descendants(self): runner = AgentRunner( trace_store=self.store, llm_call=AsyncMock(), ) root_event = asyncio.Event() runner._cancel_events["root"] = root_event runner._recursive_active_traces["root"] = root_event child_event = runner.register_recursive_child("root", "child") grandchild_event = runner.register_recursive_child("child", "grandchild") sibling_event = runner.register_recursive_child("root", "sibling") self.assertTrue(await runner.stop("child")) self.assertTrue(child_event.is_set()) self.assertTrue(grandchild_event.is_set()) self.assertFalse(root_event.is_set()) self.assertFalse(sibling_event.is_set()) async def test_legacy_stop_keeps_single_trace_behavior(self): runner = AgentRunner(trace_store=self.store, llm_call=AsyncMock()) root_event = asyncio.Event() child_event = asyncio.Event() runner._cancel_events.update({"legacy-root": root_event, "legacy-child": child_event}) self.assertTrue(await runner.stop("legacy-root")) self.assertTrue(root_event.is_set()) self.assertFalse(child_event.is_set()) async def test_legacy_stop_during_llm_keeps_preexisting_completion_behavior(self): llm_started = asyncio.Event() release_llm = asyncio.Event() late_tool_calls = 0 async def legacy_late_tool(): nonlocal late_tool_calls late_tool_calls += 1 return "called" registry = get_tool_registry() registry.register(legacy_late_tool, groups=["legacy-cancel-test"]) async def fake_llm(**kwargs): llm_started.set() await release_llm.wait() return { "content": "legacy completed", "tool_calls": [{ "id": "legacy-late-call", "type": "function", "function": {"name": "legacy_late_tool", "arguments": "{}"}, }], "finish_reason": "tool_calls", } runner = AgentRunner( trace_store=self.store, tool_registry=registry, llm_call=fake_llm, ) with patch.dict(os.environ, {"AGENT_MODE": "legacy"}, clear=False): task = asyncio.create_task(runner.run_result( messages=[{"role": "user", "content": "legacy"}], config=RunConfig(enable_research_flow=False), )) await asyncio.wait_for(llm_started.wait(), timeout=1) root = (await self.store.list_traces(limit=10))[0] self.assertTrue(await runner.stop(root.trace_id)) release_llm.set() result = await asyncio.wait_for(task, timeout=2) self.assertEqual("stopped", result["status"]) self.assertEqual("legacy completed", result["summary"]) self.assertEqual(1, late_tool_calls) self.assertEqual("stopped", (await self.store.get_trace(root.trace_id)).status) async def test_stop_during_llm_discards_returned_tool_call(self): dangerous_calls = 0 llm_started = asyncio.Event() release_llm = asyncio.Event() registry = get_tool_registry() async def fake_llm(**kwargs): llm_started.set() await release_llm.wait() return { "content": "late response", "tool_calls": [{ "id": "late-tool", "type": "function", "function": {"name": "dangerous_cancel_test", "arguments": "{}"}, }], "finish_reason": "tool_calls", } async def dangerous_cancel_test(): nonlocal dangerous_calls dangerous_calls += 1 return "called" registry.register(dangerous_cancel_test, groups=["cancel-test"]) runner = AgentRunner( trace_store=self.store, tool_registry=registry, llm_call=fake_llm, ) with patch.dict(os.environ, {"AGENT_MODE": "recursive"}, clear=False): task = asyncio.create_task(runner.run_result( messages=[{"role": "user", "content": "cancel while llm runs"}], config=RunConfig( tools=["dangerous_cancel_test"], tool_groups=[], enable_research_flow=False, root_completion_criteria=["stopping prevents late tool execution"], ), )) await asyncio.wait_for(llm_started.wait(), timeout=1) traces = await self.store.list_traces(limit=10) root_id = traces[0].trace_id self.assertTrue(await runner.stop(root_id)) release_llm.set() result = await asyncio.wait_for(task, timeout=2) self.assertEqual("stopped", result["status"]) self.assertEqual("Agent execution stopped.", result["summary"]) self.assertEqual(0, dangerous_calls) self.assertEqual("stopped", (await self.store.get_trace(root_id)).status) async def test_root_stop_cancels_running_and_queued_children_with_reports(self): child_llm_started = 0 two_children_started = asyncio.Event() release_children = asyncio.Event() async def fake_llm(**kwargs): nonlocal child_llm_started messages = kwargs["messages"] user_text = "\n".join( str(message.get("content", "")) for message in messages if message.get("role") == "user" ) if "root-cancel-batch" in user_text: return { "content": "", "tool_calls": [{ "id": "delegate-batch", "type": "function", "function": { "name": "agent", "arguments": json.dumps({ "task_brief": [brief(index) for index in range(4)], }), }, }], "finish_reason": "tool_calls", } child_llm_started += 1 if child_llm_started == 2: two_children_started.set() await release_children.wait() return {"content": "late child result", "finish_reason": "stop"} runner = AgentRunner( trace_store=self.store, tool_registry=get_tool_registry(), llm_call=fake_llm, ) with patch.dict(os.environ, {"AGENT_MODE": "recursive"}, clear=False): task = asyncio.create_task(runner.run_result( messages=[{"role": "user", "content": "root-cancel-batch"}], config=RunConfig( tools=["agent"], tool_groups=[], child_execution_mode="parallel", max_parallel_children=2, enable_research_flow=False, root_completion_criteria=["stopping propagates through the child tree"], ), )) await asyncio.wait_for(two_children_started.wait(), timeout=2) traces = await self.store.list_traces(limit=10) root = next(trace for trace in traces if trace.parent_trace_id is None) self.assertTrue(await runner.stop(root.trace_id)) release_children.set() result = await asyncio.wait_for(task, timeout=3) traces = await self.store.list_traces(limit=10) children = [trace for trace in traces if trace.parent_trace_id == root.trace_id] self.assertEqual(4, len(children)) self.assertEqual(2, child_llm_started) self.assertTrue(all(child.status == "stopped" for child in children)) self.assertEqual(2, sum(child.total_messages == 0 for child in children)) self.assertEqual("stopped", result["status"]) persisted_root = await self.store.get_trace(root.trace_id) self.assertEqual("stopped", persisted_root.status) pending = ensure_task_protocol(persisted_root.context)["pending_reviews"] self.assertEqual(sorted(child.trace_id for child in children), list(pending)) self.assertTrue(all( entry["task_report"]["outcome"] == "failed" for entry in pending.values() )) self.assertEqual({}, runner._cancel_events) self.assertEqual({}, runner._recursive_active_traces) self.assertEqual({}, runner._active_children) self.assertEqual({}, runner._active_parents) async def test_stopped_report_replaces_satisfied_report_only_once(self): child_id = "stale-report-child" context = { "agent_mode": "recursive", "agent_mode_revision": 2, "agent_depth": 1, "root_trace_id": "root", } state = ensure_task_protocol(context) state["task_report"] = { "child_trace_id": child_id, "summary": "previously satisfied", "outcome": "satisfied", "validation": {"hard_passed": True, "open_issues": []}, "next_step_suggestion": { "direction": "NONE", "reason": "nothing else", }, "outputs": [], "evidence": [], "remaining_issues": [], } await self.store.create_trace(Trace( trace_id=child_id, mode="agent", task="child", parent_trace_id="root", status="stopped", context=context, )) for _ in range(2): report = await _load_or_create_task_report( self.store, child_id, "stopped by parent", "stopped", 10, ) self.assertEqual("failed", report.outcome) persisted = await self.store.get_trace(child_id) self.assertEqual( 1, len(ensure_task_protocol(persisted.context)["report_history"]), ) async def test_stop_api_uses_the_runner_that_started_the_root(self): trace_id = "example-root" actual_runner = AsyncMock() actual_runner.stop.return_value = True run_api._running_runners[trace_id] = actual_runner try: response = await run_api.stop_trace(trace_id) finally: run_api._running_runners.pop(trace_id, None) run_api._running_tasks.pop(trace_id, None) actual_runner.stop.assert_awaited_once_with(trace_id) self.assertEqual("stopping", response.status) async def test_immediate_resume_waits_for_old_event_cleanup(self): trace_id = "resume-root" runner = AgentRunner(trace_store=self.store, llm_call=AsyncMock()) old_event = asyncio.Event() old_event.set() runner._cancel_events[trace_id] = old_event runner._recursive_active_traces[trace_id] = old_event started = asyncio.Event() async def old_run(): started.set() try: await asyncio.Event().wait() finally: runner.release_recursive_trace(trace_id, old_event) if run_api._running_runners.get(trace_id) is runner: run_api._running_runners.pop(trace_id, None) old_task = asyncio.create_task(old_run()) run_api._running_tasks[trace_id] = old_task run_api._running_runners[trace_id] = runner await started.wait() await run_api._cancel_and_wait_for_running_task(trace_id) new_event = runner._cancel_events.setdefault(trace_id, asyncio.Event()) self.assertIsNot(old_event, new_event) self.assertFalse(new_event.is_set()) self.assertNotIn(trace_id, run_api._running_tasks) self.assertNotIn(trace_id, run_api._running_runners) runner._cancel_events.pop(trace_id, None) if __name__ == "__main__": unittest.main()