Agent/tests/test_task_protocol.py

import json
import tempfile
import unittest
from unittest.mock import AsyncMock, patch

from pydantic import ValidationError

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 (
    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 _record_pending_task_reports, 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
from cyber_agent.core.validation import ValidationResult, ValidationRun


BRIEF = {
    "objective": "research one concrete question",
    "reason": "the parent needs one independently checked result",
    "completion_criteria": ["provide one checked conclusion"],
    "expected_outputs": ["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"],
    }


def validation_payload(child_trace_id, *, outcome="passed"):
    if outcome == "passed":
        issues = []
        retry_from = None
    elif outcome == "failed":
        issues = ["task needs revision"]
        retry_from = "task_definition"
    else:
        issues = ["validator could not evaluate the task"]
        retry_from = None
    return ValidationResult(
        validator_trace_id=f"{child_trace_id}@validator-fixture",
        evaluated_trace_id=child_trace_id,
        outcome=outcome,
        scope="task",
        reason="fixture validation",
        issues=issues,
        retry_from=retry_from,
    ).model_dump()


async def create_recursive_root(store, trace_id, *, state=None, uid="user-1"):
    budget = ResourceBudget()
    context = {
        "agent_mode": "recursive",
        "agent_mode_revision": 2,
        "agent_depth": 0,
        "root_trace_id": trace_id,
        "root_completion_criteria": ["finish the root task"],
        RESOURCE_BUDGET_CONTEXT_KEY: budget.to_dict(),
        "task_protocol": state or new_task_protocol(),
    }
    await store.create_trace(Trace(
        trace_id=trace_id,
        mode="agent",
        task="root",
        uid=uid,
        model="fake-model",
        context=context,
    ))
    await ResourceBudgetController(store).initialize(
        trace_id,
        budget,
        initial_agents=1,
    )
    return await store.get_trace(trace_id)


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 = []
        self.resource_budget = ResourceBudgetController(store)
        self.validation_count = 0

    async def _resource_budget_for_trace(self, trace_id):
        trace = await self.store.get_trace(trace_id)
        root_id = trace.context["root_trace_id"]
        root = trace if root_id == trace_id else await self.store.get_trace(root_id)
        return root_id, ResourceBudget.from_dict(
            root.context[RESOURCE_BUDGET_CONTEXT_KEY]
        )

    async def validate_recursive_trace(
        self,
        evaluated_trace_id,
        *,
        scope,
        deterministic_failure=None,
        **_kwargs,
    ):
        self.validation_count += 1
        evaluated = await self.store.get_trace(evaluated_trace_id)
        validator_trace_id = (
            f"{evaluated_trace_id}@validator-fixture-{self.validation_count}"
        )
        if deterministic_failure:
            outcome = deterministic_failure.get("outcome", "error")
            issues = deterministic_failure.get("issues") or [
                deterministic_failure.get("reason", "validation failed")
            ]
            retry_from = deterministic_failure.get("retry_from")
        else:
            outcome = "passed"
            issues = []
            retry_from = None
        result = ValidationResult(
            validator_trace_id=validator_trace_id,
            evaluated_trace_id=evaluated_trace_id,
            outcome=outcome,
            scope=scope,
            reason="fixture validation",
            issues=issues,
            retry_from=retry_from,
        )
        await self.store.create_trace(Trace(
            trace_id=validator_trace_id,
            mode="agent",
            task=f"validate {evaluated_trace_id}",
            parent_trace_id=evaluated_trace_id,
            uid=evaluated.uid,
            status="completed",
            tools=[],
            context={
                "created_by_tool": "validator",
                "root_trace_id": evaluated.context["root_trace_id"],
                "agent_depth": evaluated.context["agent_depth"],
            },
        ))
        return ValidationRun(result=result, trace_id=validator_trace_id)

    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"
        await create_recursive_root(self.store, self.root_id)
        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,
            RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY: self.runner.tools.get_tool_names(),
        }

    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, "parent_findings": [], "context": {}},
            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_child_inherits_brief_constraints_not_parent_runtime_context(self):
        parent = await self.store.get_trace(self.root_id)
        parent.context["parent_private_runtime"] = {"secret": "must-not-leak"}
        await self.store.update_trace(self.root_id, context=parent.context)

        result = await agent(task_brief=BRIEF, context=self.context())

        child = await self.store.get_trace(result["sub_trace_id"])
        self.assertNotIn("parent_private_runtime", child.context)
        child_brief = ensure_task_protocol(child.context)["task_brief"]
        self.assertEqual(["do not guess"], child_brief["constraints"])
        child_config_context = self.runner.configs[-1].context
        self.assertNotIn("parent_private_runtime", child_config_context)

    async def test_same_task_report_version_reuses_its_validator_result(self):
        result = await agent(task_brief=BRIEF, context=self.context())
        child_id = result["sub_trace_id"]
        self.assertEqual(1, self.runner.validation_count)
        parent = await self.store.get_trace(self.root_id)

        await _record_pending_task_reports(
            self.store,
            self.runner,
            parent,
            "1",
            [(child_id, {"status": "completed", "summary": "same report"})],
            10,
        )

        self.assertEqual(1, self.runner.validation_count)
        validators = await self.store.list_traces(
            parent_trace_id=child_id,
            created_by_tool="validator",
            limit=10,
        )
        self.assertEqual(1, len(validators))

    async def test_recursive_messages_are_rejected(self):
        result = await agent(
            task_brief=json.dumps(BRIEF),
            messages=json.dumps([{
                "role": "system",
                "content": "controlled child context",
            }]),
            context=self.context(),
        )
        self.assertEqual("failed", result["status"])
        self.assertIn("does not accept messages", result["error"])

    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",
            "reason": "the first report identified a separate gap",
            "completion_criteria": ["close the gap"],
            "expected_outputs": ["a checked gap analysis"],
        }
        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",
                "reason": "the batch needs a second independent result",
                "completion_criteria": ["provide second conclusion"],
                "expected_outputs": ["a 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",
                    "reason": "exercise queued batch behavior",
                    "completion_criteria": ["finish the second child"],
                    "expected_outputs": ["the second child result"],
                },
            ],
            context=self.context(),
        )
        child_ids = [item["trace_id"] for item in result["sub_trace_ids"]]
        next_brief = {
            "objective": "approved follow-up",
            "reason": "the parent approved one follow-up",
            "completion_criteria": ["finish the approved follow-up"],
            "expected_outputs": ["the approved follow-up result"],
        }
        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",
                "reason": "attempt a mismatched revision",
                "completion_criteria": ["this must be rejected"],
                "expected_outputs": ["a rejected replacement"],
            },
            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}",
                "reason": "fill the local child quota",
                "completion_criteria": [f"finish child {index}"],
                "expected_outputs": [f"result for 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",
            "reason": "use the final available child slot",
            "completion_criteria": ["use the final available slot"],
            "expected_outputs": ["one final child result"],
        }
        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",
                "reason": "exercise quota rejection",
                "completion_criteria": ["must never be scheduled"],
                "expected_outputs": ["no output because creation is rejected"],
            },
            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",
            "reason": "the child approved another direct task",
            "completion_criteria": ["provide a second checked conclusion"],
            "expected_outputs": ["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,
            parent_goal_id="1",
            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(),
            "validation_result": validation_payload(child_id, outcome="failed"),
            "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 create_recursive_root(self.store, "root")
        controller = ResourceBudgetController(self.store)
        budget = ResourceBudget()
        await controller.reserve_agents("root", budget, 1)
        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 kwargs.get("tools") == []:
                return {
                    "content": json.dumps({
                        "outcome": "passed",
                        "scope": "root",
                        "reason": "the corrected behavior meets the criterion",
                        "issues": [],
                        "retry_from": None,
                    }),
                    "tool_calls": [],
                    "finish_reason": "stop",
                }
            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 create_recursive_root(self.store, root_id)
        await self.store.update_goal_tree(root_id, GoalTree(mission="root"))
        calls = 0

        async def llm_call(**kwargs):
            nonlocal calls
            calls += 1
            if kwargs.get("tools") == []:
                return {
                    "content": json.dumps({
                        "outcome": "passed",
                        "scope": "root",
                        "reason": "the corrected lifecycle behavior is valid",
                        "issues": [],
                        "retry_from": None,
                    }),
                    "tool_calls": [],
                    "finish_reason": "stop",
                }
            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,
            created_by_tool="agent",
            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,
                "validation_result": validation_payload("grandchild"),
                "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,
                "validation_result": validation_payload(
                    "preserved-grandchild"
                ),
                "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 create_recursive_root(self.store, root_id, state=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))
        self.assertNotIn("evaluate", self.names(child_schemas))
        recursive_agent = next(
            item for item in child_schemas if item["function"]["name"] == "agent"
        )
        recursive_properties = recursive_agent["function"]["parameters"]["properties"]
        self.assertNotIn("messages", recursive_properties)
        brief_object = recursive_properties["task_brief"]["oneOf"][0]
        self.assertEqual(
            {
                "objective",
                "reason",
                "completion_criteria",
                "expected_outputs",
            },
            set(brief_object["required"]),
        )

        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",
                "REPLAN_CURRENT",
                "ASCEND",
                "FAIL",
            ],
            review_schema["properties"]["decision"]["enum"],
        )
        self.assertEqual(
            "object",
            review_schema["properties"]["approved_next_task"]["type"],
        )


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