Agent/cyber_agent/core/context_policy.py

"""Recursive 任务上下文的规范化、版本引用与逐层授权策略。

Runner、任务协议工具和 ``agent`` 都只调用这里的纯函数：根目标保持不变，
TaskBrief 只继承硬约束，完整协议内容则通过当前 Trace 的不可变本地快照按需读取。
"""

from __future__ import annotations

from copy import deepcopy
from hashlib import sha256
import json
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, ValidationError

from cyber_agent.core.task_protocol import (
    ContextRef,
    RootTaskAnchor,
    TaskBrief,
    current_task_progress,
)


MAX_ROOT_TASK_ANCHOR_CHARS = 4_000
MAX_TASK_BRIEF_CHARS = 16_000
MAX_CONTEXT_REFS = 8
MAX_CONTEXT_SNAPSHOT_CHARS = 16_000
MAX_CONTEXT_SNAPSHOTS_CHARS = 64_000

ROOT_TASK_ANCHOR_KEY = "root_task_anchor"
ROOT_TASK_ANCHOR_HASH_KEY = "root_task_anchor_hash"
CONTEXT_ACCESS_KEY = "context_access"

ContextSnapshotKind = Literal[
    "task_brief",
    "reviewed_task_result",
    "application_resource",
]


class ContextPolicyError(ValueError):
    """Recursive 上下文无法安全通过委托、审核或读取边界。"""


class ContextSnapshot(BaseModel):
    """当前 Trace持有的一份框架生成、内容寻址的只读协议快照。"""

    model_config = ConfigDict(extra="forbid", str_strip_whitespace=True)

    ref_id: str = Field(pattern=r"^ctx_[0-9a-f]{24}$")
    version: str = Field(pattern=r"^[0-9a-f]{64}$")
    kind: ContextSnapshotKind
    summary: str = Field(min_length=1, max_length=500)
    source_trace_id: str = Field(min_length=1)
    root_trace_id: str = Field(min_length=1)
    uid: str | None = None
    content: dict[str, Any]
    granted_at_sequence: int = Field(ge=0)

    def to_ref(self) -> ContextRef:
        return ContextRef(ref_id=self.ref_id, version=self.version)


def canonical_json(value: Any) -> str:
    """用唯一 JSON 形态计算大小、哈希和协议精确匹配。"""
    try:
        return json.dumps(
            value,
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
            allow_nan=False,
        )
    except (TypeError, ValueError) as exc:
        raise ContextPolicyError(f"Context must be JSON serializable: {exc}") from exc


def context_chars(value: Any) -> int:
    return len(canonical_json(value))


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


def _stable_unique_refs(items: list[ContextRef]) -> list[ContextRef]:
    return list({(item.ref_id, item.version): item for item in items}.values())


def normalize_root_task_anchor(
    anchor: RootTaskAnchor | dict[str, Any] | str,
    *,
    max_chars: int = MAX_ROOT_TASK_ANCHOR_CHARS,
) -> RootTaskAnchor:
    """校验并规范化调用方提供的根目标锚点。"""
    if max_chars <= 0:
        raise ContextPolicyError("max_chars must be greater than zero")
    try:
        raw = json.loads(anchor) if isinstance(anchor, str) else anchor
        parsed = RootTaskAnchor.model_validate(raw)
    except (json.JSONDecodeError, TypeError, ValidationError) as exc:
        raise ContextPolicyError(f"Invalid RootTaskAnchor: {exc}") from exc
    normalized = parsed.model_copy(update={
        "completion_criteria": _stable_unique(parsed.completion_criteria),
        "constraints": _stable_unique(parsed.constraints),
    })
    try:
        encoded_chars = context_chars(normalized.model_dump(mode="json"))
    except Exception as exc:
        raise ContextPolicyError(
            f"RootTaskAnchor must be JSON serializable: {exc}"
        ) from exc
    if encoded_chars > max_chars:
        raise ContextPolicyError(
            f"RootTaskAnchor exceeds the {max_chars}-character context limit"
        )
    return normalized


def root_task_anchor_hash(anchor: RootTaskAnchor | dict[str, Any]) -> str:
    normalized = normalize_root_task_anchor(anchor)
    return sha256(
        canonical_json(normalized.model_dump(mode="json")).encode("utf-8")
    ).hexdigest()


def persist_root_task_anchor(
    context: dict[str, Any],
    anchor: RootTaskAnchor | dict[str, Any] | str,
) -> RootTaskAnchor:
    """把唯一规范化 Anchor和哈希写入 Trace context。"""
    normalized = normalize_root_task_anchor(anchor)
    context[ROOT_TASK_ANCHOR_KEY] = normalized.model_dump(mode="json")
    context[ROOT_TASK_ANCHOR_HASH_KEY] = root_task_anchor_hash(normalized)
    return normalized


def require_root_task_anchor(context: dict[str, Any]) -> RootTaskAnchor:
    """读取并校验持久化 Anchor；缺失或被篡改时失败关闭。"""
    raw = context.get(ROOT_TASK_ANCHOR_KEY)
    if raw is None:
        raise ContextPolicyError(
            "This Recursive trace predates root_task_anchor; create a new trace"
        )
    normalized = normalize_root_task_anchor(raw)
    expected = root_task_anchor_hash(normalized)
    if context.get(ROOT_TASK_ANCHOR_HASH_KEY) != expected:
        raise ContextPolicyError("Persisted RootTaskAnchor hash does not match its content")
    return normalized


def require_matching_root_task_anchor(
    root_context: dict[str, Any],
    current_context: dict[str, Any],
) -> RootTaskAnchor:
    """从根 Trace读取权威 Anchor，并确认当前节点副本完全一致。

    Runner续跑、Sub-Agent创建和父级审核共用，避免某一入口只验证
    当前 Trace自身哈希，却接受与任务树根不同的自洽副本。
    """
    root = require_root_task_anchor(root_context)
    current = require_root_task_anchor(current_context)
    if (
        current_context.get(ROOT_TASK_ANCHOR_HASH_KEY)
        != root_context.get(ROOT_TASK_ANCHOR_HASH_KEY)
        or current.model_dump(mode="json") != root.model_dump(mode="json")
    ):
        raise ContextPolicyError(
            "Persisted RootTaskAnchor does not match the root Trace"
        )
    return root


def normalize_task_brief(
    task_brief: TaskBrief | dict[str, Any] | str,
    *,
    parent_task_brief: TaskBrief | dict[str, Any] | None = None,
    root_task_anchor: RootTaskAnchor | dict[str, Any] | None = None,
    max_chars: int = MAX_TASK_BRIEF_CHARS,
) -> TaskBrief:
    """规范化 TaskBrief，只继承根与直接父级的硬约束。"""
    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
        )
        anchor = (
            normalize_root_task_anchor(root_task_anchor)
            if root_task_anchor 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([
            *(anchor.constraints if anchor else []),
            *(parent.constraints if parent else []),
            *brief.constraints,
        ]),
        "context_refs": _stable_unique_refs(brief.context_refs),
        "validation_scopes": [
            scope
            for scope in ("evidence", "hypothesis", "output", "task")
            if scope in set(brief.validation_scopes)
        ],
    })
    try:
        encoded_chars = context_chars(normalized.model_dump(mode="json"))
    except Exception as exc:
        raise ContextPolicyError(
            f"TaskBrief must be JSON serializable: {exc}"
        ) from exc
    if encoded_chars > 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,
    root_task_anchor: RootTaskAnchor | dict[str, Any] | None = None,
) -> bool:
    """以同一上下文策略规范化后精确比较实际和批准的 TaskBrief。"""
    actual_brief = normalize_task_brief(
        actual,
        parent_task_brief=parent_task_brief,
        root_task_anchor=root_task_anchor,
    )
    approved_brief = normalize_task_brief(
        approved,
        parent_task_brief=parent_task_brief,
        root_task_anchor=root_task_anchor,
    )
    return actual_brief.model_dump(mode="json") == approved_brief.model_dump(
        mode="json"
    )


def _context_ref_id(
    kind: ContextSnapshotKind,
    source_trace_id: str,
    version: str,
) -> str:
    seed = f"{kind}\0{source_trace_id}\0{version}".encode("utf-8")
    return f"ctx_{sha256(seed).hexdigest()[:24]}"


def create_context_snapshot(
    *,
    kind: ContextSnapshotKind,
    summary: str,
    source_trace_id: str,
    root_trace_id: str,
    uid: str | None,
    content: dict[str, Any],
    granted_at_sequence: int,
) -> ContextSnapshot:
    """为 TaskBrief或审核结论创建内容寻址的不可变快照。"""
    encoded = canonical_json(content)
    if len(encoded) > MAX_CONTEXT_SNAPSHOT_CHARS:
        raise ContextPolicyError(
            f"Context snapshot exceeds the {MAX_CONTEXT_SNAPSHOT_CHARS}-character limit"
        )
    version = sha256(encoded.encode("utf-8")).hexdigest()
    return ContextSnapshot(
        ref_id=_context_ref_id(kind, source_trace_id, version),
        version=version,
        kind=kind,
        summary=summary.strip()[:500] or kind,
        source_trace_id=source_trace_id,
        root_trace_id=root_trace_id,
        uid=uid,
        content=deepcopy(content),
        granted_at_sequence=granted_at_sequence,
    )


def _validated_snapshot(raw: ContextSnapshot | dict[str, Any]) -> ContextSnapshot:
    try:
        snapshot = ContextSnapshot.model_validate(raw)
    except ValidationError as exc:
        raise ContextPolicyError(f"Invalid context snapshot: {exc}") from exc
    encoded = canonical_json(snapshot.content)
    if sha256(encoded.encode("utf-8")).hexdigest() != snapshot.version:
        raise ContextPolicyError("Context snapshot version does not match its content")
    if snapshot.ref_id != _context_ref_id(
        snapshot.kind,
        snapshot.source_trace_id,
        snapshot.version,
    ):
        raise ContextPolicyError("Context snapshot ref_id does not match its source")
    return snapshot


def ensure_context_access(context: dict[str, Any]) -> dict[str, Any]:
    """幂等获取当前 Trace本地授权目录。

    缺失时可初始化；已存在但结构被破坏时失败关闭，避免读取工具
    把非法授权目录当成空目录静默修复。
    """
    access = context.get(CONTEXT_ACCESS_KEY)
    if access is None:
        access = {}
        context[CONTEXT_ACCESS_KEY] = access
    elif not isinstance(access, dict):
        raise ContextPolicyError("Invalid context_access directory")
    snapshots = access.setdefault("snapshots", {})
    metrics = access.setdefault("metrics", {})
    if not isinstance(snapshots, dict) or not isinstance(metrics, dict):
        raise ContextPolicyError("Invalid context_access directory")
    return access


def _metrics(
    *,
    root_task_anchor: RootTaskAnchor | dict[str, Any],
    task_brief: TaskBrief | dict[str, Any] | None,
    snapshots: list[ContextSnapshot],
) -> dict[str, int]:
    return {
        "root_anchor_chars": context_chars(
            RootTaskAnchor.model_validate(root_task_anchor).model_dump(mode="json")
        ),
        "task_brief_chars": (
            context_chars(TaskBrief.model_validate(task_brief).model_dump(mode="json"))
            if task_brief is not None
            else 0
        ),
        "authorized_ref_count": len(snapshots),
        "snapshot_chars": sum(context_chars(item.content) for item in snapshots),
    }


def replace_context_access(
    context: dict[str, Any],
    snapshots: list[ContextSnapshot | dict[str, Any]],
    *,
    root_task_anchor: RootTaskAnchor | dict[str, Any],
    task_brief: TaskBrief | dict[str, Any] | None,
) -> None:
    """原子替换当前 Trace的全部本地授权快照和静态用量。"""
    normalized = [_validated_snapshot(item) for item in snapshots]
    unique = {item.ref_id: item for item in normalized}
    normalized = list(unique.values())
    if len(normalized) > MAX_CONTEXT_REFS:
        raise ContextPolicyError(
            f"A Trace may authorize at most {MAX_CONTEXT_REFS} context references"
        )
    total_chars = sum(context_chars(item.content) for item in normalized)
    if total_chars > MAX_CONTEXT_SNAPSHOTS_CHARS:
        raise ContextPolicyError(
            f"Context snapshots exceed the {MAX_CONTEXT_SNAPSHOTS_CHARS}-character total limit"
        )
    context[CONTEXT_ACCESS_KEY] = {
        "snapshots": {
            item.ref_id: item.model_dump(mode="json") for item in normalized
        },
        "metrics": _metrics(
            root_task_anchor=root_task_anchor,
            task_brief=task_brief,
            snapshots=normalized,
        ),
    }


def add_context_snapshot(
    context: dict[str, Any],
    snapshot: ContextSnapshot,
    *,
    root_task_anchor: RootTaskAnchor | dict[str, Any],
    task_brief: TaskBrief | dict[str, Any] | None,
) -> ContextRef:
    """向当前 Trace增加一个框架生成的本地快照。"""
    access = ensure_context_access(context)
    snapshots = list(access["snapshots"].values())
    snapshots.append(snapshot)
    replace_context_access(
        context,
        snapshots,
        root_task_anchor=root_task_anchor,
        task_brief=task_brief,
    )
    return snapshot.to_ref()


def build_child_context_access(
    *,
    parent_context: dict[str, Any],
    parent_trace_id: str,
    root_trace_id: str,
    uid: str | None,
    parent_task_state: dict[str, Any],
    child_task_brief: TaskBrief,
    root_task_anchor: RootTaskAnchor,
    granted_at_sequence: int = 0,
) -> dict[str, Any]:
    """解析显式转授并为新孩子生成完整的本地授权目录。

    直接父 TaskBrief自动占一个引用；其他祖先或审核结论必须由父级在
    ``child_task_brief.context_refs`` 中逐层明确转授。
    """
    snapshots: list[ContextSnapshot] = []
    parent_brief = parent_task_state.get("task_brief")
    if parent_brief is not None:
        parent_brief_model = TaskBrief.model_validate(parent_brief)
        snapshots.append(create_context_snapshot(
            kind="task_brief",
            summary=parent_brief_model.objective,
            source_trace_id=parent_trace_id,
            root_trace_id=root_trace_id,
            uid=uid,
            content=parent_brief_model.model_dump(mode="json"),
            granted_at_sequence=granted_at_sequence,
        ))

    parent_snapshots = ensure_context_access(parent_context)["snapshots"]
    for ref in child_task_brief.context_refs:
        raw = parent_snapshots.get(ref.ref_id)
        if raw is None:
            raise ContextPolicyError(
                f"Context reference is not authorized for the parent: {ref.ref_id}"
            )
        snapshot = _validated_snapshot(raw)
        if snapshot.ref_id != ref.ref_id:
            raise ContextPolicyError(
                f"Context reference directory entry does not match its key: {ref.ref_id}"
            )
        if snapshot.version != ref.version:
            raise ContextPolicyError(
                f"Context reference version is not authorized: {ref.ref_id}"
            )
        if snapshot.root_trace_id != root_trace_id or snapshot.uid != uid:
            raise ContextPolicyError("Context reference belongs to another tree or owner")
        resource_meta = (
            snapshot.content.get("application_resource")
            if snapshot.kind == "application_resource"
            and isinstance(snapshot.content, dict)
            else None
        )
        if (
            isinstance(resource_meta, dict)
            and resource_meta.get("inheritance") == "task_only"
        ):
            raise ContextPolicyError(
                f"Application context is task_only and cannot be delegated: {ref.ref_id}"
            )
        snapshots.append(snapshot.model_copy(
            update={"granted_at_sequence": granted_at_sequence}
        ))

    child_context: dict[str, Any] = {}
    replace_context_access(
        child_context,
        snapshots,
        root_task_anchor=root_task_anchor,
        task_brief=child_task_brief,
    )
    return child_context[CONTEXT_ACCESS_KEY]


def get_authorized_context_snapshot(
    context: dict[str, Any],
    *,
    ref_id: str,
    version: str,
    root_trace_id: str,
    uid: str | None,
) -> ContextSnapshot:
    """只从当前 Trace本地目录读取并复核一个授权快照。"""
    raw = ensure_context_access(context)["snapshots"].get(ref_id)
    if raw is None:
        raise ContextPolicyError(f"Context reference is not authorized: {ref_id}")
    snapshot = _validated_snapshot(raw)
    if snapshot.ref_id != ref_id:
        raise ContextPolicyError(
            f"Context reference directory entry does not match its key: {ref_id}"
        )
    if snapshot.version != version:
        raise ContextPolicyError(f"Context reference version is not authorized: {ref_id}")
    if snapshot.root_trace_id != root_trace_id or snapshot.uid != uid:
        raise ContextPolicyError("Context reference belongs to another tree or owner")
    return snapshot


def context_ref_descriptors(context: dict[str, Any]) -> list[dict[str, Any]]:
    """返回可展示但不展开内容的授权引用元数据。"""
    descriptors = []
    for raw in ensure_context_access(context)["snapshots"].values():
        snapshot = _validated_snapshot(raw)
        descriptors.append({
            "ref_id": snapshot.ref_id,
            "version": snapshot.version,
            "kind": snapshot.kind,
            "summary": snapshot.summary,
            "source_trace_id": snapshot.source_trace_id,
        })
    return descriptors


def prune_context_access(context: dict[str, Any], cutoff_sequence: int) -> None:
    """回溯时删除当前 Trace在截断点之后新获得的本地引用。"""
    access = ensure_context_access(context)
    kept = [
        _validated_snapshot(raw)
        for raw in access["snapshots"].values()
        if int(raw.get("granted_at_sequence", 0)) <= cutoff_sequence
    ]
    anchor = require_root_task_anchor(context)
    state = context.get("task_protocol") or {}
    replace_context_access(
        context,
        kept,
        root_task_anchor=anchor,
        task_brief=state.get("task_brief"),
    )


def render_recursive_context(context: dict[str, Any]) -> str:
    """为初始任务和周期刷新生成同一份 Recursive上下文摘要。"""
    anchor = require_root_task_anchor(context)
    state = context.get("task_protocol") or {}
    access = ensure_context_access(context)
    payload = {
        "root_task_anchor": anchor.model_dump(mode="json"),
        "task_brief_version": state.get("task_brief_version", 0),
        "task_progress": (
            current_task_progress(state).model_dump(mode="json")
            if state.get("task_progress_head_revision") is not None
            else None
        ),
        "available_context_refs": context_ref_descriptors(context),
        "context_usage": access.get("metrics", {}),
    }
    return "## Recursive Task Context\n\n" + canonical_json(payload)