Agent/tests/test_recursive_runtime_extensions.py

import asyncio
import json
import os
import tempfile
import unittest
from unittest.mock import patch

from cyber_agent.core.agent_mode import RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY
from cyber_agent.core.resource_budget import (
    RESOURCE_BUDGET_CONTEXT_KEY,
    ResourceBudget,
    ResourceBudgetController,
)
from cyber_agent.core.runner import AgentRunner, RunConfig
from cyber_agent.core.task_protocol import ensure_task_protocol, new_task_protocol
from cyber_agent.core.context_policy import persist_root_task_anchor
from cyber_agent.tools import get_tool_registry
from cyber_agent.tools.builtin.knowledge import KnowledgeConfig
from cyber_agent.tools.builtin.subagent import agent
from cyber_agent.tools.registry import ToolRegistry
from cyber_agent.trace.models import Trace
from cyber_agent.trace.store import FileSystemTraceStore


ROOT_CRITERIA = ["return a checked final answer"]
ROOT_ANCHOR = {
    "objective": "return a checked final answer",
    "completion_criteria": ROOT_CRITERIA,
    "constraints": ["do not invent evidence"],
}
CHILD_BRIEF = {
    "objective": "check one concrete fact",
    "reason": "the root needs independently checked evidence",
    "completion_criteria": ["return one supported conclusion"],
    "expected_outputs": ["one checked conclusion"],
    "constraints": ["do not guess"],
}


def knowledge_disabled():
    return KnowledgeConfig(
        enable_extraction=False,
        enable_completion_extraction=False,
        enable_injection=False,
    )


def with_root_anchor(context):
    persist_root_task_anchor(context, ROOT_ANCHOR)
    return context


def recursive_env(**overrides):
    values = {
        "AGENT_MODE": "recursive",
        "AGENT_RESOURCE_BUDGET_ENABLED": "true",
        "AGENT_MAX_TOTAL_AGENTS": "50",
        "AGENT_MAX_LLM_CALLS": "150",
        "AGENT_MAX_TOTAL_TOKENS": "1500000",
        "AGENT_MAX_TOTAL_COST_USD": "15",
        "AGENT_MAX_DURATION_SECONDS": "3600",
        "AGENT_RESERVED_FINAL_CALLS": "1",
    }
    values.update(overrides)
    return patch.dict(os.environ, values, clear=False)


def tool_names(schemas):
    return {
        schema["function"]["name"]
        for schema in schemas or []
    }


def is_validator_call(messages):
    return any(
        message.get("role") == "system"
        and "independent validator" in str(message.get("content", ""))
        for message in messages
    )


def validator_response(messages, outcome="passed"):
    packet = json.loads(messages[-1]["content"])
    scope = packet["validation_scope"]
    passed = outcome == "passed"
    return {
        "content": json.dumps({
            "outcome": outcome,
            "scope": scope,
            "reason": "the persisted record was checked",
            "issues": [] if passed else ["the answer lacks required support"],
            "retry_from": None if passed else "evidence",
        }),
        "tool_calls": [],
        "finish_reason": "stop",
        "prompt_tokens": 3,
        "completion_tokens": 2,
        "cost": 0.001,
    }


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

    def config(self, **overrides):
        values = {
            "tools": [],
            "tool_groups": [],
            "enable_research_flow": False,
            "root_task_anchor": ROOT_ANCHOR,
            "knowledge": knowledge_disabled(),
        }
        values.update(overrides)
        return RunConfig(**values)

    async def test_root_completes_only_after_independent_validator_passes(self):
        calls = []

        async def llm_call(**kwargs):
            calls.append(kwargs)
            if is_validator_call(kwargs["messages"]):
                return validator_response(kwargs["messages"])
            return {
                "content": "checked final answer",
                "tool_calls": [],
                "finish_reason": "stop",
                "prompt_tokens": 5,
                "completion_tokens": 4,
                "cost": 0.002,
            }

        runner = AgentRunner(trace_store=self.store, llm_call=llm_call)
        with recursive_env():
            result = await runner.run_result(
                messages=[{"role": "user", "content": "answer carefully"}],
                config=self.config(),
            )

        self.assertEqual("completed", result["status"])
        root = await self.store.get_trace(result["trace_id"])
        state = ensure_task_protocol(root.context)
        self.assertTrue(state["root_validation_passed"])
        self.assertEqual(1, state["root_validation_attempts"])
        validators = await self.store.list_traces(
            parent_trace_id=root.trace_id,
            created_by_tool="validator",
            limit=10,
        )
        self.assertEqual(1, len(validators))
        self.assertEqual(root.trace_id, validators[0].context["root_trace_id"])
        self.assertEqual([], validators[0].tools)
        usage = await ResourceBudgetController(self.store).get_usage(root.trace_id)
        self.assertEqual(1, usage.total_agents)
        self.assertEqual(2, usage.llm_calls)
        self.assertEqual(14, usage.total_tokens)
        self.assertEqual(2, len(calls))

    async def test_root_gets_one_revision_then_fails_on_second_rejection(self):
        ordinary_calls = 0
        validation_calls = 0

        async def llm_call(**kwargs):
            nonlocal ordinary_calls, validation_calls
            if is_validator_call(kwargs["messages"]):
                validation_calls += 1
                return validator_response(kwargs["messages"], outcome="failed")
            ordinary_calls += 1
            return {
                "content": f"candidate-{ordinary_calls}",
                "tool_calls": [],
                "finish_reason": "stop",
            }

        runner = AgentRunner(trace_store=self.store, llm_call=llm_call)
        with recursive_env():
            result = await runner.run_result(
                messages=[{"role": "user", "content": "answer carefully"}],
                config=self.config(),
            )

        self.assertEqual("failed", result["status"])
        self.assertEqual(2, ordinary_calls)
        self.assertEqual(2, validation_calls)
        root = await self.store.get_trace(result["trace_id"])
        state = ensure_task_protocol(root.context)
        self.assertEqual(2, state["root_validation_attempts"])
        self.assertEqual(
            ["failed", "failed"],
            [item["outcome"] for item in state["root_validation_history"]],
        )
        self.assertFalse(state["root_validation_passed"])

    async def test_root_can_complete_after_one_validator_requested_revision(self):
        ordinary_calls = 0
        validation_outcomes = iter(["failed", "passed"])
        second_candidate_messages = None

        async def llm_call(**kwargs):
            nonlocal ordinary_calls, second_candidate_messages
            if is_validator_call(kwargs["messages"]):
                return validator_response(
                    kwargs["messages"],
                    outcome=next(validation_outcomes),
                )
            ordinary_calls += 1
            if ordinary_calls == 2:
                second_candidate_messages = kwargs["messages"]
            return {
                "content": f"candidate-{ordinary_calls}",
                "tool_calls": [],
                "finish_reason": "stop",
            }

        runner = AgentRunner(trace_store=self.store, llm_call=llm_call)
        with recursive_env():
            result = await runner.run_result(
                messages=[{"role": "user", "content": "revise if needed"}],
                config=self.config(),
            )

        self.assertEqual("completed", result["status"])
        root = await self.store.get_trace(result["trace_id"])
        state = ensure_task_protocol(root.context)
        self.assertEqual(2, state["root_validation_attempts"])
        self.assertEqual(
            ["failed", "passed"],
            [item["outcome"] for item in state["root_validation_history"]],
        )
        self.assertTrue(any(
            message.get("role") == "user"
            and "framework-owned ValidationResult" in str(message.get("content"))
            for message in second_candidate_messages
        ))
        persisted_messages = await self.store.get_trace_messages(root.trace_id)
        self.assertTrue(any(
            message.role == "user"
            and "framework-owned ValidationResult" in str(message.content)
            for message in persisted_messages
        ))

        calls_after_completion = 0

        async def must_not_call(**_kwargs):
            nonlocal calls_after_completion
            calls_after_completion += 1
            raise AssertionError("a third validation attempt must not start")

        resumed = AgentRunner(trace_store=self.store, llm_call=must_not_call)
        with recursive_env():
            with self.assertRaisesRegex(ValueError, "two independent"):
                await resumed.run_result(
                    [],
                    self.config(trace_id=root.trace_id),
                )
        self.assertEqual(0, calls_after_completion)

    async def test_root_validator_reads_current_tool_result_and_candidate(self):
        registry = ToolRegistry()

        async def evidence_tool():
            return {"verified": "persisted-tool-result"}

        registry.register(evidence_tool, groups=["test"])
        ordinary_calls = 0
        validator_packet = None

        async def llm_call(**kwargs):
            nonlocal ordinary_calls, validator_packet
            if is_validator_call(kwargs["messages"]):
                validator_packet = json.loads(kwargs["messages"][-1]["content"])
                return validator_response(kwargs["messages"])
            ordinary_calls += 1
            if ordinary_calls == 1:
                return {
                    "content": "",
                    "tool_calls": [{
                        "id": "evidence-call",
                        "type": "function",
                        "function": {
                            "name": "evidence_tool",
                            "arguments": "{}",
                        },
                    }],
                    "finish_reason": "tool_calls",
                }
            return {
                "content": "candidate-from-real-tool-result",
                "tool_calls": [],
                "finish_reason": "stop",
            }

        runner = AgentRunner(
            trace_store=self.store,
            tool_registry=registry,
            llm_call=llm_call,
        )
        with recursive_env():
            result = await runner.run_result(
                [{"role": "user", "content": "use the evidence tool"}],
                self.config(tools=["evidence_tool"]),
            )

        self.assertEqual("completed", result["status"])
        trajectory = validator_packet["trajectory"]
        self.assertTrue(any(
            item.get("role") == "tool"
            and "persisted-tool-result" in str(item.get("content"))
            for item in trajectory
        ))
        self.assertTrue(any(
            item.get("role") == "assistant"
            and "candidate-from-real-tool-result" in str(item.get("content"))
            for item in trajectory
        ))

    async def test_new_recursive_root_requires_explicit_task_anchor(self):
        async def llm_call(**_kwargs):
            raise AssertionError("LLM must not be called")

        runner = AgentRunner(trace_store=self.store, llm_call=llm_call)
        with recursive_env():
            with self.assertRaisesRegex(ValueError, "root_task_anchor"):
                await runner.run_result(
                    messages=[{"role": "user", "content": "missing criteria"}],
                    config=self.config(root_task_anchor=None),
                )
        self.assertEqual([], await self.store.list_traces(limit=10))

    async def test_old_recursive_revision_two_trace_is_not_silently_migrated(self):
        trace_id = "old-recursive-experiment"
        await self.store.create_trace(Trace(
            trace_id=trace_id,
            mode="agent",
            context={
                "agent_mode": "recursive",
                "agent_mode_revision": 2,
                "agent_depth": 0,
                "root_trace_id": trace_id,
                "task_protocol": {},
            },
        ))
        runner = AgentRunner(trace_store=self.store, llm_call=lambda **_: None)
        with self.assertRaisesRegex(ValueError, "create a new trace"):
            await runner.run_result([], RunConfig(trace_id=trace_id))

    async def test_validator_trace_cannot_be_continued_as_an_agent(self):
        validator_id = "root@validator-fixture"
        await self.store.create_trace(Trace(
            trace_id=validator_id,
            mode="agent",
            agent_type="validator",
            tools=[],
            status="completed",
            context=with_root_anchor({
                "created_by_tool": "validator",
                "agent_mode": "recursive",
                "agent_mode_revision": 2,
                "root_trace_id": "root",
            }),
        ))
        calls = 0

        async def llm_call(**_kwargs):
            nonlocal calls
            calls += 1
            return {"content": "must not run", "tool_calls": []}

        runner = AgentRunner(trace_store=self.store, llm_call=llm_call)
        with self.assertRaisesRegex(ValueError, "cannot be continued"):
            await runner.run_result(
                [{"role": "user", "content": "gain tools"}],
                RunConfig(trace_id=validator_id),
            )
        self.assertEqual(0, calls)
        validator = await self.store.get_trace(validator_id)
        self.assertEqual("completed", validator.status)
        self.assertEqual([], validator.tools)


class RecursiveBudgetIntegrationTest(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_response_over_token_budget_never_executes_its_tool_call(self):
        dangerous_calls = 0
        registry = ToolRegistry()

        async def dangerous():
            nonlocal dangerous_calls
            dangerous_calls += 1
            return "must not run"

        registry.register(dangerous, groups=["test"])

        async def llm_call(**_kwargs):
            return {
                "content": "over-budget response",
                "tool_calls": [{
                    "id": "dangerous-call",
                    "type": "function",
                    "function": {"name": "dangerous", "arguments": "{}"},
                }],
                "finish_reason": "tool_calls",
                "prompt_tokens": 8,
                "completion_tokens": 8,
            }

        runner = AgentRunner(
            trace_store=self.store,
            tool_registry=registry,
            llm_call=llm_call,
        )
        config = RunConfig(
            tools=["dangerous"],
            tool_groups=[],
            enable_research_flow=False,
            root_task_anchor=ROOT_ANCHOR,
            knowledge=knowledge_disabled(),
        )
        with recursive_env(AGENT_MAX_TOTAL_TOKENS="10"):
            result = await runner.run_result(
                [{"role": "user", "content": "do not overspend"}],
                config,
            )

        self.assertEqual("failed", result["status"])
        self.assertEqual(0, dangerous_calls)
        root = await self.store.get_trace(result["trace_id"])
        self.assertEqual("budget_exhausted:tokens", root.context["termination_reason"])
        usage = await ResourceBudgetController(self.store).get_usage(root.trace_id)
        self.assertEqual(16, usage.total_tokens)
        self.assertEqual("budget_exhausted:tokens", usage.exhausted_reason)
        assistant = [
            message
            for message in await self.store.get_trace_messages(root.trace_id)
            if message.role == "assistant"
        ]
        self.assertIsNone(assistant[-1].content["tool_calls"])

    async def test_agent_budget_denial_blocks_child_but_preserves_root_validator(self):
        ordinary_calls = 0

        async def llm_call(**kwargs):
            nonlocal ordinary_calls
            if is_validator_call(kwargs["messages"]):
                return validator_response(kwargs["messages"])
            ordinary_calls += 1
            if ordinary_calls == 1:
                return {
                    "content": "",
                    "tool_calls": [{
                        "id": "delegate",
                        "type": "function",
                        "function": {
                            "name": "agent",
                            "arguments": json.dumps({"task_brief": CHILD_BRIEF}),
                        },
                    }],
                    "finish_reason": "tool_calls",
                }
            return {
                "content": "root can still finish without a child",
                "tool_calls": [],
                "finish_reason": "stop",
            }

        runner = AgentRunner(
            trace_store=self.store,
            tool_registry=get_tool_registry(),
            llm_call=llm_call,
        )
        config = RunConfig(
            tools=["agent"],
            tool_groups=[],
            enable_research_flow=False,
            root_task_anchor=ROOT_ANCHOR,
            knowledge=knowledge_disabled(),
        )
        with recursive_env(AGENT_MAX_TOTAL_AGENTS="1"):
            result = await runner.run_result(
                [{"role": "user", "content": "try one delegated check"}],
                config,
            )

        self.assertEqual("completed", result["status"])
        root = await self.store.get_trace(result["trace_id"])
        children = await self.store.list_traces(
            parent_trace_id=root.trace_id,
            created_by_tool="agent",
            limit=10,
        )
        self.assertEqual([], children)
        usage = await ResourceBudgetController(self.store).get_usage(root.trace_id)
        self.assertEqual(1, usage.total_agents)
        self.assertEqual("budget_exhausted:agents", usage.exhausted_reason)
        self.assertTrue(ensure_task_protocol(root.context)["root_validation_passed"])

    async def test_persisted_child_is_not_refunded_when_goal_write_fails(self):
        root_id = "partial-child-create-root"
        budget = ResourceBudget()
        await self.store.create_trace(Trace(
            trace_id=root_id,
            mode="agent",
            uid="user-1",
            model="fake",
            context=with_root_anchor({
                "agent_mode": "recursive",
                "agent_mode_revision": 2,
                "agent_depth": 0,
                "root_trace_id": root_id,
                RESOURCE_BUDGET_CONTEXT_KEY: budget.to_dict(),
                "task_protocol": new_task_protocol(),
            }),
        ))
        await ResourceBudgetController(self.store).initialize(root_id, budget)
        original_update_goal_tree = self.store.update_goal_tree

        async def fail_child_goal_write(trace_id, goal_tree):
            if trace_id != root_id:
                raise RuntimeError("simulated child GoalTree write failure")
            return await original_update_goal_tree(trace_id, goal_tree)

        self.store.update_goal_tree = fail_child_goal_write
        runner = AgentRunner(
            trace_store=self.store,
            tool_registry=get_tool_registry(),
            llm_call=lambda **_: None,
        )

        with self.assertRaisesRegex(RuntimeError, "GoalTree write failure"):
            await agent(
                task_brief=CHILD_BRIEF,
                context={
                    "store": self.store,
                    "trace_id": root_id,
                    "runner": runner,
                    RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY: ["agent"],
                },
            )

        children = await self.store.list_traces(
            parent_trace_id=root_id,
            created_by_tool="agent",
            limit=10,
        )
        self.assertEqual(1, len(children))
        self.assertEqual("failed", children[0].status)
        self.assertIn("Child initialization failed", children[0].error_message)
        usage = await ResourceBudgetController(self.store).get_usage(root_id)
        self.assertEqual(2, usage.total_agents)

    async def test_parent_goal_projection_failure_closes_precreated_child(self):
        root_id = "parent-goal-projection-root"
        budget = ResourceBudget()
        await self.store.create_trace(Trace(
            trace_id=root_id,
            mode="agent",
            uid="user-1",
            model="fake",
            context=with_root_anchor({
                "agent_mode": "recursive",
                "agent_mode_revision": 2,
                "agent_depth": 0,
                "root_trace_id": root_id,
                RESOURCE_BUDGET_CONTEXT_KEY: budget.to_dict(),
                "task_protocol": new_task_protocol(),
            }),
        ))
        await ResourceBudgetController(self.store).initialize(root_id, budget)
        runner = AgentRunner(
            trace_store=self.store,
            tool_registry=get_tool_registry(),
            llm_call=lambda **_: None,
        )

        async def fail_parent_projection(*_args, **_kwargs):
            raise RuntimeError("simulated parent Goal projection failure")

        with patch(
            "cyber_agent.tools.builtin.subagent._update_goal_start",
            new=fail_parent_projection,
        ):
            with self.assertRaisesRegex(RuntimeError, "projection failure"):
                await agent(
                    task_brief=CHILD_BRIEF,
                    context={
                        "store": self.store,
                        "trace_id": root_id,
                        "runner": runner,
                        RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY: ["agent"],
                    },
                )

        children = await self.store.list_traces(
            parent_trace_id=root_id,
            created_by_tool="agent",
            limit=10,
        )
        self.assertEqual(1, len(children))
        self.assertEqual("failed", children[0].status)
        usage = await ResourceBudgetController(self.store).get_usage(root_id)
        self.assertEqual(2, usage.total_agents)

    async def test_approved_action_write_failure_closes_precreated_child(self):
        root_id = "approved-action-write-root"
        budget = ResourceBudget()
        context = with_root_anchor({
            "agent_mode": "recursive",
            "agent_mode_revision": 2,
            "agent_depth": 0,
            "root_trace_id": root_id,
            RESOURCE_BUDGET_CONTEXT_KEY: budget.to_dict(),
            "task_protocol": new_task_protocol(),
        })
        state = ensure_task_protocol(context)
        state["next_actions"] = [{
            "decision": "ACCEPT_REFINE",
            "child_trace_id": "reviewed-child",
            "goal_id": None,
            "task_brief": CHILD_BRIEF,
            "created_at_sequence": 1,
        }]
        await self.store.create_trace(Trace(
            trace_id=root_id,
            mode="agent",
            uid="user-1",
            model="fake",
            context=context,
        ))
        await ResourceBudgetController(self.store).initialize(root_id, budget)
        runner = AgentRunner(
            trace_store=self.store,
            tool_registry=get_tool_registry(),
            llm_call=lambda **_: None,
        )
        original_update_trace = self.store.update_trace
        rejected_parent_write = False

        async def fail_approved_action_write(trace_id, **updates):
            nonlocal rejected_parent_write
            if trace_id == root_id and "context" in updates and not rejected_parent_write:
                rejected_parent_write = True
                raise RuntimeError("simulated approved action write failure")
            return await original_update_trace(trace_id, **updates)

        self.store.update_trace = fail_approved_action_write
        with self.assertRaisesRegex(RuntimeError, "approved action write failure"):
            await agent(
                task_brief=CHILD_BRIEF,
                context={
                    "store": self.store,
                    "trace_id": root_id,
                    "runner": runner,
                    RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY: ["agent"],
                },
            )

        children = await self.store.list_traces(
            parent_trace_id=root_id,
            created_by_tool="agent",
            limit=10,
        )
        self.assertEqual(1, len(children))
        self.assertEqual("failed", children[0].status)
        usage = await ResourceBudgetController(self.store).get_usage(root_id)
        self.assertEqual(2, usage.total_agents)


class RecursiveLifecycleIntegrationTest(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_child_report_is_validated_before_parent_review(self):
        child_submitted = False
        root_delegated = False
        root_reviewed = False

        async def llm_call(**kwargs):
            nonlocal child_submitted, root_delegated, root_reviewed
            messages = kwargs["messages"]
            if is_validator_call(messages):
                return validator_response(messages)
            names = tool_names(kwargs.get("tools"))
            user_text = "\n".join(
                str(message.get("content", ""))
                for message in messages
                if message.get("role") == "user"
            )
            if "submit_task_report" in names and not child_submitted:
                child_submitted = True
                return {
                    "content": "",
                    "tool_calls": [{
                        "id": "submit-report",
                        "type": "function",
                        "function": {
                            "name": "submit_task_report",
                            "arguments": json.dumps({
                                "task_report": {
                                    "summary": "one fact was checked",
                                    "outcome": "satisfied",
                                    "validation": {
                                        "hard_passed": True,
                                        "open_issues": [],
                                    },
                                    "next_step_suggestion": {
                                        "direction": "NONE",
                                        "reason": "the delegated check is complete",
                                    },
                                    "outputs": [{"kind": "answer", "value": "checked"}],
                                    "evidence": [{"kind": "tool", "value": "persisted"}],
                                    "remaining_issues": [],
                                }
                            }),
                        },
                    }],
                    "finish_reason": "tool_calls",
                }
            if "# Task Brief" in user_text:
                return {
                    "content": "structured child report submitted",
                    "tool_calls": [],
                    "finish_reason": "stop",
                }
            if names == {"review_task_result"}:
                roots = [
                    trace
                    for trace in await self.store.list_traces(limit=20)
                    if trace.parent_trace_id is None
                ]
                root = roots[0]
                pending = ensure_task_protocol(root.context)["pending_reviews"]
                child_id = next(iter(pending))
                root_reviewed = True
                return {
                    "content": "",
                    "tool_calls": [{
                        "id": "review-report",
                        "type": "function",
                        "function": {
                            "name": "review_task_result",
                            "arguments": json.dumps({
                                "child_trace_id": child_id,
                                "decision": "ASCEND",
                                "reason": "the validated child satisfies this goal",
                            }),
                        },
                    }],
                    "finish_reason": "tool_calls",
                }
            if "agent" in names and not root_delegated:
                root_delegated = True
                return {
                    "content": "",
                    "tool_calls": [{
                        "id": "delegate-child",
                        "type": "function",
                        "function": {
                            "name": "agent",
                            "arguments": json.dumps({"task_brief": CHILD_BRIEF}),
                        },
                    }],
                    "finish_reason": "tool_calls",
                }
            return {
                "content": "final answer based on the validated child",
                "tool_calls": [],
                "finish_reason": "stop",
            }

        runner = AgentRunner(
            trace_store=self.store,
            tool_registry=get_tool_registry(),
            llm_call=llm_call,
        )
        config = RunConfig(
            tools=["agent", "submit_task_report", "review_task_result"],
            tool_groups=[],
            enable_research_flow=False,
            root_task_anchor=ROOT_ANCHOR,
            knowledge=knowledge_disabled(),
        )
        with recursive_env():
            result = await runner.run_result(
                [{"role": "user", "content": "perform one delegated check"}],
                config,
            )

        self.assertEqual("completed", result["status"])
        self.assertTrue(root_reviewed)
        root = await self.store.get_trace(result["trace_id"])
        business_children = await self.store.list_traces(
            parent_trace_id=root.trace_id,
            created_by_tool="agent",
            limit=10,
        )
        self.assertEqual(1, len(business_children))
        child = business_children[0]
        child_validators = await self.store.list_traces(
            parent_trace_id=child.trace_id,
            created_by_tool="validator",
            limit=10,
        )
        root_validators = await self.store.list_traces(
            parent_trace_id=root.trace_id,
            created_by_tool="validator",
            limit=10,
        )
        self.assertEqual(1, len(child_validators))
        self.assertEqual(1, len(root_validators))
        self.assertEqual(
            child.trace_id,
            child_validators[0].context["evaluated_trace_id"],
        )
        child_state = ensure_task_protocol(child.context)
        self.assertEqual(
            child_validators[0].trace_id,
            child_state["task_report_validation"]["validation_result"][
                "validator_trace_id"
            ],
        )
        root_state = ensure_task_protocol(root.context)
        self.assertEqual({}, root_state["pending_reviews"])
        self.assertEqual(child.trace_id, root_state["reviews"][0]["child_trace_id"])
        usage = await ResourceBudgetController(self.store).get_usage(root.trace_id)
        self.assertEqual(2, usage.total_agents)
        self.assertEqual(7, usage.llm_calls)

    async def test_parent_stop_reaches_validator_after_child_runner_finished(self):
        root_id = "validator-cancel-root"
        child_id = "validator-cancel-child"
        budget = ResourceBudget()
        await self.store.create_trace(Trace(
            trace_id=root_id,
            mode="agent",
            model="fake-model",
            context=with_root_anchor({
                "agent_mode": "recursive",
                "agent_mode_revision": 2,
                "agent_depth": 0,
                "root_trace_id": root_id,
                RESOURCE_BUDGET_CONTEXT_KEY: budget.to_dict(),
                "task_protocol": new_task_protocol(),
            }),
        ))
        await self.store.create_trace(Trace(
            trace_id=child_id,
            mode="agent",
            model="fake-model",
            parent_trace_id=root_id,
            context=with_root_anchor({
                "agent_mode": "recursive",
                "agent_mode_revision": 2,
                "agent_depth": 1,
                "root_trace_id": root_id,
                "created_by_tool": "agent",
                "task_protocol": new_task_protocol(CHILD_BRIEF),
            }),
        ))
        await ResourceBudgetController(self.store).initialize(
            root_id,
            budget,
            initial_agents=2,
        )
        started = asyncio.Event()
        release = asyncio.Event()

        async def llm_call(**kwargs):
            started.set()
            await release.wait()
            return validator_response(kwargs["messages"])

        runner = AgentRunner(trace_store=self.store, llm_call=llm_call)
        root_event = asyncio.Event()
        runner._cancel_events[root_id] = root_event
        runner._recursive_active_traces[root_id] = root_event
        validation_task = asyncio.create_task(runner.validate_recursive_trace(
            child_id,
            scope="task",
            task_brief=CHILD_BRIEF,
            task_report={"summary": "candidate"},
        ))
        await asyncio.wait_for(started.wait(), timeout=1)

        self.assertIn(child_id, runner._active_children[root_id])
        self.assertTrue(await runner.stop(root_id))
        release.set()
        validation_run = await asyncio.wait_for(validation_task, timeout=1)

        self.assertEqual("error", validation_run.result.outcome)
        validator = await self.store.get_trace(validation_run.trace_id)
        self.assertEqual("failed", validator.status)
        self.assertEqual(child_id, validator.parent_trace_id)
        self.assertEqual(root_id, validator.context["root_trace_id"])
        self.assertNotIn(child_id, runner._recursive_active_traces)
        self.assertNotIn(validation_run.trace_id, runner._recursive_active_traces)


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