Agent/cyber_agent/core/run_snapshot.py

"""Versioned, hash-bound snapshots for reproducible Agent runs."""

from __future__ import annotations

import json
from dataclasses import asdict, is_dataclass
from hashlib import sha256
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field


RUN_CONFIG_SNAPSHOT_CONTEXT_KEY = "run_config_snapshot"
RUN_CONFIG_SNAPSHOT_HASH_CONTEXT_KEY = "run_config_snapshot_hash"


class RunConfigSnapshotError(ValueError):
    """A persisted run snapshot is absent, malformed, or was modified."""


class RunConfigSnapshotV1(BaseModel):
    """Persisted behavior-affecting fields from ``RunConfig``.

    Runtime routing fields (trace_id, parent IDs, rewind position and forced side
    branches) are intentionally excluded because they describe one invocation,
    not the immutable behavior of the run.
    """

    model_config = ConfigDict(extra="forbid")

    schema_version: Literal[1] = 1
    model: str
    temperature: float
    max_iterations: int = Field(gt=0)
    extra_llm_params: dict[str, Any]

    tools: list[str] | None
    tool_groups: list[str] | None
    exclude_tools: list[str]
    auto_execute_tools: bool

    agent_type: str
    uid: str | None
    skills: list[str] | None
    system_prompt_hash: str | None = Field(
        default=None,
        pattern=r"^[0-9a-f]{64}$",
    )

    enable_memory: bool
    memory: dict[str, Any] | None
    memory_identity: str | None = Field(
        default=None,
        pattern=r"^memory-v1:[0-9a-f]{64}$",
    )
    knowledge: dict[str, Any]

    parallel_tool_execution: bool
    child_execution_mode: Literal["sequential", "parallel"]
    max_parallel_children: int = Field(gt=0)
    side_branch_max_turns: int = Field(gt=0)
    goal_compression: Literal["none", "on_complete", "on_overflow"]

    enable_prompt_caching: bool
    enable_research_flow: bool
    project_name: str | None = None
    custom_context: dict[str, Any]
    legacy_inferred: bool = False

    @classmethod
    def from_run_config(
        cls,
        config: Any,
        *,
        memory_identity: str | None,
        system_prompt_hash: str | None = None,
        legacy_inferred: bool = False,
    ) -> "RunConfigSnapshotV1":
        def dump(value: Any) -> Any:
            if value is None:
                return None
            if is_dataclass(value):
                return asdict(value)
            if isinstance(value, BaseModel):
                return value.model_dump(mode="json")
            if isinstance(value, dict):
                return dict(value)
            raise TypeError(f"run snapshot field is not serializable: {type(value)!r}")

        context = config.context or {}
        return cls(
            model=config.model,
            temperature=config.temperature,
            max_iterations=config.max_iterations,
            extra_llm_params=dict(config.extra_llm_params),
            tools=list(config.tools) if config.tools is not None else None,
            tool_groups=(
                list(config.tool_groups) if config.tool_groups is not None else None
            ),
            exclude_tools=list(config.exclude_tools),
            auto_execute_tools=config.auto_execute_tools,
            agent_type=config.agent_type,
            uid=config.uid,
            skills=list(config.skills) if config.skills is not None else None,
            system_prompt_hash=system_prompt_hash,
            enable_memory=config.enable_memory,
            memory=dump(config.memory),
            memory_identity=memory_identity,
            knowledge=dump(config.knowledge),
            parallel_tool_execution=config.parallel_tool_execution,
            child_execution_mode=config.child_execution_mode,
            max_parallel_children=config.max_parallel_children,
            side_branch_max_turns=config.side_branch_max_turns,
            goal_compression=config.goal_compression,
            enable_prompt_caching=config.enable_prompt_caching,
            enable_research_flow=config.enable_research_flow,
            project_name=context.get("project_name"),
            custom_context=dict(context),
            legacy_inferred=legacy_inferred,
        )

    @property
    def snapshot_hash(self) -> str:
        payload = json.dumps(
            self.model_dump(mode="json"),
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
            allow_nan=False,
        )
        return sha256(payload.encode("utf-8")).hexdigest()


class RunConfigSnapshotV2(RunConfigSnapshotV1):
    """Application-bound snapshot; V1's serialized shape remains unchanged."""

    schema_version: Literal[2] = 2
    application_ref: dict[str, Any]
    role_id: str = Field(min_length=1)
    role_hash: str = Field(pattern=r"^[0-9a-f]{64}$")
    effective_run_limits: dict[str, Any]

    @classmethod
    def from_run_config(
        cls,
        config: Any,
        *,
        memory_identity: str | None,
        system_prompt_hash: str | None = None,
        legacy_inferred: bool = False,
    ) -> "RunConfigSnapshotV2":
        base = RunConfigSnapshotV1.from_run_config(
            config,
            memory_identity=memory_identity,
            system_prompt_hash=system_prompt_hash,
            legacy_inferred=legacy_inferred,
        ).model_dump(mode="json")
        base.pop("schema_version", None)
        application_ref = config.application_ref
        if isinstance(application_ref, BaseModel):
            application_ref = application_ref.model_dump(mode="json")
        if not isinstance(application_ref, dict):
            raise TypeError("application_ref is required for a V2 run snapshot")
        if not config.role_id or not config.role_hash:
            raise TypeError("role_id and role_hash are required for a V2 run snapshot")
        return cls(
            **base,
            application_ref=dict(application_ref),
            role_id=config.role_id,
            role_hash=config.role_hash,
            effective_run_limits=dict(config.effective_run_limits),
        )


RunConfigSnapshot = RunConfigSnapshotV1 | RunConfigSnapshotV2


def persist_run_config_snapshot(
    context: dict[str, Any],
    snapshot: RunConfigSnapshot,
) -> None:
    context[RUN_CONFIG_SNAPSHOT_CONTEXT_KEY] = snapshot.model_dump(mode="json")
    context[RUN_CONFIG_SNAPSHOT_HASH_CONTEXT_KEY] = snapshot.snapshot_hash


def load_run_config_snapshot(context: dict[str, Any]) -> RunConfigSnapshot:
    raw = context.get(RUN_CONFIG_SNAPSHOT_CONTEXT_KEY)
    digest = context.get(RUN_CONFIG_SNAPSHOT_HASH_CONTEXT_KEY)
    if raw is None or digest is None:
        raise RunConfigSnapshotError("run config snapshot is missing")
    try:
        schema_version = raw.get("schema_version") if isinstance(raw, dict) else None
        if schema_version == 1:
            snapshot = RunConfigSnapshotV1.model_validate(raw)
        elif schema_version == 2:
            snapshot = RunConfigSnapshotV2.model_validate(raw)
        else:
            raise ValueError(f"unsupported schema_version: {schema_version!r}")
    except Exception as exc:
        raise RunConfigSnapshotError(f"invalid run config snapshot: {exc}") from exc
    if snapshot.snapshot_hash != digest:
        raise RunConfigSnapshotError("run config snapshot hash mismatch")
    return snapshot