Agent/cyber_agent/tools/builtin/task_protocol.py

"""Recursive 任务报告与父级审核工具。

子 Agent 通过 ``submit_task_report`` 提交结构化结果；Validator 验收并汇合后，
直接父 Agent 必须通过 ``review_task_result`` 批准下一步，两者均不向 Legacy 暴露。
"""

from __future__ import annotations

from copy import deepcopy
from datetime import datetime
from typing import Any

from pydantic import ValidationError

from cyber_agent.core.agent_mode import require_mutable_trace_policy
from cyber_agent.core.context_policy import (
    ContextPolicyError,
    add_context_snapshot,
    build_child_context_access,
    create_context_snapshot,
    normalize_task_brief,
    replace_context_access,
    require_matching_root_task_anchor,
    require_root_task_anchor,
)
from cyber_agent.core.task_protocol import (
    TaskReport,
    TaskReportSubmission,
    TaskReview,
    ReviewDecision,
    TaskBrief,
    TaskProgress,
    task_progress_readiness_error,
    ensure_task_protocol,
)
from cyber_agent.core.task_protocol_service import TaskProtocolStateError
from cyber_agent.core.validation import ValidationResult
from cyber_agent.tools import tool


_ALLOWED_REVIEW_DECISIONS_BY_VALIDATION: dict[str, set[ReviewDecision]] = {
    "passed": {
        "ACCEPT_REFINE",
        "REJECT_REFINE",
        "REVISE_CHILD",
        "DESCEND_AGAIN",
        "ASCEND",
        "FAIL",
    },
    "failed": {"REVISE_CHILD", "REPLAN_CURRENT", "FAIL"},
    "unknown": {"REVISE_CHILD", "REPLAN_CURRENT", "FAIL"},
    "error": {"REVISE_CHILD", "FAIL"},
}


def _error(message: str) -> dict[str, Any]:
    return {"status": "failed", "error": message}


@tool(
    description=(
        "Replace the current Recursive task's structured progress with one "
        "complete, version-checked snapshot."
    ),
    hidden_params=["context"],
    groups=["core"],
)
async def update_task_progress(
    expected_revision: int,
    progress: TaskProgress,
    context: dict | None = None,
) -> dict[str, Any]:
    """追加当前 Trace 的下一版 TaskProgress。"""
    if not context:
        return _error("context is required")
    if context.get("side_branch"):
        return _error("Side branches cannot update TaskProgress")
    service = context.get("task_protocol_service")
    trace_id = context.get("trace_id")
    if service is None or not trace_id:
        return _error("task_protocol_service and trace_id are required")
    try:
        revision = await service.update_progress(
            trace_id,
            expected_revision=expected_revision,
            progress=progress,
            effective_at_sequence=int(context.get("sequence", 0) or 0),
        )
    except (TaskProtocolStateError, ValueError) as exc:
        return _error(str(exc))
    return {
        "status": "completed",
        "task_progress": revision.model_dump(mode="json"),
    }


@tool(
    description="Submit the current Recursive child Agent's structured task report before finishing.",
    hidden_params=["context"],
    groups=["core"],
)
async def submit_task_report(
    task_report: TaskReportSubmission,
    context: dict | None = None,
) -> dict[str, Any]:
    """提交当前 Recursive 非根 Agent 的结构化任务报告。

    子 Agent 在结束前调用；Runner 会拒绝未提交报告或尚有待审核孩子的返回。
    """
    if not context:
        return _error("context is required")
    store = context.get("store")
    trace_id = context.get("trace_id")
    if not store or not trace_id:
        return _error("store and trace_id are required")

    trace = await store.get_trace(trace_id)
    if not trace:
        return _error(f"Trace not found: {trace_id}")
    try:
        policy = require_mutable_trace_policy(trace.context)
    except ValueError as exc:
        return _error(str(exc))
    if not policy.requires_task_protocol:
        return _error("submit_task_report is only available in Recursive revision 2")
    if not trace.parent_trace_id:
        return _error("Root Agent does not submit a TaskReport")

    root_trace_id = trace.context.get("root_trace_id") or trace.trace_id
    root = trace if root_trace_id == trace.trace_id else await store.get_trace(
        root_trace_id
    )
    if not root:
        return _error(f"Recursive root Trace not found: {root_trace_id}")
    try:
        require_matching_root_task_anchor(root.context, trace.context)
    except ContextPolicyError as exc:
        return _error(str(exc))

    state = ensure_task_protocol(trace.context)
    if state["pending_reviews"]:
        return _error("Review all child TaskReports before submitting this TaskReport")
    if state["next_actions"]:
        return _error("Execute every approved next action before submitting this TaskReport")
    if state["task_report"] is not None:
        return _error("A TaskReport has already been submitted")

    try:
        parsed_submission = (
            TaskReportSubmission.model_validate_json(task_report)
            if isinstance(task_report, str)
            else TaskReportSubmission.model_validate(task_report)
        )
    except ValidationError as exc:
        return _error(f"Invalid TaskReport: {exc}")
    if policy.requires_task_progress and parsed_submission.outcome == "satisfied":
        readiness_error = task_progress_readiness_error(state)
        if readiness_error:
            return _error(readiness_error)

    try:
        report = TaskReport(
            child_trace_id=trace_id,
            **parsed_submission.model_dump(),
        )
    except ValidationError as exc:
        return _error(f"Invalid TaskReport: {exc}")

    state["task_report"] = report.model_dump()
    state["task_report_submitted_at_sequence"] = context.get("sequence", 0)
    state["task_report_progress_revision"] = state.get(
        "task_progress_head_revision"
    )
    state["task_report_validation"] = None
    state["protocol_correction_attempts"] = 0
    await store.update_trace(trace_id, context=trace.context)
    goal_id = context.get("goal_id")
    if goal_id:
        goal_status = (
            "failed"
            if report.outcome in {"failed", "protocol_error"}
            else "completed"
        )
        await store.update_goal(
            trace_id,
            goal_id,
            cascade_completion=False,
            status=goal_status,
            summary=report.summary,
        )
        tree = context.get("goal_tree")
        goal = tree.find(goal_id) if tree else None
        if goal:
            goal.status = goal_status
            goal.summary = report.summary
    return {"status": "completed", "task_report": report.model_dump()}


@tool(
    description="Review one pending direct child's TaskReport and approve the parent's next action.",
    hidden_params=["context"],
    groups=["core"],
)
async def review_task_result(
    child_trace_id: str,
    decision: ReviewDecision,
    reason: str,
    approved_next_task: TaskBrief | None = None,
    context: dict | None = None,
) -> dict[str, Any]:
    """审核一份直属孩子的待处理 ``TaskReport``。

    父 Agent 在 ``pending_review`` 阶段调用；决策受 Validator 结果约束并持久化下一动作。
    """
    if not context:
        return _error("context is required")
    store = context.get("store")
    parent_trace_id = context.get("trace_id")
    if not store or not parent_trace_id:
        return _error("store and trace_id are required")

    parent = await store.get_trace(parent_trace_id)
    child = await store.get_trace(child_trace_id)
    if not parent or not child:
        return _error("Parent or child Trace not found")
    try:
        policy = require_mutable_trace_policy(parent.context)
    except ValueError as exc:
        return _error(str(exc))
    if not policy.requires_task_protocol:
        return _error("review_task_result is only available in Recursive revision 2")
    if child.parent_trace_id != parent_trace_id or child.uid != parent.uid:
        return _error("TaskReport must belong to a direct child with the same owner")

    state = ensure_task_protocol(parent.context)
    pending = state["pending_reviews"]
    entry = pending.get(child_trace_id)
    if not entry:
        return _error("This direct child has no pending TaskReport")
    origin_goal_id = child.parent_goal_id
    if entry.get("goal_id") != origin_goal_id:
        return _error(
            "Pending TaskReport goal_id does not match the child's originating Goal"
        )

    revision_task_brief = None
    child_state = None
    try:
        root_trace_id = parent.context.get("root_trace_id") or parent.trace_id
        root = parent if root_trace_id == parent.trace_id else await store.get_trace(
            root_trace_id
        )
        if not root:
            return _error(f"Recursive root Trace not found: {root_trace_id}")
        root_task_anchor = require_matching_root_task_anchor(
            root.context,
            parent.context,
        )
        require_matching_root_task_anchor(root.context, child.context)
        parent_task_brief = state.get("task_brief")
        if isinstance(approved_next_task, str):
            approved_next_task = TaskBrief.model_validate_json(approved_next_task)
        if approved_next_task is not None:
            approved_next_task = normalize_task_brief(
                approved_next_task,
                parent_task_brief=parent_task_brief,
                root_task_anchor=root_task_anchor,
            )
        validation_result = ValidationResult.model_validate(
            entry.get("validation_result")
        )
        if validation_result.evaluated_trace_id != child_trace_id:
            return _error(
                "ValidationResult evaluated_trace_id does not match the direct child"
            )
        child_validation_state = ensure_task_protocol(child.context).get(
            "task_report_validation"
        )
        if not isinstance(child_validation_state, dict):
            return _error("Child has no persisted framework validation cache")
        if (
            child_validation_state.get("plan_hash") != validation_result.plan_hash
            or child_validation_state.get("aggregate_result")
            != validation_result.model_dump(mode="json")
        ):
            return _error(
                "Pending ValidationResult does not match the child's authoritative cache"
            )
        allowed_decisions = _ALLOWED_REVIEW_DECISIONS_BY_VALIDATION[
            validation_result.outcome
        ]
        if decision not in allowed_decisions:
            return _error(
                f"Validation {validation_result.outcome} only allows: "
                f"{', '.join(sorted(allowed_decisions))}"
            )
        review = TaskReview.model_validate({
            "parent_trace_id": parent_trace_id,
            "child_trace_id": child_trace_id,
            "decision": decision,
            "reason": reason,
            "approved_next_task": approved_next_task,
            "retry_from": (
                validation_result.retry_from
                if decision == "REPLAN_CURRENT"
                else None
            ),
        })
        report = TaskReport.model_validate(entry["task_report"])
        if report.child_trace_id != child_trace_id:
            return _error(
                "Pending TaskReport child_trace_id does not match the direct child"
            )
        if review.decision == "REVISE_CHILD":
            child_state = ensure_task_protocol(child.context)
            revision_task_brief = normalize_task_brief(
                review.approved_next_task or child_state.get("task_brief"),
                parent_task_brief=parent_task_brief,
                root_task_anchor=root_task_anchor,
            )
    except (ContextPolicyError, ValidationError) as exc:
        return _error(f"Invalid TaskReview: {exc}")

    if report.outcome in {"failed", "protocol_error"} and review.decision not in {
        "REVISE_CHILD",
        "FAIL",
    }:
        return _error("Failed/protocol_error reports only allow REVISE_CHILD or FAIL")
    if review.decision == "REPLAN_CURRENT":
        goal_id = origin_goal_id
        persisted_tree = await store.get_goal_tree(parent_trace_id)
        if not goal_id or not persisted_tree or not persisted_tree.find(goal_id):
            return _error(
                "REPLAN_CURRENT requires the direct parent's originating Goal"
            )
    if review.decision in {"ASCEND", "FAIL"} and len(pending) != 1:
        return _error("ASCEND/FAIL is only allowed for the final pending report")
    if review.decision in {"ASCEND", "FAIL"} and state["next_actions"]:
        return _error("ASCEND/FAIL requires every approved next action to finish first")
    if review.decision in {"ACCEPT_REFINE", "DESCEND_AGAIN"}:
        child_limit = policy.max_children_per_parent
        if child_limit is not None:
            existing_children = await store.list_traces(
                parent_trace_id=parent_trace_id,
                created_by_tool="agent",
                limit=child_limit + 1,
            )
            reserved_children = sum(
                action.get("decision") in {"ACCEPT_REFINE", "DESCEND_AGAIN"}
                for action in state["next_actions"]
            )
            if len(existing_children) + reserved_children + 1 > child_limit:
                return _error(
                    "Cannot approve another child: local child Agent limit "
                    f"would be exceeded (existing={len(existing_children)}, "
                    f"reserved={reserved_children}, max={child_limit})"
                )

    action_task_brief = revision_task_brief or review.approved_next_task
    if action_task_brief is not None:
        try:
            build_child_context_access(
                parent_context=parent.context,
                parent_trace_id=parent_trace_id,
                root_trace_id=root_trace_id,
                uid=parent.uid,
                parent_task_state=state,
                child_task_brief=action_task_brief,
                root_task_anchor=root_task_anchor,
            )
        except ContextPolicyError as exc:
            return _error(f"Approved TaskBrief context is invalid: {exc}")

    reviewed_at_sequence = int(context.get("sequence", 0) or 0)
    context_ref = None
    context_ref_warning = None
    try:
        snapshot = create_context_snapshot(
            kind="reviewed_task_result",
            summary=review.reason,
            source_trace_id=child_trace_id,
            root_trace_id=root_trace_id,
            uid=parent.uid,
            content={
                "task_report": deepcopy(entry["task_report"]),
                "validation_result": deepcopy(entry["validation_result"]),
                "task_review": review.model_dump(mode="json"),
            },
            granted_at_sequence=reviewed_at_sequence,
        )
        context_ref = add_context_snapshot(
            parent.context,
            snapshot,
            root_task_anchor=root_task_anchor,
            task_brief=state.get("task_brief"),
        )
    except ContextPolicyError as exc:
        context_ref_warning = str(exc)

    if context_ref is not None and action_task_brief is not None:
        try:
            action_task_brief = normalize_task_brief(
                action_task_brief.model_copy(update={
                    "context_refs": [
                        *action_task_brief.context_refs,
                        context_ref,
                    ],
                }),
                parent_task_brief=parent_task_brief,
                root_task_anchor=root_task_anchor,
            )
            build_child_context_access(
                parent_context=parent.context,
                parent_trace_id=parent_trace_id,
                root_trace_id=root_trace_id,
                uid=parent.uid,
                parent_task_state=state,
                child_task_brief=action_task_brief,
                root_task_anchor=root_task_anchor,
            )
        except ContextPolicyError as exc:
            action_task_brief = revision_task_brief or review.approved_next_task
            context_ref_warning = (
                f"{context_ref_warning}; {exc}"
                if context_ref_warning
                else str(exc)
            )

    pending.pop(child_trace_id)
    review_record = {
        **review.model_dump(),
        "pending_review": deepcopy(entry),
        "reviewed_at": datetime.now().isoformat(),
        "reviewed_at_sequence": reviewed_at_sequence,
        "context_ref": context_ref.model_dump() if context_ref else None,
        "context_ref_warning": context_ref_warning,
    }
    state["reviews"].append(review_record)
    state["protocol_correction_attempts"] = 0

    if review.decision == "REVISE_CHILD":
        assert child_state is not None and revision_task_brief is not None
        old_report = child_state.get("task_report")
        if old_report:
            child_state["report_history"].append(old_report)
        child_state["task_report"] = None
        child_state["task_report_submitted_at_sequence"] = None
        child_state["task_report_progress_revision"] = None
        child_state["task_report_validation"] = None
        child_state["protocol_correction_attempts"] = 0
        replace_context_access(
            child.context,
            [],
            root_task_anchor=require_root_task_anchor(child.context),
            task_brief=child_state.get("task_brief"),
        )
        await store.update_trace(child_trace_id, context=child.context)

    if review.decision in {"ACCEPT_REFINE", "DESCEND_AGAIN", "REVISE_CHILD"}:
        state["next_actions"].append({
            "decision": review.decision,
            "child_trace_id": child_trace_id,
            "goal_id": origin_goal_id,
            "task_brief": action_task_brief.model_dump() if action_task_brief else None,
            "created_at_sequence": reviewed_at_sequence,
        })

    resolved_replan = None
    if review.decision == "REPLAN_CURRENT":
        state["pending_replans"].append({
            "child_trace_id": child_trace_id,
            "goal_id": origin_goal_id,
            "retry_from": review.retry_from,
            "requested_at_sequence": context.get("sequence", 0),
            "received_at_sequence": entry.get("received_at_sequence"),
        })
    elif review.decision == "FAIL":
        state["pending_replans"].clear()
    if not pending and state["pending_replans"]:
        resolved_replan = state["pending_replans"][0]
        state["pending_replans"].clear()

    await store.update_trace(parent_trace_id, context=parent.context)

    goal_id = origin_goal_id
    if goal_id:
        if pending:
            next_status = "pending_review"
        elif resolved_replan:
            next_status = "in_progress"
            goal_id = resolved_replan["goal_id"]
        elif review.decision == "ASCEND":
            next_status = "completed"
        elif review.decision == "FAIL":
            next_status = "failed"
        else:
            next_status = "in_progress"
        await store.update_goal(
            parent_trace_id,
            goal_id,
            cascade_completion=False,
            status=next_status,
            summary=reason if next_status in {"completed", "failed"} else None,
        )
        tree = context.get("goal_tree")
        goal = tree.find(goal_id) if tree else None
        if goal:
            goal.status = next_status
            if next_status in {"completed", "failed"}:
                goal.summary = reason
                if tree.current_id == goal_id:
                    tree.current_id = goal.parent_id
            elif resolved_replan:
                tree.current_id = goal_id

        if resolved_replan:
            persisted_tree = await store.get_goal_tree(parent_trace_id)
            persisted_goal = persisted_tree.find(goal_id) if persisted_tree else None
            if persisted_goal:
                persisted_goal.status = "in_progress"
                persisted_tree.current_id = goal_id
                await store.update_goal_tree(parent_trace_id, persisted_tree)

    return {
        "status": "completed",
        "task_review": review.model_dump(),
        "context_ref": context_ref.model_dump() if context_ref else None,
        "context_ref_warning": context_ref_warning,
        "next_action_task_brief": (
            action_task_brief.model_dump(mode="json")
            if action_task_brief is not None
            else None
        ),
        "pending_review_count": len(pending),
        "next_action_count": len(state["next_actions"]),
    }