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


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


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


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

    @field_validator("completion_criteria", "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

    @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"
            )
        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,
        "report_history": [],
        "pending_reviews": {},
        "reviews": [],
        "next_actions": [],
        "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 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 pending_review_entry(
    *,
    goal_id: str | None,
    report: TaskReport,
    received_at_sequence: int,
) -> dict[str, Any]:
    return {
        "goal_id": goal_id,
        "task_report": report.model_dump(),
        "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)
    parts = [
        "# Task Brief",
        f"\n## Objective\n{task_brief.objective}",
        f"\n## Completion Criteria\n{criteria}",
    ]
    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)