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.task_protocol import (
    TaskReport,
    TaskReportSubmission,
    TaskReview,
    ReviewDecision,
    TaskBrief,
    ensure_task_protocol,
)
from cyber_agent.tools import tool


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["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")

    revision_task_brief = None
    child_state = None
    try:
        if isinstance(approved_next_task, str):
            approved_next_task = TaskBrief.model_validate_json(approved_next_task)
        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,
        })
        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 = TaskBrief.model_validate(
                review.approved_next_task or child_state.get("task_brief")
            )
    except 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 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["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": entry.get("goal_id"),
            "task_brief": action_task_brief.model_dump() if action_task_brief else None,
            "created_at_sequence": context.get("sequence", 0),
        })

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

    goal_id = entry.get("goal_id")
    if goal_id:
        if pending:
            next_status = "pending_review"
        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

    return {
        "status": "completed",
        "task_review": review.model_dump(),
        "pending_review_count": len(pending),
        "next_action_count": len(state["next_actions"]),
    }