|
|
@@ -0,0 +1,1195 @@
|
|
|
+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()
|
SamLee