Agent/tests/test_recursive_subtree_cancellation.py

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