Agent/cyber_agent/core/task_protocol.py

"""Recursive revision 2 parent/child task protocol."""

from __future__ import annotations

from copy import deepcopy
from datetime import datetime
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

from cyber_agent.core.validation import RetryFrom


TaskOutcome = Literal["satisfied", "partial", "failed", "protocol_error"]
NextStepDirection = Literal["DESCEND", "REFINE", "ASCEND", "NONE"]
ReviewDecision = Literal[
    "ACCEPT_REFINE",
    "REJECT_REFINE",
    "REVISE_CHILD",
    "DESCEND_AGAIN",
    "REPLAN_CURRENT",
    "ASCEND",
    "FAIL",
]


class StrictProtocolModel(BaseModel):
    model_config = ConfigDict(extra="forbid", str_strip_whitespace=True)


class TaskBrief(StrictProtocolModel):
    objective: str = Field(min_length=1)
    reason: str = Field(min_length=1)
    completion_criteria: list[str] = Field(min_length=1)
    expected_outputs: list[str] = Field(min_length=1)
    parent_findings: list[str] = Field(default_factory=list)
    context: dict[str, Any] = Field(default_factory=dict)
    constraints: list[str] = Field(default_factory=list)

    @field_validator(
        "completion_criteria",
        "expected_outputs",
        "parent_findings",
        "constraints",
    )
    @classmethod
    def validate_non_empty_items(cls, items: list[str]) -> list[str]:
        if any(not item.strip() for item in items):
            raise ValueError("items must be non-empty strings")
        return items


class Validation(StrictProtocolModel):
    hard_passed: bool
    open_issues: list[str] = Field(default_factory=list)

    @field_validator("open_issues")
    @classmethod
    def validate_open_issues(cls, items: list[str]) -> list[str]:
        if any(not item.strip() for item in items):
            raise ValueError("open_issues must contain non-empty strings")
        return items


class NextStepSuggestion(StrictProtocolModel):
    direction: NextStepDirection
    reason: str = Field(min_length=1)
    suggested_next_task: TaskBrief | None = None

    @model_validator(mode="after")
    def validate_suggested_task(self) -> "NextStepSuggestion":
        if self.direction in {"DESCEND", "REFINE"} and not self.suggested_next_task:
            raise ValueError("DESCEND/REFINE requires suggested_next_task")
        if self.direction in {"ASCEND", "NONE"} and self.suggested_next_task:
            raise ValueError("ASCEND/NONE does not accept suggested_next_task")
        return self


class TaskReportSubmission(StrictProtocolModel):
    summary: str = Field(min_length=1)
    outcome: TaskOutcome
    validation: Validation
    next_step_suggestion: NextStepSuggestion
    outputs: list[dict[str, Any]] = Field(default_factory=list)
    evidence: list[dict[str, Any]] = Field(default_factory=list)
    remaining_issues: list[str] = Field(default_factory=list)

    @field_validator("remaining_issues")
    @classmethod
    def validate_remaining_issues(cls, items: list[str]) -> list[str]:
        if any(not item.strip() for item in items):
            raise ValueError("remaining_issues must contain non-empty strings")
        return items

    @model_validator(mode="after")
    def validate_outcome_consistency(self) -> "TaskReportSubmission":
        if self.outcome == "satisfied" and (
            not self.validation.hard_passed
            or self.validation.open_issues
            or self.remaining_issues
        ):
            raise ValueError(
                "satisfied requires hard_passed with no open or remaining issues"
            )
        if self.outcome in {"failed", "protocol_error"} and self.validation.hard_passed:
            raise ValueError(f"{self.outcome} cannot have hard_passed=true")
        return self


class TaskReport(TaskReportSubmission):
    child_trace_id: str = Field(min_length=1)


class TaskReview(StrictProtocolModel):
    parent_trace_id: str = Field(min_length=1)
    child_trace_id: str = Field(min_length=1)
    decision: ReviewDecision
    reason: str = Field(min_length=1)
    approved_next_task: TaskBrief | None = None
    retry_from: RetryFrom | None = None

    @model_validator(mode="after")
    def validate_approved_next_task(self) -> "TaskReview":
        if self.decision in {"ACCEPT_REFINE", "DESCEND_AGAIN"}:
            if not self.approved_next_task:
                raise ValueError(f"{self.decision} requires approved_next_task")
        elif self.decision != "REVISE_CHILD" and self.approved_next_task:
            raise ValueError(
                f"{self.decision} does not accept approved_next_task"
            )
        if self.decision == "REPLAN_CURRENT":
            if self.retry_from is None:
                raise ValueError("REPLAN_CURRENT requires retry_from")
        elif self.retry_from is not None:
            raise ValueError(f"{self.decision} does not accept retry_from")
        return self


def new_task_protocol(task_brief: TaskBrief | dict | None = None) -> dict[str, Any]:
    brief = TaskBrief.model_validate(task_brief).model_dump() if task_brief else None
    return {
        "task_brief": brief,
        "task_report": None,
        "task_report_submitted_at_sequence": None,
        "task_report_validation": None,
        "report_history": [],
        "pending_reviews": {},
        "reviews": [],
        "next_actions": [],
        "pending_replans": [],
        "root_validation_attempts": 0,
        "root_validation_history": [],
        "root_validation_passed": False,
        "protocol_correction_attempts": 0,
    }


def ensure_task_protocol(context: dict[str, Any]) -> dict[str, Any]:
    state = context.get("task_protocol")
    if not isinstance(state, dict):
        state = new_task_protocol()
        context["task_protocol"] = state
    defaults = new_task_protocol()
    for key, value in defaults.items():
        state.setdefault(key, deepcopy(value))
    return state


def rebuild_pending_replans(state: dict[str, Any]) -> list[dict[str, Any]]:
    """Rebuild unresolved replans after review history is rewound."""
    pending_batches = {
        entry.get("received_at_sequence")
        for entry in state.get("pending_reviews", {}).values()
    }
    pending_batches.discard(None)
    if not pending_batches:
        return []
    replans = []
    for review in state.get("reviews", []):
        pending_entry = review.get("pending_review", {})
        if (
            review.get("decision") == "REPLAN_CURRENT"
            and pending_entry.get("received_at_sequence") in pending_batches
        ):
            replans.append({
                "child_trace_id": review.get("child_trace_id"),
                "goal_id": pending_entry.get("goal_id"),
                "retry_from": review.get("retry_from"),
                "requested_at_sequence": review.get("reviewed_at_sequence", 0),
                "received_at_sequence": pending_entry.get(
                    "received_at_sequence"
                ),
            })
    return replans


def protocol_error_report(child_trace_id: str, reason: str) -> TaskReport:
    return TaskReport(
        child_trace_id=child_trace_id,
        summary="Task protocol failed before a valid report was submitted.",
        outcome="protocol_error",
        validation=Validation(hard_passed=False, open_issues=[reason]),
        next_step_suggestion=NextStepSuggestion(
            direction="ASCEND",
            reason="Parent review is required to revise or fail this task.",
        ),
        remaining_issues=[reason],
    )


def stopped_task_report(child_trace_id: str, reason: str) -> TaskReport:
    return TaskReport(
        child_trace_id=child_trace_id,
        summary="Child Agent execution was stopped before completion.",
        outcome="failed",
        validation=Validation(hard_passed=False, open_issues=[reason]),
        next_step_suggestion=NextStepSuggestion(
            direction="NONE",
            reason="Parent review must revise the stopped child or fail the task.",
        ),
        remaining_issues=[reason],
    )


def pending_review_entry(
    *,
    goal_id: str | None,
    report: TaskReport,
    validation_result: dict[str, Any],
    received_at_sequence: int,
) -> dict[str, Any]:
    return {
        "goal_id": goal_id,
        "task_report": report.model_dump(),
        "validation_result": validation_result,
        "received_at_sequence": received_at_sequence,
        "received_at": datetime.now().isoformat(),
    }


def format_task_brief(task_brief: TaskBrief) -> str:
    criteria = "\n".join(
        f"{index}. {item}"
        for index, item in enumerate(task_brief.completion_criteria, start=1)
    )
    constraints = "\n".join(f"- {item}" for item in task_brief.constraints)
    expected_outputs = "\n".join(
        f"- {item}" for item in task_brief.expected_outputs
    )
    parent_findings = "\n".join(
        f"- {item}" for item in task_brief.parent_findings
    )
    parts = [
        "# Task Brief",
        f"\n## Objective\n{task_brief.objective}",
        f"\n## Reason\n{task_brief.reason}",
        f"\n## Completion Criteria\n{criteria}",
        f"\n## Expected Outputs\n{expected_outputs}",
    ]
    if parent_findings:
        parts.append(f"\n## Parent Findings\n{parent_findings}")
    if task_brief.context:
        parts.append(f"\n## Context\n{task_brief.context}")
    if constraints:
        parts.append(f"\n## Constraints\n{constraints}")
    parts.append(
        "\nSubmit a validated TaskReport with submit_task_report before finishing."
    )
    return "\n".join(parts)