Agent/cyber_agent/core/context_policy.py

"""Minimal, explicit context inheritance for Recursive task delegation."""

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):
    """A TaskBrief cannot cross the Recursive delegation boundary safely."""


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:
    """Validate a TaskBrief and inherit only the parent's hard constraints."""
    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:
    """Compare actual and approved briefs after applying the same policy."""
    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"
    )