code/self_healing_finetuning/rollback/anchor.py

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Any

from observability.audit_log import utc_now


def _canonical(payload: dict[str, Any]) -> str:
    return json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True)


def _checksum(payload: dict[str, Any]) -> str:
    body = {key: value for key, value in payload.items() if key != "checksum_sha256"}
    return hashlib.sha256(_canonical(body).encode("utf-8")).hexdigest()


def anchor_path(implementation_root: Path, env: str) -> Path:
    return implementation_root / "rollback" / "last_good" / f"{env}_anchor.json"


def validate_rollback_anchor(anchor: dict[str, Any]) -> list[str]:
    required = [
        "anchor_id",
        "env",
        "type",
        "status",
        "model_id",
        "deployment_id",
        "source",
        "rollback_artifacts",
        "created_at",
        "checksum_sha256",
    ]
    errors = [f"rollback anchor missing required field: {field}" for field in required if field not in anchor]
    if anchor.get("type") != "last_good":
        errors.append(f"rollback anchor type is not last_good: {anchor.get('type')}")
    if anchor.get("status") != "active":
        errors.append(f"rollback anchor is not active: {anchor.get('status')}")
    if "checksum_sha256" in anchor and anchor["checksum_sha256"] != _checksum(anchor):
        errors.append("rollback anchor checksum mismatch")
    return errors


def ensure_last_good_anchor(implementation_root: Path, *, env: str, model_id: str) -> dict[str, Any]:
    path = anchor_path(implementation_root, env)
    if path.exists():
        return json.loads(path.read_text(encoding="utf-8"))

    anchor = {
        "anchor_id": f"last_good_{env}_bootstrap",
        "env": env,
        "type": "last_good",
        "status": "active",
        "model_id": model_id,
        "deployment_id": f"bootstrap_{env}_baseline",
        "source": "initial_linvest21_fingpt_bootstrap_anchor",
        "rollback_artifacts": {
            "model_repo": model_id,
            "path": str(implementation_root / "rollback" / "last_good"),
            "restore_mode": "switch_endpoint_or_registry_pointer",
        },
        "created_at": utc_now(),
    }
    anchor["checksum_sha256"] = _checksum(anchor)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(anchor, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return anchor


def load_last_good_anchor(implementation_root: Path, *, env: str) -> dict[str, Any] | None:
    path = anchor_path(implementation_root, env)
    if not path.exists():
        return None
    return json.loads(path.read_text(encoding="utf-8"))