Agent/cyber_agent/core/context_policy.py

"""Recursive 任务委托的最小上下文继承策略。

`agent()` 创建子任务和父 Agent 批准后续任务时共用本模块，只规范化 TaskBrief
并继承父级硬约束，不复制父级完整消息、GoalTree 或内部运行对象。
"""

from __future__ import annotations

import json
from typing import Any

from pydantic import ValidationError

from cyber_agent.core.task_protocol import TaskBrief


MAX_TASK_BRIEF_CHARS = 16_000


class ContextPolicyError(ValueError):
    """TaskBrief 无法安全穿过 Recursive 委托边界。

    `agent()` 或 `review_task_result` 捕获该异常后拒绝创建/批准不受控的子任务。
    """


def _stable_unique(items: list[str]) -> list[str]:
    return list(dict.fromkeys(items))


def normalize_task_brief(
    task_brief: TaskBrief | dict[str, Any] | str,
    *,
    parent_task_brief: TaskBrief | dict[str, Any] | None = None,
    max_chars: int = MAX_TASK_BRIEF_CHARS,
) -> TaskBrief:
    """校验并规范化 TaskBrief，仅向下继承父级硬约束。

    `agent()` 解析任务、`review_task_result` 批准下一步和续跑任务时调用，同时执行 JSON 与 16K 边界检查。
    """
    if max_chars <= 0:
        raise ContextPolicyError("max_chars must be greater than zero")
    try:
        raw_brief = (
            json.loads(task_brief)
            if isinstance(task_brief, str)
            else task_brief
        )
        brief = TaskBrief.model_validate(raw_brief)
        parent = (
            TaskBrief.model_validate(parent_task_brief)
            if parent_task_brief is not None
            else None
        )
    except (json.JSONDecodeError, TypeError, ValidationError) as exc:
        raise ContextPolicyError(f"Invalid TaskBrief: {exc}") from exc

    normalized = brief.model_copy(update={
        "completion_criteria": _stable_unique(brief.completion_criteria),
        "expected_outputs": _stable_unique(brief.expected_outputs),
        "parent_findings": _stable_unique(brief.parent_findings),
        "constraints": _stable_unique([
            *(parent.constraints if parent else []),
            *brief.constraints,
        ]),
    })
    try:
        encoded = json.dumps(
            normalized.model_dump(mode="json"),
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
            allow_nan=False,
        )
    except (TypeError, ValueError) as exc:
        raise ContextPolicyError(
            f"TaskBrief must be JSON serializable: {exc}"
        ) from exc
    if len(encoded) > max_chars:
        raise ContextPolicyError(
            f"TaskBrief exceeds the {max_chars}-character context limit"
        )
    return normalized


def task_briefs_match(
    actual: TaskBrief | dict[str, Any] | str,
    approved: TaskBrief | dict[str, Any] | str,
    *,
    parent_task_brief: TaskBrief | dict[str, Any] | None = None,
) -> bool:
    """用同一继承策略规范化后，精确比较实际与已批准 TaskBrief。

    Sub-Agent 在消费 `next_actions` 或执行 `REVISE_CHILD` 前调用，防止模型篡改父级审批内容。
    """
    actual_brief = normalize_task_brief(
        actual,
        parent_task_brief=parent_task_brief,
    )
    approved_brief = normalize_task_brief(
        approved,
        parent_task_brief=parent_task_brief,
    )
    return actual_brief.model_dump(mode="json") == approved_brief.model_dump(
        mode="json"
    )