Agent/tests/test_task_protocol.py

import asyncio
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.run_snapshot import (
    RunConfigSnapshotV1,
    persist_run_config_snapshot,
)
from cyber_agent.core.context_policy import (
    add_context_snapshot,
    create_context_snapshot,
    persist_root_task_anchor,
)
from cyber_agent.core.task_protocol import (
    TaskProgress,
    TaskBrief,
    TaskReport,
    TaskReview,
    append_task_progress_revision,
    ensure_task_protocol,
    initialize_task_progress,
    new_task_protocol,
    replace_task_brief,
    task_contract_ref,
)
from cyber_agent.core.task_protocol_service import TaskProtocolService
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,
    update_task_progress,
)
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 (
    ScopeValidationResult,
    ValidationCheck,
    ValidationPolicy,
    ValidationResult,
    ValidationRun,
    ValidatorSettings,
    persist_validation_policy,
)


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"],
}
ROOT_ANCHOR = {
    "objective": "finish the root task",
    "completion_criteria": ["finish the root task"],
    "constraints": ["do not guess"],
}


def with_root_anchor(context):
    persist_root_task_anchor(context, ROOT_ANCHOR)
    state = ensure_task_protocol(context)
    if state.get("task_progress_head_revision") is None:
        initialize_task_progress(
            state,
            root_task_anchor_hash=context.get("root_task_anchor_hash"),
        )
    current = state.get("task_progress_head_revision")
    if current == 1:
        append_task_progress_revision(
            state,
            TaskProgress(phase="ready_to_submit"),
            expected_revision=current,
            effective_at_sequence=0,
            contract_ref=task_contract_ref(
                state,
                root_task_anchor_hash=context.get("root_task_anchor_hash"),
            ),
        )
    return context


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"):
    plan_hash = "a" * 64
    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
    scope_result = ScopeValidationResult(
        validator_trace_id=f"{child_trace_id}@validator-fixture",
        scope="task",
        outcome=outcome,
        checks=[ValidationCheck(
            check_id="task.fixture",
            status=("unknown" if outcome == "error" else outcome),
            issue=None if outcome == "passed" else issues[0],
        )],
        reason="fixture validation",
        retry_from=retry_from,
        plan_hash=plan_hash,
    )
    return ValidationResult(
        evaluated_trace_id=child_trace_id,
        outcome=outcome,
        scope_results=[scope_result],
        issues=issues,
        retry_from=retry_from,
        plan_hash=plan_hash,
    ).model_dump(mode="json")


def passed_validator_response(kwargs):
    """按当前 ValidationPlan 返回完整的通过结果。"""
    messages = kwargs.get("messages") or []
    packet = json.loads(messages[1]["content"])
    scope = packet["validation_scope"]
    checks = [
        {
            "check_id": item["check_id"],
            "status": "passed",
            "evidence_refs": [],
            "issue": None,
        }
        for item in packet["validation_plan"]["checks"]
        if item["scope"] == scope
    ]
    return {
        "content": json.dumps({
            "outcome": "passed",
            "scope": scope,
            "checks": checks,
            "reason": "all planned checks passed",
            "retry_from": None,
        }),
        "tool_calls": [],
        "finish_reason": "stop",
    }


async def create_recursive_root(store, trace_id, *, state=None, uid="user-1"):
    budget = ResourceBudget()
    context = {
        "agent_mode": "recursive",
        "agent_mode_revision": 3,
        "agent_depth": 0,
        "root_trace_id": trace_id,
        RESOURCE_BUDGET_CONTEXT_KEY: budget.to_dict(),
        "task_protocol": state or new_task_protocol(),
    }
    with_root_anchor(context)
    persist_validation_policy(
        context,
        ValidationPolicy(),
        ValidatorSettings(),
    )
    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",
            "update_task_progress",
            "submit_task_report",
            "review_task_result",
            "read_context_ref",
        ]


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,
    ):
        evaluated = await self.store.get_trace(evaluated_trace_id)
        child_state = ensure_task_protocol(evaluated.context)
        cached = child_state.get("task_report_validation")
        if isinstance(cached, dict) and cached.get("aggregate_result"):
            result = ValidationResult.model_validate(cached["aggregate_result"])
            return ValidationRun(
                result=result,
                trace_ids=[
                    item.validator_trace_id for item in result.scope_results
                ],
                cached=True,
            )
        self.validation_count += 1
        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.model_validate(
            validation_payload(evaluated_trace_id, outcome=outcome)
        )
        child_state["task_report_validation"] = {
            "validation_plan": None,
            "plan_hash": result.plan_hash,
            "scope_results": [
                item.model_dump(mode="json") for item in result.scope_results
            ],
            "aggregate_result": result.model_dump(mode="json"),
            "validated_at_sequence": evaluated.head_sequence,
        }
        await self.store.update_trace(
            evaluated_trace_id,
            context=evaluated.context,
        )
        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=with_root_anchor({
                "created_by_tool": "validator",
                "root_trace_id": evaluated.context["root_trace_id"],
                "agent_depth": evaluated.context["agent_depth"],
            }),
        ))
        return ValidationRun(result=result, trace_ids=[validator_trace_id])

    async def run_result(self, messages, config, on_event=None):
        self.configs.append(config)
        trace = await self.store.get_trace(config.trace_id)
        current_revision = ensure_task_protocol(trace.context)[
            "task_progress_head_revision"
        ]
        progressed = await update_task_progress(
            expected_revision=current_revision,
            progress={
                "phase": "ready_to_submit",
                "questions": [],
                "blockers": [],
                "findings": [],
                "hypotheses": [],
                "work_items": [],
                "decision_rationale": "The fixture completed its assigned work.",
            },
            context={
                "store": self.store,
                "task_protocol_service": TaskProtocolService(self.store),
                "trace_id": config.trace_id,
                "sequence": 2,
            },
        )
        if progressed["status"] != "completed":
            raise AssertionError(progressed)
        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": {},
                "context_refs": [],
                "validation_scopes": [],
            },
            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", "read_context_ref"},
            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_persisted_review_command_replays_idempotently_and_rejects_tampering(self):
        result = await agent(task_brief=BRIEF, context=self.context())
        review_context = {
            **self.context(),
            "tool_call_id": "durable-review-command",
        }
        first = await review_task_result(
            child_trace_id=result["sub_trace_id"],
            decision="ASCEND",
            reason="The exact persisted report is accepted",
            context=review_context,
        )
        self.assertEqual("completed", first["status"])

        replayed = await review_task_result(
            child_trace_id=result["sub_trace_id"],
            decision="ASCEND",
            reason="The exact persisted report is accepted",
            context={**review_context, "sequence": 999},
        )
        self.assertEqual(first["task_review"], replayed["task_review"])
        parent = await self.store.get_trace(self.root_id)
        state = ensure_task_protocol(parent.context)
        self.assertEqual(1, len(state["reviews"]))
        self.assertEqual({}, state["pending_reviews"])

        tampered = await review_task_result(
            child_trace_id=result["sub_trace_id"],
            decision="ASCEND",
            reason="Changed after the durable command was committed",
            context={**review_context, "sequence": 1_000},
        )
        self.assertEqual("failed", tampered["status"])
        self.assertIn("payload changed", tampered["error"])

    async def test_revise_child_projection_recovers_after_parent_commit_crash(self):
        result = await agent(task_brief=BRIEF, context=self.context())
        child_id = result["sub_trace_id"]
        review_context = {
            **self.context(),
            "tool_call_id": "durable-revise-child",
        }
        original_update_trace = self.store.update_trace
        interrupted = False

        async def interrupt_child_projection(trace_id, **updates):
            nonlocal interrupted
            context = updates.get("context")
            child_state = (
                context.get("task_protocol", {})
                if isinstance(context, dict)
                else {}
            )
            if (
                not interrupted
                and trace_id == child_id
                and child_state.get("task_report") is None
            ):
                interrupted = True
                raise asyncio.CancelledError()
            return await original_update_trace(trace_id, **updates)

        with patch.object(
            self.store,
            "update_trace",
            side_effect=interrupt_child_projection,
        ):
            with self.assertRaises(asyncio.CancelledError):
                await review_task_result(
                    child_trace_id=child_id,
                    decision="REVISE_CHILD",
                    reason="Revise only the child's current output",
                    context=review_context,
                )
        self.assertTrue(interrupted)
        parent = await self.store.get_trace(self.root_id)
        parent_state = ensure_task_protocol(parent.context)
        self.assertEqual(1, len(parent_state["reviews"]))
        self.assertEqual(1, len(parent_state["next_actions"]))
        self.assertEqual({}, parent_state["pending_reviews"])
        child = await self.store.get_trace(child_id)
        self.assertIsNotNone(ensure_task_protocol(child.context)["task_report"])

        replayed = await review_task_result(
            child_trace_id=child_id,
            decision="REVISE_CHILD",
            reason="Revise only the child's current output",
            context={**review_context, "sequence": 999},
        )
        self.assertEqual("completed", replayed["status"])
        parent = await self.store.get_trace(self.root_id)
        parent_state = ensure_task_protocol(parent.context)
        self.assertEqual(1, len(parent_state["reviews"]))
        self.assertEqual(1, len(parent_state["next_actions"]))
        child = await self.store.get_trace(child_id)
        child_state = ensure_task_protocol(child.context)
        self.assertIsNone(child_state["task_report"])
        self.assertIsNone(child_state["task_report_validation"])
        self.assertEqual(1, len(child_state["report_history"]))

    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"])
        self.assertIsNotNone(reviewed["context_ref"])
        self.assertIn(
            reviewed["context_ref"],
            reviewed["next_action_task_brief"]["context_refs"],
        )
        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", "read_context_ref"},
            {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=reviewed["next_action_task_brief"],
            context=self.context(),
        )
        self.assertEqual("completed", second["status"])
        self.assertEqual(2, len(await self.children()))
        second_child = await self.store.get_trace(second["sub_trace_id"])
        self.assertIn(
            reviewed["context_ref"]["ref_id"],
            second_child.context["context_access"]["snapshots"],
        )
        parent = await self.store.get_trace(self.root_id)
        self.assertEqual([], ensure_task_protocol(parent.context)["next_actions"])

    async def test_review_rejects_an_unusable_context_ref_without_consuming_report(self):
        first = await agent(task_brief=BRIEF, context=self.context())
        invalid_next = {
            "objective": "use an unauthorized context ref",
            "reason": "exercise approval validation",
            "completion_criteria": ["must not be approved"],
            "expected_outputs": ["nothing"],
            "context_refs": [{
                "ref_id": "ctx_000000000000000000000000",
                "version": "0" * 64,
            }],
        }
        reviewed = await review_task_result(
            child_trace_id=first["sub_trace_id"],
            decision="ACCEPT_REFINE",
            reason="try an invalid transfer",
            approved_next_task=invalid_next,
            context=self.context(),
        )
        self.assertEqual("failed", reviewed["status"])
        self.assertIn("not authorized", reviewed["error"])
        parent = await self.store.get_trace(self.root_id)
        state = ensure_task_protocol(parent.context)
        self.assertIn(first["sub_trace_id"], state["pending_reviews"])
        self.assertEqual([], state["reviews"])

    async def test_oversized_review_snapshot_warns_without_blocking_review(self):
        first = await agent(task_brief=BRIEF, context=self.context())
        parent = await self.store.get_trace(self.root_id)
        state = ensure_task_protocol(parent.context)
        state["pending_reviews"][first["sub_trace_id"]]["task_report"]["outputs"] = [
            {"payload": "x" * 16_000}
        ]
        await self.store.update_trace(self.root_id, context=parent.context)

        reviewed = await review_task_result(
            child_trace_id=first["sub_trace_id"],
            decision="ASCEND",
            reason="The task passed but its transferable snapshot is too large",
            context=self.context(),
        )
        self.assertEqual("completed", reviewed["status"])
        self.assertIsNone(reviewed["context_ref"])
        self.assertIn("16000", reviewed["context_ref_warning"])
        parent = await self.store.get_trace(self.root_id)
        self.assertEqual({}, ensure_task_protocol(parent.context)["pending_reviews"])

    async def test_batch_context_refs_are_validated_before_any_trace_is_created(self):
        invalid_ref_brief = {
            **BRIEF,
            "objective": "second child with an unauthorized reference",
            "context_refs": [{
                "ref_id": "ctx_000000000000000000000000",
                "version": "0" * 64,
            }],
        }
        result = await agent(
            task_brief=[BRIEF, invalid_ref_brief],
            context=self.context(),
        )
        self.assertEqual("failed", result["status"])
        self.assertIn("not authorized", result["error"])
        self.assertEqual([], await self.children())

    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_submit_and_review_reject_child_anchor_that_differs_from_root(self):
        child_context = with_root_anchor({
            "agent_mode": "recursive",
            "agent_mode_revision": 3,
            "agent_depth": 1,
            "root_trace_id": self.root_id,
            "created_by_tool": "agent",
            "task_protocol": new_task_protocol(BRIEF),
        })
        persist_root_task_anchor(child_context, {
            **ROOT_ANCHOR,
            "objective": "self-consistent but unauthorized root objective",
        })
        child_id = "mismatched-anchor-child"
        await self.store.create_trace(Trace(
            trace_id=child_id,
            mode="agent",
            task="child",
            parent_trace_id=self.root_id,
            parent_goal_id="1",
            uid="user-1",
            model="fake-model",
            context=child_context,
        ))
        submitted = await submit_task_report(
            task_report=report_payload(),
            context={"store": self.store, "trace_id": child_id},
        )
        self.assertEqual("failed", submitted["status"])
        self.assertIn("does not match", submitted["error"])

        parent = await self.store.get_trace(self.root_id)
        parent_state = ensure_task_protocol(parent.context)
        parent_state["pending_reviews"][child_id] = {
            "goal_id": "1",
            "task_report": TaskReport(
                child_trace_id=child_id,
                **report_payload(),
            ).model_dump(),
            "validation_result": validation_payload(child_id),
            "received_at_sequence": 1,
        }
        await self.store.update_trace(self.root_id, context=parent.context)
        reviewed = await review_task_result(
            child_trace_id=child_id,
            decision="ASCEND",
            reason="must not accept a different root anchor",
            context=self.context(),
        )
        self.assertEqual("failed", reviewed["status"])
        self.assertIn("does not match", reviewed["error"])
        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"])

    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_is_rejected_by_recursive_revision_2(self):
        remote_call = AsyncMock(return_value={
            "status": "completed",
            "summary": "remote result",
        })
        with patch(
            "cyber_agent.tools.builtin.subagent._run_remote_agent",
            new=remote_call,
        ):
            result = await agent(
                task="remote free text",
                agent_type="remote_librarian",
                context=self.context(),
            )
        self.assertEqual("failed", result["status"])
        self.assertIn("does not support remote agents", result["error"])
        remote_call.assert_not_awaited()

    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"])
        cleared_child = await self.store.get_trace(child_id)
        self.assertEqual({}, cleared_child.context["context_access"]["snapshots"])

        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_brief = reviewed["next_action_task_brief"]
        revised = await agent(
            task_brief=revised_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(revised_brief).model_dump(),
            ensure_task_protocol(child.context)["task_brief"],
        )
        self.assertEqual(2, ensure_task_protocol(child.context)["task_brief_version"])
        self.assertEqual(
            1,
            len(ensure_task_protocol(child.context)["task_brief_history"]),
        )

    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", "read_context_ref"},
            {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 = with_root_anchor({
            "agent_mode": "recursive",
            "agent_mode_revision": 3,
            "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 = with_root_anchor({
            "agent_mode": "recursive",
            "agent_mode_revision": 3,
            "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=reviewed["next_action_task_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 = with_root_anchor({
            "agent_mode": "recursive",
            "agent_mode_revision": 3,
            "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()
        failed_validation = validation_payload(child_id, outcome="failed")
        child_context["task_protocol"]["task_report_validation"] = {
            "validation_plan": None,
            "plan_hash": failed_validation["plan_hash"],
            "scope_results": failed_validation["scope_results"],
            "aggregate_result": failed_validation,
            "validated_at_sequence": 2,
        }
        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": failed_validation,
            "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=with_root_anchor({
                "agent_mode": "recursive",
                "agent_mode_revision": 3,
                "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 bind_snapshot(self, trace_id, config):
        """Make manually-created Recursive fixtures equivalent to new traces."""
        trace = await self.store.get_trace(trace_id)
        config.model = trace.model or config.model
        config.uid = trace.uid
        config.agent_type = trace.agent_type or "default"
        persist_run_config_snapshot(
            trace.context,
            RunConfigSnapshotV1.from_run_config(
                config,
                memory_identity=None,
            ),
        )
        await self.store.update_trace(trace_id, context=trace.context)
        return config

    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)
        config = await self.bind_snapshot(
            self.trace_id,
            RunConfig(
                trace_id=self.trace_id,
                max_iterations=8,
                knowledge=KnowledgeConfig(
                    enable_extraction=False,
                    enable_completion_extraction=False,
                    enable_injection=False,
                ),
            ),
        )
        result = await runner.run_result(messages=[], config=config)
        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
            messages = kwargs.get("messages") or []
            if messages and "recursive_validation_protocol" in messages[0]["content"]:
                return passed_validator_response(kwargs)
            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)
        config = await self.bind_snapshot(
            self.trace_id,
            RunConfig(
                trace_id=self.trace_id,
                max_iterations=4,
                knowledge=KnowledgeConfig(
                    enable_extraction=False,
                    enable_completion_extraction=False,
                    enable_injection=False,
                ),
            ),
        )
        result = await runner.run_result(messages=[], config=config)
        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
            messages = kwargs.get("messages") or []
            if messages and "recursive_validation_protocol" in messages[0]["content"]:
                return passed_validator_response(kwargs)
            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)
        config = await self.bind_snapshot(
            root_id,
            RunConfig(
                trace_id=root_id,
                max_iterations=3,
                knowledge=KnowledgeConfig(
                    enable_extraction=False,
                    enable_completion_extraction=False,
                    enable_injection=False,
                ),
            ),
        )
        result = await runner.run_result(messages=[], config=config)
        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)
        original_brief = TaskBrief.model_validate(state["task_brief"])
        revised_brief = original_brief.model_copy(
            update={"objective": "revised after the rewind cutoff"}
        )
        replace_task_brief(
            state,
            revised_brief,
            effective_at_sequence=8,
        )
        late_snapshot = create_context_snapshot(
            kind="reviewed_task_result",
            summary="granted after cutoff",
            source_trace_id="late-child",
            root_trace_id="root",
            uid="user-1",
            content={"result": "must be removed by rewind"},
            granted_at_sequence=8,
        )
        add_context_snapshot(
            trace.context,
            late_snapshot,
            root_task_anchor=ROOT_ANCHOR,
            task_brief=revised_brief,
        )
        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"])
        self.assertEqual(1, state["task_brief_version"])
        self.assertEqual(original_brief.model_dump(), state["task_brief"])
        self.assertEqual([], state["task_brief_history"])
        self.assertEqual({}, rewound.context["context_access"]["snapshots"])
        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_rewind_rebuilds_root_validation_history_and_cache(self):
        root_id = "root-validation-rewind"
        state = new_task_protocol()
        passed = validation_payload(root_id, outcome="passed")
        failed = validation_payload(root_id, outcome="failed")
        state["root_validation_history"] = [
            {**passed, "validated_at_sequence": 2},
            {**failed, "validated_at_sequence": 5},
        ]
        state["root_validation_attempts"] = 2
        state["root_validation_passed"] = False
        state["task_report_validation"] = {
            "plan_hash": failed["plan_hash"],
            "aggregate_result": failed,
            "scope_results": failed["scope_results"],
            "validated_at_sequence": 5,
        }
        await create_recursive_root(self.store, root_id, state=state)
        tree = GoalTree(
            mission="root",
            goals=[Goal(id="1", description="root", status="in_progress")],
            current_id="1",
        )
        await self.store.update_goal_tree(root_id, tree)
        for sequence in (1, 2, 3):
            await self.store.add_message(Message.create(
                trace_id=root_id,
                role="user",
                sequence=sequence,
                content=f"message-{sequence}",
            ))

        await AgentRunner(trace_store=self.store)._rewind(root_id, 3, tree)

        rewound = await self.store.get_trace(root_id)
        state = ensure_task_protocol(rewound.context)
        self.assertEqual(1, state["root_validation_attempts"])
        self.assertTrue(state["root_validation_passed"])
        self.assertEqual(
            ["passed"],
            [item["outcome"] for item in state["root_validation_history"]],
        )
        self.assertIsNone(state["task_report_validation"])

    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)
        config = await self.bind_snapshot(
            root_id,
            RunConfig(
                trace_id=root_id,
                max_iterations=5,
                knowledge=KnowledgeConfig(
                    enable_extraction=False,
                    enable_completion_extraction=False,
                    enable_injection=False,
                ),
            ),
        )
        result = await runner.run_result(messages=[], config=config)
        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": 3,
                "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", "read_context_ref"},
            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()