Agent/cyber_agent/core/task_protocol.py

"""Recursive revision 2 的父子任务协议与持久化状态。

Sub-Agent 使用 TaskBrief 下发任务，子 Agent 通过 TaskReport 上报，父 Agent 使用
TaskReview 审核。Runner 和两个协议工具共同把待审、后续动作和回退状态保存在 Trace context。
"""

from __future__ import annotations

from copy import deepcopy
from datetime import datetime
import json
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 RootTaskAnchor(StrictProtocolModel):
    """Recursive 根任务不可变的目标锚点。

    调用方创建根 Trace 时提供；Runner 固化后由每层孩子和 Validator读取，
    中间 Agent不能重新概括或覆盖它。
    """

    objective: str = Field(min_length=1)
    completion_criteria: list[str] = Field(min_length=1)
    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 ContextRef(StrictProtocolModel):
    """模型可在 TaskBrief 中转交的最小上下文引用句柄。

    来源、类型、摘要和内容由框架保管；模型只能选择已获授权的 ID与版本。
    """

    ref_id: str = Field(min_length=1)
    version: str = Field(pattern=r"^[0-9a-f]{64}$")


class TaskBrief(StrictProtocolModel):
    """父 Agent 下发给直接孩子的结构化任务说明。

    `agent()` 在创建/续跑子 Trace 前规范化，Runner 再将格式化内容作为子 Agent 的任务上下文。
    """

    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)
    context_refs: list[ContextRef] = 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):
    """子 Agent 在 TaskReport 中提交的执行者自检。

    `submit_task_report` 校验并持久化它；它不等于框架独立 Validator 产生的权威 ValidationResult。
    """

    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):
    """子 Agent 对直接父级的下一步建议，本身不会创建同级任务。

    父 Agent 在 `review_task_result` 中结合独立验收决定是否采纳、修订或上行。
    """

    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):
    """子 Agent 可提交的 TaskReport 负载，不允许模型填写 Trace ID。

    `submit_task_report` 接收该结构，检查结果与自检是否一致后由框架注入子 Trace ID。
    """

    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` 后的权威子任务报告。

    Sub-Agent 汇合孩子结果后将其与 ValidationResult 一起写入父 Trace 的待审队列。
    """

    child_trace_id: str = Field(min_length=1)


class TaskReview(StrictProtocolModel):
    """父 Agent 对一份直接孩子报告的结构化审核结果。

    `review_task_result` 根据 ValidationResult 的决策矩阵构建它，然后持久化修订、新任务或回退动作。
    """

    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]:
    """创建一份完整的 Recursive 任务协议初始状态。

    Runner 初始化根 Trace，Sub-Agent 创建子 Trace 时调用，状态随 Trace context 持久化。
    """
    brief = TaskBrief.model_validate(task_brief).model_dump() if task_brief else None
    return {
        "task_brief": brief,
        "task_brief_version": 1 if brief else 0,
        "task_brief_effective_at_sequence": 0,
        "task_brief_history": [],
        "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 replace_task_brief(
    state: dict[str, Any],
    task_brief: TaskBrief | dict[str, Any],
    *,
    effective_at_sequence: int,
) -> int:
    """以唯一入口替换当前 TaskBrief并保留旧版本。

    `review_task_result`批准修订后，`agent(continue_from=...)`真正恢复孩子时调用，
    避免审核阶段提前改写孩子，也避免覆盖后丢失任务合同的审计版本。
    """
    brief = TaskBrief.model_validate(task_brief).model_dump(mode="json")
    current = state.get("task_brief")
    current_version = state.get("task_brief_version", 0)
    if not isinstance(current_version, int) or current_version < 0:
        current_version = 0
    if current is not None:
        history_entry = {
            "version": current_version or 1,
            "task_brief": deepcopy(current),
            "effective_at_sequence": state.get(
                "task_brief_effective_at_sequence",
                0,
            ),
        }
        history = state.setdefault("task_brief_history", [])
        if history_entry not in history:
            history.append(history_entry)
    next_version = max(current_version, 0) + 1 if current is not None else 1
    state["task_brief"] = brief
    state["task_brief_version"] = next_version
    state["task_brief_effective_at_sequence"] = max(
        0,
        int(effective_at_sequence),
    )
    return next_version


def ensure_task_protocol(context: dict[str, Any]) -> dict[str, Any]:
    """获取 Trace context 中的协议状态，并幂等补齐缺失字段。

    Runner、Sub-Agent 和报告/审核工具在读写任务状态前统一调用此入口。
    """
    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]]:
    """在 Trace 回溯截断审核历史后，重建尚未解决的重规划请求。

    `AgentRunner._rewind` 清理截断点之后的协议记录时调用，避免丢失同批次的 `REPLAN_CURRENT`。
    """
    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:
    """在子 Agent 未能提交合法报告时生成框架级失败报告。

    Runner 的完成门禁和 Sub-Agent 结果汇合阶段调用，使父级仍能审核并恢复。
    """
    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:
    """为被协作式停止的子 Agent 生成 `failed` TaskReport。

    Sub-Agent 的停止汇合路径调用，然后交由父 Agent 选择修订孩子或失败结束。
    """
    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]:
    """组装父 Trace `pending_reviews` 中的一份待审记录。

    Sub-Agent 完成子 Trace 验收后调用，将 TaskReport、ValidationResult、原 Goal 和序列绑定。
    """
    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,
    *,
    available_context_refs: list[dict[str, Any]] | None = None,
) -> str:
    """将结构化 TaskBrief 格式化为子 Agent 可执行的任务文本。

    ``agent`` 完成规范化后调用，生成 Runner 作为 user message 传给子 Agent 的任务说明。
    """
    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(
            "\n## Context\n"
            + json.dumps(
                task_brief.context,
                ensure_ascii=False,
                sort_keys=True,
                separators=(",", ":"),
            )
        )
    if constraints:
        parts.append(f"\n## Constraints\n{constraints}")
    if available_context_refs:
        parts.append(
            "\n## Available Context References\n"
            + json.dumps(
                available_context_refs,
                ensure_ascii=False,
                sort_keys=True,
                separators=(",", ":"),
            )
        )
    parts.append(
        "\nSubmit a validated TaskReport with submit_task_report before finishing."
    )
    return "\n".join(parts)