Agent/tests/test_recursive_child_scheduling.py

import asyncio
import tempfile
import unittest

from cyber_agent.core.agent_mode import (
    RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY,
    RECURSIVE_CHILD_EXECUTION_MODE_CONTEXT_KEY,
    RECURSIVE_MAX_PARALLEL_CHILDREN_CONTEXT_KEY,
)
from cyber_agent.core.runner import AgentRunner
from cyber_agent.tools.builtin.subagent import agent
from cyber_agent.trace.goal_models import GoalTree
from cyber_agent.trace.models import Trace
from cyber_agent.trace.store import FileSystemTraceStore


class SchedulingToolRegistry:
    def get_tool_names(self, current_url=None, groups=None):
        return ["agent", "read_file"]


class CountingRunner(AgentRunner):
    def __init__(self, store, delays=None):
        super().__init__(
            trace_store=store,
            tool_registry=SchedulingToolRegistry(),
            llm_call=lambda **kwargs: None,
        )
        self.delays = delays or {}
        self.active = 0
        self.max_active = 0
        self.start_order = []
        self.finish_order = []
        self.configs = []

    async def run_result(self, messages, config, on_event=None):
        self.configs.append(config)
        self.start_order.append(config.name)
        self.active += 1
        self.max_active = max(self.max_active, self.active)
        try:
            await asyncio.sleep(self.delays.get(config.name, 0.01))
            self.finish_order.append(config.name)
            await self.trace_store.update_trace(config.trace_id, status="completed")
            return {
                "status": "completed",
                "summary": f"result:{config.name}",
                "saved_knowledge_ids": [],
                "stats": {"total_messages": 1, "total_tokens": 1, "total_cost": 0.0},
            }
        finally:
            self.active -= 1
            event = self._cancel_events.get(config.trace_id)
            self.release_recursive_trace(config.trace_id, event)


class RecursiveChildSchedulingTest(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 create_root(self, trace_id, mode="recursive"):
        await self.store.create_trace(Trace(
            trace_id=trace_id,
            mode="agent",
            task=trace_id,
            uid="user-1",
            model="fake",
            context={
                "agent_mode": mode,
                "agent_mode_revision": 1,
                "agent_depth": 0,
                "root_trace_id": trace_id,
            },
        ))
        await self.store.update_goal_tree(trace_id, GoalTree(mission=trace_id))

    def context(self, runner, trace_id, mode=None, maximum=None):
        context = {
            "store": self.store,
            "trace_id": trace_id,
            "runner": runner,
            RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY: ["agent", "read_file"],
        }
        if mode is not None:
            context[RECURSIVE_CHILD_EXECUTION_MODE_CONTEXT_KEY] = mode
        if maximum is not None:
            context[RECURSIVE_MAX_PARALLEL_CHILDREN_CONTEXT_KEY] = maximum
        return context

    async def test_recursive_defaults_to_strict_sequential(self):
        await self.create_root("sequential-root")
        runner = CountingRunner(self.store)
        result = await agent(
            task=["task-1", "task-2", "task-3"],
            context=self.context(runner, "sequential-root"),
        )
        self.assertEqual("completed", result["status"])
        self.assertEqual(1, runner.max_active)
        self.assertEqual(["task-1", "task-2", "task-3"], runner.start_order)

    async def test_parallel_mode_never_exceeds_two_and_keeps_input_order(self):
        await self.create_root("parallel-root")
        delays = {f"task-{index}": (7 - index) * 0.005 for index in range(1, 7)}
        runner = CountingRunner(self.store, delays)
        tasks = [f"task-{index}" for index in range(1, 7)]
        result = await agent(
            task=tasks,
            context=self.context(runner, "parallel-root", "parallel", 2),
        )
        self.assertEqual("completed", result["status"])
        self.assertEqual(2, runner.max_active)
        self.assertNotEqual(tasks, runner.finish_order)
        positions = [result["summary"].index(f"result:{task}") for task in tasks]
        self.assertEqual(sorted(positions), positions)
        self.assertTrue(all(config.child_execution_mode == "parallel" for config in runner.configs))
        self.assertTrue(all(config.max_parallel_children == 2 for config in runner.configs))

    async def test_invalid_recursive_settings_fail_before_trace_creation(self):
        invalid_values = [
            ("invalid", 2),
            ("parallel", 0),
            ("parallel", -1),
            ("parallel", True),
            ("parallel", 3),
        ]
        for index, (mode, maximum) in enumerate(invalid_values):
            root_id = f"invalid-{index}"
            await self.create_root(root_id)
            runner = CountingRunner(self.store)
            result = await agent(
                task=["one", "two"],
                context=self.context(runner, root_id, mode, maximum),
            )
            self.assertEqual("failed", result["status"])
            children = await self.store.list_traces(parent_trace_id=root_id, limit=10)
            self.assertEqual([], children)

    async def test_legacy_ignores_new_settings_and_keeps_batch_parallel(self):
        await self.create_root("legacy-root", mode="legacy")
        runner = CountingRunner(self.store)
        result = await agent(
            task=["one", "two", "three"],
            context=self.context(runner, "legacy-root", "invalid", 99),
        )
        self.assertEqual("completed", result["status"])
        self.assertEqual(3, runner.max_active)


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