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, initialize_task_progress, 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.builtin.task_protocol import update_task_progress from cyber_agent.tools.registry import ToolRegistry from cyber_agent.trace.models import Trace from cyber_agent.trace.store import FileSystemTraceStore from cyber_agent.core.validation import ( ValidationPolicy, ValidatorSettings, persist_validation_policy, ) 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) state = ensure_task_protocol(context) if state.get("task_progress_head_revision") is None: initialize_task_progress( state, root_task_anchor_hash=context["root_task_anchor_hash"], ) 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 "recursive_validation_protocol" in str(message.get("content", "")) for message in messages ) def has_tool_result(messages, name): return any( message.get("role") == "tool" and message.get("name") == name for message in messages ) def ready_progress_response(call_id="progress-ready"): return { "content": "", "tool_calls": [{ "id": call_id, "type": "function", "function": { "name": "update_task_progress", "arguments": json.dumps({ "expected_revision": 1, "progress": { "phase": "ready_to_submit", "questions": [], "blockers": [], "findings": [], "hypotheses": [], "work_items": [], "decision_rationale": "The fixture completed the checked work.", }, }), }, }], "finish_reason": "tool_calls", } def validator_response(messages, outcome="passed"): packet = next( json.loads(message["content"]) for message in messages if message.get("role") == "user" and "recursive_validation_protocol" in str(message.get("content", "")) ) scope = packet["validation_scope"] passed = outcome == "passed" checks = [ { "check_id": item["check_id"], "status": outcome, "evidence_refs": [], "issue": None if passed else "the answer lacks required support", } for item in packet["validation_plan"]["checks"] if item["scope"] == scope ] return { "content": json.dumps({ "outcome": outcome, "scope": scope, "checks": checks, "reason": "the persisted record was checked", "retry_from": None if passed else ( "task_definition" if scope in {"task", "root"} else scope ), }), "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": ["update_task_progress"], "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"]) if not has_tool_result(kwargs["messages"], "update_task_progress"): return ready_progress_response() 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( {"validator_web_search", "validator_open_url"}, tool_names(validators[0].tools), ) usage = await ResourceBudgetController(self.store).get_usage(root.trace_id) self.assertEqual(1, usage.total_agents) self.assertEqual(3, usage.llm_calls) self.assertEqual(14, usage.total_tokens) self.assertEqual(3, 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") if not has_tool_result(kwargs["messages"], "update_task_progress"): return ready_progress_response() 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), ) if not has_tool_result(kwargs["messages"], "update_task_progress"): return ready_progress_response() 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"]) registry.register( update_task_progress, hidden_params=["context"], groups=["core"], ) ordinary_calls = 0 validator_packet = None async def llm_call(**kwargs): nonlocal ordinary_calls, validator_packet if is_validator_call(kwargs["messages"]): validator_packet = next( json.loads(message["content"]) for message in kwargs["messages"] if message.get("role") == "user" and "recursive_validation_protocol" in str(message.get("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", } if not has_tool_result(kwargs["messages"], "update_task_progress"): return ready_progress_response() 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", "update_task_progress"]), ) 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, "recursive_revision_2_read_only"): 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", } if not has_tool_result(kwargs["messages"], "update_task_progress"): return ready_progress_response() 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", "update_task_progress"], 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": 3, "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": 3, "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": 3, "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 ( "update_task_progress" in names and not has_tool_result(messages, "update_task_progress") and ("# Task Brief" in user_text or root_reviewed) ): return ready_progress_response( "child-progress" if "# Task Brief" in user_text else "root-progress" ) 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", "update_task_progress", "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) aggregate = child_state["task_report_validation"]["aggregate_result"] self.assertEqual( child_validators[0].trace_id, aggregate["scope_results"][0]["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(9, 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(), }), )) root = await self.store.get_trace(root_id) persist_validation_policy( root.context, ValidationPolicy(), ValidatorSettings(), ) await self.store.update_trace(root_id, context=root.context) 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()