Agent/cyber_agent/tools/builtin/task_protocol.py

"""Tools for Recursive revision 2 task reporting and parent review."""

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 policy_from_context
from cyber_agent.core.context_policy import ContextPolicyError, normalize_task_brief
from cyber_agent.core.task_protocol import (
    TaskReport,
    TaskReportSubmission,
    TaskReview,
    ReviewDecision,
    TaskBrief,
    ensure_task_protocol,
)
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"},
    "error": {"REVISE_CHILD", "FAIL"},
}


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


@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]:
    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}")
    policy = policy_from_context(trace.context)
    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")

    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:
        submission = (
            TaskReportSubmission.model_validate_json(task_report)
            if isinstance(task_report, str)
            else TaskReportSubmission.model_validate(task_report)
        )
        report = TaskReport(
            child_trace_id=trace_id,
            **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_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]:
    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")
    policy = policy_from_context(parent.context)
    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:
        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,
            )
        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"
            )
        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,
            )
    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})"
                )

    pending.pop(child_trace_id)
    review_record = {
        **review.model_dump(),
        "pending_review": deepcopy(entry),
        "reviewed_at": datetime.now().isoformat(),
        "reviewed_at_sequence": context.get("sequence", 0),
    }
    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_validation"] = None
        child_state["protocol_correction_attempts"] = 0
        child_state["task_brief"] = revision_task_brief.model_dump()
        await store.update_trace(child_trace_id, context=child.context)

    if review.decision in {"ACCEPT_REFINE", "DESCEND_AGAIN", "REVISE_CHILD"}:
        action_task_brief = revision_task_brief or review.approved_next_task
        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": context.get("sequence", 0),
        })

    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(),
        "pending_review_count": len(pending),
        "next_action_count": len(state["next_actions"]),
    }