import json import tempfile import unittest from unittest.mock import AsyncMock, patch from pydantic import ValidationError from cyber_agent.core.runner import AgentRunner, RunConfig from cyber_agent.core.task_protocol import ( TaskBrief, TaskReport, TaskReview, ensure_task_protocol, new_task_protocol, ) 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 ( review_task_result, submit_task_report, ) from cyber_agent.trace.goal_models import Goal, GoalTree from cyber_agent.trace.goal_tool import goal_tool from cyber_agent.trace.models import Message, Trace from cyber_agent.trace.store import FileSystemTraceStore BRIEF = { "objective": "research one concrete question", "completion_criteria": ["provide one checked conclusion"], "constraints": ["do not guess"], } def report_payload(*, outcome="satisfied"): return { "summary": "checked result", "outcome": outcome, "validation": { "hard_passed": outcome == "satisfied", "open_issues": [] if outcome == "satisfied" else ["still failing"], }, "next_step_suggestion": { "direction": "NONE", "reason": "No additional task is needed", }, "outputs": [{"kind": "note", "value": "result"}], "evidence": [{"kind": "test", "value": "passed"}], "remaining_issues": [] if outcome == "satisfied" else ["still failing"], } class FakeToolRegistry: def get_tool_names(self): return [ "agent", "evaluate", "bash_command", "goal", "read_file", "submit_task_report", "review_task_result", ] class ReportingRunner: def __init__(self, store): self.store = store self.tools = FakeToolRegistry() self.debug = False self.stdin_check = None self._cancel_events = {} self.configs = [] async def run_result(self, messages, config, on_event=None): self.configs.append(config) submitted = await submit_task_report( task_report=report_payload(), context={ "store": self.store, "trace_id": config.trace_id, "sequence": 3, }, ) if submitted["status"] != "completed": raise AssertionError(submitted) await self.store.update_trace(config.trace_id, status="completed") return { "status": "completed", "summary": "child submitted TaskReport", "saved_knowledge_ids": [], "stats": {"total_messages": 1, "total_tokens": 1, "total_cost": 0.0}, } class FailedAfterReportingRunner(ReportingRunner): async def run_result(self, messages, config, on_event=None): result = await super().run_result(messages, config, on_event) await self.store.update_trace( config.trace_id, status="failed", error_message="runner failed after report submission", ) return { **result, "status": "failed", "error": "runner failed after report submission", } class ForgedIdentityRunner(ReportingRunner): async def run_result(self, messages, config, on_event=None): trace = await self.store.get_trace(config.trace_id) state = ensure_task_protocol(trace.context) state["task_report"] = TaskReport( child_trace_id="someone-elses-trace", **report_payload(), ).model_dump() await self.store.update_trace( config.trace_id, context=trace.context, status="completed", ) return { "status": "completed", "summary": "forged report identity", "saved_knowledge_ids": [], "stats": {"total_messages": 1, "total_tokens": 1, "total_cost": 0.0}, } class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.temp_dir = tempfile.TemporaryDirectory() self.store = FileSystemTraceStore(self.temp_dir.name) self.runner = ReportingRunner(self.store) self.root_id = "recursive-root" context = { "agent_mode": "recursive", "agent_mode_revision": 2, "agent_depth": 0, "root_trace_id": self.root_id, "task_protocol": new_task_protocol(), } await self.store.create_trace(Trace( trace_id=self.root_id, mode="agent", task="root", uid="user-1", model="fake-model", context=context, )) self.tree = GoalTree( mission="root", goals=[Goal(id="1", description="delegate", status="in_progress")], current_id="1", ) await self.store.update_goal_tree(self.root_id, self.tree) async def asyncTearDown(self): self.temp_dir.cleanup() def context(self, trace_id=None, goal_id="1", tree=None): return { "store": self.store, "trace_id": trace_id or self.root_id, "goal_id": goal_id, "runner": self.runner, "goal_tree": tree or self.tree, "sequence": 10, } async def children(self): return await self.store.list_traces( parent_trace_id=self.root_id, created_by_tool="agent", limit=20, ) async def test_models_validate_and_round_trip(self): brief = TaskBrief.model_validate(BRIEF) self.assertEqual(BRIEF, brief.model_dump(exclude_none=True)) with self.assertRaises(ValidationError): TaskBrief(objective="", completion_criteria=[]) with self.assertRaises(ValidationError): TaskReview( parent_trace_id="p", child_trace_id="c", decision="ACCEPT_REFINE", reason="continue", ) invalid_suggestion = report_payload() invalid_suggestion["next_step_suggestion"] = { "direction": "NONE", "reason": "No next task should be attached", "suggested_next_task": BRIEF, } with self.assertRaises(ValidationError): TaskReport(child_trace_id="child", **invalid_suggestion) root_submission = await submit_task_report( task_report=report_payload(), context={"store": self.store, "trace_id": self.root_id}, ) self.assertEqual("failed", root_submission["status"]) self.assertIn("Root Agent", root_submission["error"]) async def test_child_report_blocks_new_child_until_parent_review(self): result = await agent(task_brief=BRIEF, context=self.context()) self.assertEqual("completed", result["status"]) child_id = result["sub_trace_id"] self.assertEqual(child_id, result["task_report"]["child_trace_id"]) parent = await self.store.get_trace(self.root_id) state = ensure_task_protocol(parent.context) self.assertEqual([child_id], list(state["pending_reviews"])) tree = await self.store.get_goal_tree(self.root_id) self.assertEqual("pending_review", tree.find("1").status) restarted_runner = AgentRunner(tool_registry=get_tool_registry()) restored = await self.store.get_trace(self.root_id) restored_names = { schema["function"]["name"] for schema in restarted_runner._get_runtime_tool_schemas( RunConfig(), restored ) } self.assertEqual({"review_task_result"}, restored_names) blocked = await agent(task_brief=BRIEF, context=self.context()) self.assertEqual("failed", blocked["status"]) self.assertIn("Review pending", blocked["error"]) async def test_qwen_json_string_containers_are_normalized(self): result = await agent( task_brief=json.dumps(BRIEF), messages=json.dumps([{ "role": "system", "content": "controlled child context", }]), context=self.context(), ) self.assertEqual("completed", result["status"]) child = await self.store.get_trace(result["sub_trace_id"]) self.assertEqual( TaskBrief.model_validate(BRIEF).model_dump(), ensure_task_protocol(child.context)["task_brief"], ) async def test_messages_reject_mixed_dimensions(self): mixed_messages = [ {"role": "system", "content": "shared context"}, [{"role": "user", "content": "per-agent context"}], ] for messages in (mixed_messages, json.dumps(mixed_messages)): with self.subTest(messages_type=type(messages).__name__): result = await agent( task_brief=BRIEF, messages=messages, context=self.context(), ) self.assertEqual("failed", result["status"]) self.assertIn("without mixed dimensions", result["error"]) async def test_no_focused_goal_still_persists_and_reviews_report(self): result = await agent( task_brief=BRIEF, context=self.context(goal_id=None), ) parent = await self.store.get_trace(self.root_id) entry = ensure_task_protocol(parent.context)["pending_reviews"][ result["sub_trace_id"] ] self.assertIsNone(entry["goal_id"]) reviewed = await review_task_result( child_trace_id=result["sub_trace_id"], decision="ASCEND", reason="Standalone delegation is complete", context=self.context(goal_id=None), ) self.assertEqual("completed", reviewed["status"]) async def test_accept_refine_approves_exactly_one_next_child(self): first = await agent(task_brief=BRIEF, context=self.context()) next_brief = { "objective": "check the independent gap", "completion_criteria": ["close the gap"], } parent = await self.store.get_trace(self.root_id) ensure_task_protocol(parent.context)["protocol_correction_attempts"] = 2 await self.store.update_trace(self.root_id, context=parent.context) reviewed = await review_task_result( child_trace_id=first["sub_trace_id"], decision="ACCEPT_REFINE", reason="The gap belongs to this parent", approved_next_task=json.dumps(next_brief), context=self.context(), ) self.assertEqual("completed", reviewed["status"]) parent = await self.store.get_trace(self.root_id) self.assertEqual( 0, ensure_task_protocol(parent.context)["protocol_correction_attempts"], ) schemas = AgentRunner( tool_registry=get_tool_registry() )._get_runtime_tool_schemas(RunConfig(), parent) self.assertEqual( {"agent"}, {schema["function"]["name"] for schema in schemas}, ) wrong = await agent(task_brief=BRIEF, context=self.context()) self.assertEqual("failed", wrong["status"]) self.assertIn("approved next task", wrong["error"]) second = await agent(task_brief=next_brief, context=self.context()) self.assertEqual("completed", second["status"]) self.assertEqual(2, len(await self.children())) parent = await self.store.get_trace(self.root_id) self.assertEqual([], ensure_task_protocol(parent.context)["next_actions"]) async def test_batch_reports_are_reviewed_one_by_one(self): briefs = [ BRIEF, { "objective": "research second question", "completion_criteria": ["provide second conclusion"], }, ] result = await agent(task_brief=briefs, context=self.context()) child_ids = [item["trace_id"] for item in result["sub_trace_ids"]] self.assertEqual(2, len(result["task_reports"])) premature = await review_task_result( child_trace_id=child_ids[0], decision="ASCEND", reason="Cannot finish while another report is pending", context=self.context(), ) self.assertEqual("failed", premature["status"]) self.assertIn("final pending report", premature["error"]) first_review = await review_task_result( child_trace_id=child_ids[0], decision="REJECT_REFINE", reason="No extra task needed for this report", context=self.context(), ) self.assertEqual(1, first_review["pending_review_count"]) tree = await self.store.get_goal_tree(self.root_id) self.assertEqual("pending_review", tree.find("1").status) final_review = await review_task_result( child_trace_id=child_ids[1], decision="ASCEND", reason="All completion criteria are satisfied", context=self.context(), ) self.assertEqual(0, final_review["pending_review_count"]) tree = await self.store.get_goal_tree(self.root_id) self.assertEqual("completed", tree.find("1").status) duplicate = await review_task_result( child_trace_id=child_ids[1], decision="ASCEND", reason="Duplicate review must fail", context=self.context(), ) self.assertEqual("failed", duplicate["status"]) self.assertIn("no pending TaskReport", duplicate["error"]) async def test_terminal_review_waits_for_queued_next_action(self): result = await agent( task_brief=[ BRIEF, { "objective": "second batch child", "completion_criteria": ["finish the second child"], }, ], context=self.context(), ) child_ids = [item["trace_id"] for item in result["sub_trace_ids"]] next_brief = { "objective": "approved follow-up", "completion_criteria": ["finish the approved follow-up"], } accepted = await review_task_result( child_trace_id=child_ids[0], decision="ACCEPT_REFINE", reason="A separate follow-up is required", approved_next_task=next_brief, context=self.context(), ) self.assertEqual("completed", accepted["status"]) blocked = await review_task_result( child_trace_id=child_ids[1], decision="ASCEND", reason="Cannot ascend before the accepted follow-up", context=self.context(), ) self.assertEqual("failed", blocked["status"]) self.assertIn("approved next action", blocked["error"]) tree = await self.store.get_goal_tree(self.root_id) self.assertEqual("pending_review", tree.find("1").status) final_non_terminal = await review_task_result( child_trace_id=child_ids[1], decision="REJECT_REFINE", reason="This report adds no other follow-up", context=self.context(), ) self.assertEqual("completed", final_non_terminal["status"]) tree = await self.store.get_goal_tree(self.root_id) self.assertEqual("in_progress", tree.find("1").status) async def test_runner_failure_overrides_stale_satisfied_report(self): self.runner = FailedAfterReportingRunner(self.store) result = await agent(task_brief=BRIEF, context=self.context()) self.assertEqual("failed", result["status"]) self.assertEqual("protocol_error", result["task_report"]["outcome"]) child = await self.store.get_trace(result["sub_trace_id"]) child_state = ensure_task_protocol(child.context) self.assertEqual(1, len(child_state["report_history"])) self.assertEqual("satisfied", child_state["report_history"][0]["outcome"]) invalid_ascend = await review_task_result( child_trace_id=result["sub_trace_id"], decision="ASCEND", reason="A failed execution cannot ascend", context=self.context(), ) self.assertEqual("failed", invalid_ascend["status"]) self.assertIn("only allow", invalid_ascend["error"]) async def test_review_rejects_wrong_owner_even_if_report_is_injected(self): first = await agent(task_brief=BRIEF, context=self.context()) child_id = first["sub_trace_id"] await self.store.update_trace(child_id, uid="another-user") rejected = await review_task_result( child_trace_id=child_id, decision="ASCEND", reason="Owner mismatch must be rejected", context=self.context(), ) self.assertEqual("failed", rejected["status"]) self.assertIn("same owner", rejected["error"]) async def test_report_identity_must_match_direct_child(self): self.runner = ForgedIdentityRunner(self.store) result = await agent(task_brief=BRIEF, context=self.context()) child_id = result["sub_trace_id"] self.assertEqual("protocol_error", result["task_report"]["outcome"]) self.assertEqual(child_id, result["task_report"]["child_trace_id"]) parent = await self.store.get_trace(self.root_id) pending_report = ensure_task_protocol(parent.context)["pending_reviews"][ child_id ]["task_report"] pending_report["child_trace_id"] = "corrupted-after-recording" await self.store.update_trace(self.root_id, context=parent.context) rejected = await review_task_result( child_trace_id=child_id, decision="FAIL", reason="Corrupted identity cannot be reviewed", context=self.context(), ) self.assertEqual("failed", rejected["status"]) self.assertIn("does not match", rejected["error"]) async def test_remote_agent_keeps_legacy_task_interface(self): remote_result = { "status": "completed", "summary": "remote result", } with patch( "cyber_agent.tools.builtin.subagent._run_remote_agent", new=AsyncMock(return_value=remote_result), ) as remote_call: result = await agent( task="remote free text", agent_type="remote_librarian", ) self.assertEqual(remote_result, result) self.assertEqual("remote free text", remote_call.await_args.kwargs["task"]) async def test_revise_child_reuses_trace_and_preserves_report_history(self): first = await agent(task_brief=BRIEF, context=self.context()) child_id = first["sub_trace_id"] reviewed = await review_task_result( child_trace_id=child_id, decision="REVISE_CHILD", reason="Please strengthen the evidence", context=self.context(), ) self.assertEqual("completed", reviewed["status"]) wrong = await agent( task_brief={ "objective": "silently replace the reviewed task", "completion_criteria": ["this must be rejected"], }, continue_from=child_id, context=self.context(), ) self.assertEqual("failed", wrong["status"]) self.assertIn("approved revision", wrong["error"]) revised = await agent( task_brief=BRIEF, continue_from=child_id, context=self.context(), ) self.assertEqual(child_id, revised["sub_trace_id"]) self.assertTrue(revised["continue_from"]) self.assertEqual(1, len(await self.children())) child = await self.store.get_trace(child_id) self.assertEqual(1, len(ensure_task_protocol(child.context)["report_history"])) self.assertEqual( TaskBrief.model_validate(BRIEF).model_dump(), ensure_task_protocol(child.context)["task_brief"], ) async def test_review_cannot_reserve_more_children_than_quota(self): briefs = [ { "objective": f"batch child {index}", "completion_criteria": [f"finish child {index}"], } for index in range(5) ] result = await agent(task_brief=briefs, context=self.context()) child_ids = [item["trace_id"] for item in result["sub_trace_ids"]] next_brief = { "objective": "one allowed reserved child", "completion_criteria": ["use the final available slot"], } first = await review_task_result( child_trace_id=child_ids[0], decision="ACCEPT_REFINE", reason="Reserve the sixth and final child slot", approved_next_task=next_brief, context=self.context(), ) self.assertEqual("completed", first["status"]) overflow = await review_task_result( child_trace_id=child_ids[1], decision="DESCEND_AGAIN", reason="A seventh child must not be approved", approved_next_task={ "objective": "overflow child", "completion_criteria": ["must never be scheduled"], }, context=self.context(), ) self.assertEqual("failed", overflow["status"]) self.assertIn("reserved=1", overflow["error"]) parent = await self.store.get_trace(self.root_id) state = ensure_task_protocol(parent.context) self.assertEqual(1, len(state["next_actions"])) self.assertIn(child_ids[1], state["pending_reviews"]) async def test_recursive_goal_completion_does_not_cascade_or_select_sibling(self): tree = GoalTree(mission="recursive") parent = tree.add_goals(["parent"])[0] first, second = tree.add_goals(["first", "second"], parent_id=parent.id) tree.focus(parent.id) tree.focus(first.id) tree.complete( first.id, "done", auto_advance=False, cascade_parent=False, ) self.assertEqual(parent.id, tree.current_id) self.assertEqual("pending", second.status) self.assertNotEqual("completed", parent.status) async def test_rewind_before_review_restores_report_and_review_gate(self): delegated = await agent(task_brief=BRIEF, context=self.context()) child_id = delegated["sub_trace_id"] reviewed = await review_task_result( child_trace_id=child_id, decision="ASCEND", reason="Initially accepted before rewind", context={**self.context(), "sequence": 20}, ) self.assertEqual("completed", reviewed["status"]) for sequence in (1, 15): await self.store.add_message(Message.create( trace_id=self.root_id, role="user", sequence=sequence, content=f"message-{sequence}", )) tree = await self.store.get_goal_tree(self.root_id) runner = AgentRunner( trace_store=self.store, tool_registry=get_tool_registry(), ) await runner._rewind(self.root_id, 15, tree) parent = await self.store.get_trace(self.root_id) state = ensure_task_protocol(parent.context) self.assertIn(child_id, state["pending_reviews"]) self.assertEqual([], state["reviews"]) restored_tree = await self.store.get_goal_tree(self.root_id) self.assertEqual("pending_review", restored_tree.find("1").status) schemas = runner._get_runtime_tool_schemas(RunConfig(), parent) self.assertEqual( {"review_task_result"}, {schema["function"]["name"] for schema in schemas}, ) async def test_grandchild_report_is_reviewed_only_by_direct_child(self): child_id = "manual-child" child_context = { "agent_mode": "recursive", "agent_mode_revision": 2, "agent_depth": 1, "root_trace_id": self.root_id, "created_by_tool": "agent", "task_protocol": new_task_protocol(BRIEF), } await self.store.create_trace(Trace( trace_id=child_id, mode="agent", task="child", parent_trace_id=self.root_id, uid="user-1", model="fake-model", context=child_context, )) child_tree = GoalTree( mission="child", goals=[Goal(id="1", description="delegate grandchild", status="in_progress")], current_id="1", ) await self.store.update_goal_tree(child_id, child_tree) grandchild = await agent( task_brief=BRIEF, context=self.context(trace_id=child_id, tree=child_tree), ) child = await self.store.get_trace(child_id) child_state = ensure_task_protocol(child.context) self.assertIn(grandchild["sub_trace_id"], child_state["pending_reviews"]) root = await self.store.get_trace(self.root_id) self.assertEqual({}, ensure_task_protocol(root.context)["pending_reviews"]) blocked = await agent( task_brief=BRIEF, context=self.context(trace_id=child_id, tree=child_tree), ) self.assertEqual("failed", blocked["status"]) reviewed = await review_task_result( child_trace_id=grandchild["sub_trace_id"], decision="ASCEND", reason="Grandchild satisfied the child criteria", context=self.context(trace_id=child_id, tree=child_tree), ) self.assertEqual("completed", reviewed["status"]) submitted = await submit_task_report( task_report=json.dumps(report_payload()), context={"store": self.store, "trace_id": child_id, "sequence": 20}, ) self.assertEqual("completed", submitted["status"]) async def test_child_must_execute_approved_action_before_own_report(self): child_id = "child-with-approved-action" child_context = { "agent_mode": "recursive", "agent_mode_revision": 2, "agent_depth": 1, "root_trace_id": self.root_id, "created_by_tool": "agent", "task_protocol": new_task_protocol(BRIEF), } await self.store.create_trace(Trace( trace_id=child_id, mode="agent", task="child", parent_trace_id=self.root_id, uid="user-1", model="fake-model", context=child_context, )) child_tree = GoalTree( mission="child", goals=[Goal(id="1", description="delegate", status="in_progress")], current_id="1", ) await self.store.update_goal_tree(child_id, child_tree) child_tool_context = self.context(trace_id=child_id, tree=child_tree) first = await agent(task_brief=BRIEF, context=child_tool_context) next_brief = { "objective": "check a second grandchild task", "completion_criteria": ["provide a second checked conclusion"], } reviewed = await review_task_result( child_trace_id=first["sub_trace_id"], decision="ACCEPT_REFINE", reason="A separate gap remains", approved_next_task=next_brief, context=child_tool_context, ) self.assertEqual("completed", reviewed["status"]) blocked_report = await submit_task_report( task_report=report_payload(), context={"store": self.store, "trace_id": child_id, "sequence": 20}, ) self.assertEqual("failed", blocked_report["status"]) self.assertIn("approved next action", blocked_report["error"]) second = await agent(task_brief=next_brief, context=child_tool_context) self.assertEqual("completed", second["status"]) child = await self.store.get_trace(child_id) state = ensure_task_protocol(child.context) self.assertEqual([], state["next_actions"]) self.assertIn(second["sub_trace_id"], state["pending_reviews"]) async def test_goal_tool_cannot_bypass_queued_next_action(self): root = await self.store.get_trace(self.root_id) state = ensure_task_protocol(root.context) state["next_actions"] = [{ "decision": "ACCEPT_REFINE", "child_trace_id": "reviewed-child", "goal_id": "1", "task_brief": BRIEF, "created_at_sequence": 5, }] await self.store.update_trace(self.root_id, context=root.context) result = await goal_tool( tree=self.tree, store=self.store, trace_id=self.root_id, done="attempt to bypass the approved action", ) self.assertIn("已批准的下一步待执行", result) persisted_tree = await self.store.get_goal_tree(self.root_id) self.assertEqual("in_progress", persisted_tree.find("1").status) async def test_failed_report_allows_only_revise_or_fail(self): child_id = "failed-child" failed_report = TaskReport( child_trace_id=child_id, **report_payload(outcome="failed"), ) child_context = { "agent_mode": "recursive", "agent_mode_revision": 2, "agent_depth": 1, "root_trace_id": self.root_id, "created_by_tool": "agent", "task_protocol": new_task_protocol(BRIEF), } child_context["task_protocol"]["task_report"] = failed_report.model_dump() await self.store.create_trace(Trace( trace_id=child_id, mode="agent", task="failed", parent_trace_id=self.root_id, uid="user-1", context=child_context, )) root = await self.store.get_trace(self.root_id) state = ensure_task_protocol(root.context) state["pending_reviews"][child_id] = { "goal_id": "1", "task_report": failed_report.model_dump(), "received_at_sequence": 2, } await self.store.update_trace(self.root_id, context=root.context) rejected = await review_task_result( child_trace_id=child_id, decision="ASCEND", reason="incorrect", context=self.context(), ) self.assertEqual("failed", rejected["status"]) self.assertIn("only allow", rejected["error"]) class RunnerCompletionGateTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.temp_dir = tempfile.TemporaryDirectory() self.store = FileSystemTraceStore(self.temp_dir.name) self.trace_id = "gate-child" await self.store.create_trace(Trace( trace_id=self.trace_id, mode="agent", task="must report", parent_trace_id="root", uid="user-1", context={ "agent_mode": "recursive", "agent_mode_revision": 2, "agent_depth": 1, "root_trace_id": "root", "created_by_tool": "agent", "task_protocol": new_task_protocol(BRIEF), }, )) await self.store.update_goal_tree(self.trace_id, GoalTree(mission="must report")) async def asyncTearDown(self): self.temp_dir.cleanup() async def test_missing_report_becomes_protocol_error_after_two_corrections(self): calls = 0 async def llm_call(**kwargs): nonlocal calls calls += 1 return { "content": "I am done", "tool_calls": [], "finish_reason": "stop", } runner = AgentRunner(trace_store=self.store, llm_call=llm_call) result = await runner.run_result( messages=[], config=RunConfig( trace_id=self.trace_id, max_iterations=8, knowledge=KnowledgeConfig( enable_extraction=False, enable_completion_extraction=False, enable_injection=False, ), ), ) self.assertEqual("failed", result["status"]) trace = await self.store.get_trace(self.trace_id) report = ensure_task_protocol(trace.context)["task_report"] self.assertEqual("protocol_error", report["outcome"]) self.assertGreaterEqual(calls, 3) async def test_valid_report_allows_child_to_finish(self): calls = 0 async def llm_call(**kwargs): nonlocal calls calls += 1 if calls == 1: return { "content": "", "tool_calls": [{ "id": "submit-1", "type": "function", "function": { "name": "submit_task_report", "arguments": json.dumps({ "task_report": report_payload(), }), }, }], "finish_reason": "tool_calls", } return { "content": "Structured report submitted.", "tool_calls": [], "finish_reason": "stop", } runner = AgentRunner(trace_store=self.store, llm_call=llm_call) result = await runner.run_result( messages=[], config=RunConfig( trace_id=self.trace_id, max_iterations=4, knowledge=KnowledgeConfig( enable_extraction=False, enable_completion_extraction=False, enable_injection=False, ), ), ) self.assertEqual("completed", result["status"]) trace = await self.store.get_trace(self.trace_id) self.assertEqual( "satisfied", ensure_task_protocol(trace.context)["task_report"]["outcome"], ) async def test_lifecycle_tool_batch_is_rejected_before_execution(self): root_id = "batch-gate-root" await self.store.create_trace(Trace( trace_id=root_id, mode="agent", task="root", uid="user-1", context={ "agent_mode": "recursive", "agent_mode_revision": 2, "agent_depth": 0, "root_trace_id": root_id, "task_protocol": new_task_protocol(), }, )) await self.store.update_goal_tree(root_id, GoalTree(mission="root")) calls = 0 async def llm_call(**kwargs): nonlocal calls calls += 1 if calls == 1: return { "content": "", "tool_calls": [ { "id": "agent-1", "type": "function", "function": { "name": "agent", "arguments": json.dumps({ "task_brief": BRIEF, }), }, }, { "id": "context-1", "type": "function", "function": { "name": "get_current_context", "arguments": "{}", }, }, ], "finish_reason": "tool_calls", } return { "content": "Corrected the invalid batch.", "tool_calls": [], "finish_reason": "stop", } runner = AgentRunner(trace_store=self.store, llm_call=llm_call) result = await runner.run_result( messages=[], config=RunConfig( trace_id=root_id, max_iterations=3, knowledge=KnowledgeConfig( enable_extraction=False, enable_completion_extraction=False, enable_injection=False, ), ), ) self.assertEqual("completed", result["status"]) children = await self.store.list_traces(parent_trace_id=root_id, limit=10) self.assertEqual([], children) tool_messages = [ message for message in await self.store.get_trace_messages(root_id) if message.role == "tool" ] self.assertEqual(2, len(tool_messages)) self.assertTrue(all( "only tool call" in str(message.content) for message in tool_messages )) async def test_rewind_clears_protocol_state_after_cutoff(self): trace = await self.store.get_trace(self.trace_id) state = ensure_task_protocol(trace.context) report = TaskReport( child_trace_id=self.trace_id, **report_payload(), ).model_dump() state["task_report"] = report state["task_report_submitted_at_sequence"] = 5 state["pending_reviews"] = { "grandchild": { "goal_id": "1", "task_report": report, "received_at_sequence": 6, } } state["reviews"] = [{"reviewed_at_sequence": 7}] state["next_actions"] = [{"created_at_sequence": 8}] await self.store.update_trace(self.trace_id, context=trace.context) tree = GoalTree( mission="rewind", goals=[Goal(id="1", description="active", status="pending_review")], current_id="1", ) await self.store.update_goal_tree(self.trace_id, tree) for sequence in (1, 2): await self.store.add_message(Message.create( trace_id=self.trace_id, role="user", sequence=sequence, content=f"message-{sequence}", )) runner = AgentRunner(trace_store=self.store) cutoff = await runner._rewind(self.trace_id, 2, tree) self.assertEqual(3, cutoff) rewound = await self.store.get_trace(self.trace_id) state = ensure_task_protocol(rewound.context) self.assertIsNone(state["task_report"]) self.assertEqual(1, len(state["report_history"])) self.assertEqual({}, state["pending_reviews"]) self.assertEqual([], state["reviews"]) self.assertEqual([], state["next_actions"]) rewound_tree = await self.store.get_goal_tree(self.trace_id) self.assertEqual("pending", rewound_tree.find("1").status) async def test_rewind_restores_preserved_pending_review_projection(self): trace = await self.store.get_trace(self.trace_id) state = ensure_task_protocol(trace.context) report = TaskReport( child_trace_id="preserved-grandchild", **report_payload(), ).model_dump() state["pending_reviews"] = { "preserved-grandchild": { "goal_id": "1", "task_report": report, "received_at_sequence": 2, } } await self.store.update_trace(self.trace_id, context=trace.context) tree = GoalTree( mission="rewind", goals=[Goal(id="1", description="review", status="pending_review")], current_id="1", ) await self.store.update_goal_tree(self.trace_id, tree) for sequence in (1, 2): await self.store.add_message(Message.create( trace_id=self.trace_id, role="user", sequence=sequence, content=f"message-{sequence}", )) runner = AgentRunner(trace_store=self.store) await runner._rewind(self.trace_id, 2, tree) rewound = await self.store.get_trace(self.trace_id) self.assertIn( "preserved-grandchild", ensure_task_protocol(rewound.context)["pending_reviews"], ) rewound_tree = await self.store.get_goal_tree(self.trace_id) self.assertEqual("pending_review", rewound_tree.find("1").status) async def test_root_cannot_finish_with_unconsumed_next_action(self): root_id = "next-action-gate-root" state = new_task_protocol() state["next_actions"] = [{ "decision": "ACCEPT_REFINE", "child_trace_id": "reviewed-child", "goal_id": None, "task_brief": BRIEF, "created_at_sequence": 1, }] await self.store.create_trace(Trace( trace_id=root_id, mode="agent", task="root", uid="user-1", context={ "agent_mode": "recursive", "agent_mode_revision": 2, "agent_depth": 0, "root_trace_id": root_id, "task_protocol": state, }, )) await self.store.update_goal_tree(root_id, GoalTree(mission="root")) calls = 0 async def llm_call(**kwargs): nonlocal calls calls += 1 return { "content": "Trying to finish without the approved action.", "tool_calls": [], "finish_reason": "stop", } runner = AgentRunner(trace_store=self.store, llm_call=llm_call) result = await runner.run_result( messages=[], config=RunConfig( trace_id=root_id, max_iterations=5, knowledge=KnowledgeConfig( enable_extraction=False, enable_completion_extraction=False, enable_injection=False, ), ), ) self.assertEqual("failed", result["status"]) self.assertEqual(3, calls) root = await self.store.get_trace(root_id) self.assertEqual(1, len(ensure_task_protocol(root.context)["next_actions"])) class RuntimeProtocolSchemaTest(unittest.TestCase): def setUp(self): self.runner = AgentRunner(tool_registry=get_tool_registry()) def names(self, schemas): return {schema["function"]["name"] for schema in schemas} def test_legacy_and_recursive_runtime_schemas_are_isolated(self): legacy = Trace( trace_id="legacy", mode="agent", context={"agent_mode": "legacy", "agent_mode_revision": 1}, ) legacy_schemas = self.runner._get_runtime_tool_schemas(RunConfig(), legacy) self.assertNotIn("submit_task_report", self.names(legacy_schemas)) self.assertNotIn("review_task_result", self.names(legacy_schemas)) legacy_agent = next( item for item in legacy_schemas if item["function"]["name"] == "agent" ) properties = legacy_agent["function"]["parameters"]["properties"] self.assertIn("task", properties) self.assertNotIn("task_brief", properties) child = Trace( trace_id="child", mode="agent", parent_trace_id="root", context={ "agent_mode": "recursive", "agent_mode_revision": 2, "task_protocol": new_task_protocol(BRIEF), }, ) child_schemas = self.runner._get_runtime_tool_schemas(RunConfig(), child) self.assertIn("submit_task_report", self.names(child_schemas)) self.assertNotIn("review_task_result", self.names(child_schemas)) state = ensure_task_protocol(child.context) state["pending_reviews"]["grandchild"] = {"task_report": {}} review_schemas = self.runner._get_runtime_tool_schemas(RunConfig(), child) self.assertEqual({"review_task_result"}, self.names(review_schemas)) def test_protocol_tools_publish_nested_object_schemas(self): schemas = { schema["function"]["name"]: schema for schema in get_tool_registry().get_schemas([ "submit_task_report", "review_task_result", ]) } report_schema = schemas["submit_task_report"]["function"]["parameters"] task_report = report_schema["properties"]["task_report"] self.assertEqual("object", task_report["type"]) self.assertIn("validation", task_report["properties"]) self.assertIn("next_step_suggestion", task_report["properties"]) review_schema = schemas["review_task_result"]["function"]["parameters"] self.assertEqual( [ "ACCEPT_REFINE", "REJECT_REFINE", "REVISE_CHILD", "DESCEND_AGAIN", "ASCEND", "FAIL", ], review_schema["properties"]["decision"]["enum"], ) self.assertEqual( "object", review_schema["properties"]["approved_next_task"]["type"], ) if __name__ == "__main__": unittest.main()