inference/src/ax_doclayout/artifact.py

# Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
"""Immutable AXDocLayout-M v1.2 artifact receipts and bundle verification.

The Hub bundle contains nine files.  ``artifact-manifest.json`` describes the
other eight payload files, while a separately trusted receipt binds all nine
files and their content root.  Keeping the receipt outside the bundle avoids a
self-digest cycle.
"""

from __future__ import annotations

import hashlib
import json
import os
import re
import stat
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from ax_doclayout.contract import PREPROCESS_ID, TAXONOMY_ID
from ax_doclayout.errors import ArtifactError

ARTIFACT_ROOT_SCHEMA = "ax-doclayout.artifact-root/1"
ARTIFACT_RECEIPT_SCHEMA = "ax-doclayout.artifact-receipt/1"
ARTIFACT_MANIFEST_SCHEMA = "ax-doclayout.artifact-manifest/1"
MODEL_ID = "AutomatosX/AXDocLayout-M-v1.2-CUDA-TF32"
LEGACY_MODEL_ID = "AutomatosX/AXDocStruct-M-v0.1"
SOURCE_CHECKPOINT_SHA256 = (
    "sha256:e4747f71d40e14c89c561c720c043ad2be5dadf68668f071983f3ef1111a70fa"
)
STAGED_ONNX_SHA256 = (
    "sha256:2d5928719c391e424a1132f8c5eadb23e4d164d3576f00ff14010ebe55aaea81"
)
STAGED_SAFETENSORS_SHA256 = (
    "sha256:3b97ba1c0b1b0816f45e79a76bb2fd495e026d2394471b76ce1a442baa2989c4"
)
SUPPORTED_ONNX_OPSET = 16
ARTIFACT_BUNDLE_FILES = (
    "LICENSE",
    "MODEL_CARD.md",
    "NOTICE",
    "artifact-manifest.json",
    "config.json",
    "model.onnx",
    "model.safetensors",
    "preprocessor_config.json",
    "taxonomy.json",
)
ARTIFACT_MANIFEST_PAYLOAD_FILES = tuple(
    name for name in ARTIFACT_BUNDLE_FILES if name != "artifact-manifest.json"
)

_MAX_RECEIPT_BYTES = 64 * 1024
_MAX_MANIFEST_BYTES = 64 * 1024
_REVISION_RE = re.compile(r"^[0-9a-f]{40}$")
_SHA256_RE = re.compile(r"^sha256:[0-9a-f]{64}$")


@dataclass(frozen=True, slots=True)
class ArtifactFile:
    """One manifest-relative artifact file identity."""

    name: str
    sha256: str

    def __post_init__(self) -> None:
        if self.name not in ARTIFACT_BUNDLE_FILES:
            raise ArtifactError("artifact_file_name_invalid")
        if (
            not isinstance(self.sha256, str)
            or _SHA256_RE.fullmatch(self.sha256) is None
        ):
            raise ArtifactError("artifact_file_digest_invalid")

    def to_json(self) -> dict[str, str]:
        """Return the canonical JSON shape."""
        return {"name": self.name, "sha256": self.sha256}


@dataclass(frozen=True, slots=True)
class ArtifactReceipt:
    """Trusted, package-external identity for one immutable Hub bundle."""

    model_id: str
    model_revision: str
    artifact_sha256: str
    artifact_files: Sequence[ArtifactFile]
    taxonomy: str = TAXONOMY_ID
    preprocess: str = PREPROCESS_ID
    schema_version: str = ARTIFACT_RECEIPT_SCHEMA

    def __post_init__(self) -> None:
        files = tuple(self.artifact_files)
        object.__setattr__(self, "artifact_files", files)
        if self.schema_version != ARTIFACT_RECEIPT_SCHEMA:
            raise ArtifactError("artifact_receipt_schema_invalid")
        if self.model_id != MODEL_ID:
            raise ArtifactError("artifact_model_id_invalid")
        if (
            not isinstance(self.model_revision, str)
            or _REVISION_RE.fullmatch(self.model_revision) is None
        ):
            raise ArtifactError("artifact_revision_invalid")
        if self.taxonomy != TAXONOMY_ID:
            raise ArtifactError("artifact_taxonomy_invalid")
        if self.preprocess != PREPROCESS_ID:
            raise ArtifactError("artifact_preprocess_invalid")
        if (
            not isinstance(self.artifact_sha256, str)
            or _SHA256_RE.fullmatch(self.artifact_sha256) is None
        ):
            raise ArtifactError("artifact_root_digest_invalid")
        _validate_complete_files(files)
        expected_root = artifact_root_sha256(files)
        if self.artifact_sha256 != expected_root:
            raise ArtifactError("artifact_root_digest_mismatch")

    def to_json(self) -> dict[str, object]:
        """Return the strict receipt JSON shape."""
        return {
            "artifact_files": [item.to_json() for item in self.artifact_files],
            "artifact_sha256": self.artifact_sha256,
            "model_id": self.model_id,
            "model_revision": self.model_revision,
            "preprocess": self.preprocess,
            "schema_version": self.schema_version,
            "taxonomy": self.taxonomy,
        }


@dataclass(frozen=True, slots=True)
class ArtifactManifest:
    """Canonical in-bundle identity for the eight non-manifest payloads."""

    artifact_files: Sequence[ArtifactFile]
    onnx_opset: int = SUPPORTED_ONNX_OPSET
    model_id: str = MODEL_ID
    taxonomy: str = TAXONOMY_ID
    preprocess: str = PREPROCESS_ID
    schema_version: str = ARTIFACT_MANIFEST_SCHEMA

    def __post_init__(self) -> None:
        files = tuple(self.artifact_files)
        object.__setattr__(self, "artifact_files", files)
        if self.schema_version != ARTIFACT_MANIFEST_SCHEMA:
            raise ArtifactError("artifact_manifest_schema_invalid")
        if self.model_id != MODEL_ID:
            raise ArtifactError("artifact_model_id_invalid")
        if self.taxonomy != TAXONOMY_ID:
            raise ArtifactError("artifact_taxonomy_invalid")
        if self.preprocess != PREPROCESS_ID:
            raise ArtifactError("artifact_preprocess_invalid")
        if (
            isinstance(self.onnx_opset, bool)
            or not isinstance(self.onnx_opset, int)
            or self.onnx_opset != SUPPORTED_ONNX_OPSET
        ):
            raise ArtifactError("artifact_onnx_opset_unsupported")
        _validate_manifest_files(files)

    def to_json(self) -> dict[str, object]:
        """Return the strict manifest JSON shape."""
        return {
            "backends": {
                "onnx": {
                    "file": "model.onnx",
                    "format": "onnx",
                    "opset": self.onnx_opset,
                    "precision": "fp32",
                    "standard_ops_only": True,
                },
                "pytorch": {
                    "file": "model.safetensors",
                    "format": "safetensors",
                    "precision": "fp32",
                },
            },
            "files": [item.to_json() for item in self.artifact_files],
            "model_id": self.model_id,
            "preprocess": self.preprocess,
            "schema_version": self.schema_version,
            "taxonomy": self.taxonomy,
        }


def canonical_json_bytes(payload: Mapping[str, object]) -> bytes:
    """Serialize strict ASCII JSON without a trailing newline."""
    try:
        return json.dumps(
            payload,
            allow_nan=False,
            ensure_ascii=True,
            separators=(",", ":"),
            sort_keys=True,
        ).encode("ascii")
    except (TypeError, ValueError, UnicodeError) as exc:
        raise ArtifactError("artifact_json_invalid") from exc


def _validate_complete_files(files: Sequence[ArtifactFile]) -> None:
    if any(not isinstance(item, ArtifactFile) for item in files):
        raise ArtifactError("artifact_files_invalid")
    names = tuple(item.name for item in files)
    if len(names) != len(set(names)):
        raise ArtifactError("artifact_file_duplicate")
    if names != ARTIFACT_BUNDLE_FILES:
        raise ArtifactError("artifact_file_set_invalid")


def _validate_manifest_files(files: Sequence[ArtifactFile]) -> None:
    if any(not isinstance(item, ArtifactFile) for item in files):
        raise ArtifactError("artifact_files_invalid")
    names = tuple(item.name for item in files)
    if len(names) != len(set(names)):
        raise ArtifactError("artifact_file_duplicate")
    if names != ARTIFACT_MANIFEST_PAYLOAD_FILES:
        raise ArtifactError("artifact_manifest_file_set_invalid")


def artifact_root_sha256(files: Sequence[ArtifactFile]) -> str:
    """Compute the content root over the sorted complete nine-file set."""
    items = tuple(files)
    _validate_complete_files(items)
    payload = {
        "files": [item.to_json() for item in items],
        "schema_version": ARTIFACT_ROOT_SCHEMA,
    }
    return "sha256:" + hashlib.sha256(canonical_json_bytes(payload)).hexdigest()


def _open_regular_file(path: Path) -> int:
    flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0)
    try:
        descriptor = os.open(path, flags)
    except OSError as exc:
        raise ArtifactError("artifact_file_unavailable") from exc
    try:
        opened = os.fstat(descriptor)
        observed = path.lstat()
        if (
            not stat.S_ISREG(opened.st_mode)
            or not stat.S_ISREG(observed.st_mode)
            or (opened.st_dev, opened.st_ino) != (observed.st_dev, observed.st_ino)
        ):
            raise ArtifactError("artifact_file_unavailable")
    except (OSError, ArtifactError):
        os.close(descriptor)
        raise
    return descriptor


def _hash_file(path: Path) -> str:
    digest = hashlib.sha256()
    descriptor = _open_regular_file(path)
    try:
        with os.fdopen(descriptor, "rb") as handle:
            for chunk in iter(lambda: handle.read(1024 * 1024), b""):
                digest.update(chunk)
    except OSError as exc:
        raise ArtifactError("artifact_file_unavailable") from exc
    return "sha256:" + digest.hexdigest()


def _bundle_files(directory: Path) -> tuple[ArtifactFile, ...]:
    try:
        directory_mode = directory.lstat().st_mode
    except OSError as exc:
        raise ArtifactError("artifact_bundle_unavailable") from exc
    if not stat.S_ISDIR(directory_mode):
        raise ArtifactError("artifact_bundle_unavailable")
    try:
        names = tuple(sorted(path.name for path in directory.iterdir()))
    except OSError as exc:
        raise ArtifactError("artifact_bundle_unavailable") from exc
    if names != ARTIFACT_BUNDLE_FILES:
        raise ArtifactError("artifact_bundle_file_set_invalid")
    files: list[ArtifactFile] = []
    for name in ARTIFACT_BUNDLE_FILES:
        path = directory / name
        files.append(ArtifactFile(name=name, sha256=_hash_file(path)))
    return tuple(files)


def _payload_files(directory: Path) -> tuple[ArtifactFile, ...]:
    try:
        directory_mode = directory.lstat().st_mode
    except OSError as exc:
        raise ArtifactError("artifact_bundle_unavailable") from exc
    if not stat.S_ISDIR(directory_mode):
        raise ArtifactError("artifact_bundle_unavailable")
    try:
        names = tuple(sorted(path.name for path in directory.iterdir()))
    except OSError as exc:
        raise ArtifactError("artifact_bundle_unavailable") from exc
    allowed_sets = {
        ARTIFACT_MANIFEST_PAYLOAD_FILES,
        ARTIFACT_BUNDLE_FILES,
    }
    if names not in allowed_sets:
        raise ArtifactError("artifact_bundle_file_set_invalid")
    return tuple(
        ArtifactFile(name=name, sha256=_hash_file(directory / name))
        for name in ARTIFACT_MANIFEST_PAYLOAD_FILES
    )


def build_artifact_manifest(
    directory: Path,
    *,
    onnx_opset: int = SUPPORTED_ONNX_OPSET,
) -> ArtifactManifest:
    """Build the manifest from the exact eight payload files."""
    return ArtifactManifest(
        artifact_files=_payload_files(directory),
        onnx_opset=onnx_opset,
    )


def build_artifact_receipt(
    directory: Path,
    *,
    model_revision: str,
) -> ArtifactReceipt:
    """Build a receipt after a complete bundle has an immutable Hub revision."""
    verify_artifact_manifest(directory)
    files = _bundle_files(directory)
    return ArtifactReceipt(
        model_id=MODEL_ID,
        model_revision=model_revision,
        artifact_sha256=artifact_root_sha256(files),
        artifact_files=files,
    )


def verify_artifact_bundle(directory: Path, receipt: ArtifactReceipt) -> None:
    """Fail closed unless every staged byte matches a trusted receipt."""
    if not isinstance(receipt, ArtifactReceipt):
        raise ArtifactError("artifact_receipt_invalid")
    verify_artifact_manifest(directory)
    observed = _bundle_files(directory)
    if observed != receipt.artifact_files:
        raise ArtifactError("artifact_file_digest_mismatch")
    if artifact_root_sha256(observed) != receipt.artifact_sha256:
        raise ArtifactError("artifact_root_digest_mismatch")


def _reject_json_constant(_: str) -> None:
    raise ValueError("non_finite_json")


def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
    result: dict[str, Any] = {}
    for key, value in pairs:
        if key in result:
            raise ValueError("duplicate_json_key")
        result[key] = value
    return result


def _exact_keys(
    value: Mapping[str, Any],
    expected: set[str],
    code: str = "artifact_receipt_fields_invalid",
) -> None:
    if set(value) != expected:
        raise ArtifactError(code)


def _parse_artifact_file(value: object) -> ArtifactFile:
    if not isinstance(value, dict):
        raise ArtifactError("artifact_file_invalid")
    _exact_keys(value, {"name", "sha256"}, "artifact_file_invalid")
    try:
        return ArtifactFile(name=value["name"], sha256=value["sha256"])
    except (KeyError, TypeError) as exc:
        raise ArtifactError("artifact_file_invalid") from exc


def _validated_backends(value: object) -> int:
    if not isinstance(value, dict):
        raise ArtifactError("artifact_backends_invalid")
    _exact_keys(value, {"onnx", "pytorch"}, "artifact_backends_invalid")
    onnx = value["onnx"]
    pytorch = value["pytorch"]
    if not isinstance(onnx, dict) or not isinstance(pytorch, dict):
        raise ArtifactError("artifact_backends_invalid")
    _exact_keys(
        onnx,
        {"file", "format", "opset", "precision", "standard_ops_only"},
        "artifact_backends_invalid",
    )
    _exact_keys(
        pytorch,
        {"file", "format", "precision"},
        "artifact_backends_invalid",
    )
    opset = onnx["opset"]
    if (
        onnx["file"] != "model.onnx"
        or onnx["format"] != "onnx"
        or onnx["precision"] != "fp32"
        or onnx["standard_ops_only"] is not True
        or isinstance(opset, bool)
        or not isinstance(opset, int)
        or opset != SUPPORTED_ONNX_OPSET
        or pytorch
        != {
            "file": "model.safetensors",
            "format": "safetensors",
            "precision": "fp32",
        }
    ):
        raise ArtifactError("artifact_backends_invalid")
    return opset


def parse_artifact_manifest(raw: bytes) -> ArtifactManifest:
    """Parse canonical manifest bytes with exact backend and payload metadata."""
    if len(raw) > _MAX_MANIFEST_BYTES:
        raise ArtifactError("artifact_manifest_too_large")
    try:
        value = json.loads(
            raw,
            object_pairs_hook=_unique_object,
            parse_constant=_reject_json_constant,
        )
    except (TypeError, ValueError, UnicodeError, json.JSONDecodeError) as exc:
        raise ArtifactError("artifact_manifest_json_invalid") from exc
    if not isinstance(value, dict):
        raise ArtifactError("artifact_manifest_invalid")
    _exact_keys(
        value,
        {
            "backends",
            "files",
            "model_id",
            "preprocess",
            "schema_version",
            "taxonomy",
        },
        "artifact_manifest_fields_invalid",
    )
    raw_files = value["files"]
    if not isinstance(raw_files, list):
        raise ArtifactError("artifact_files_invalid")
    manifest = ArtifactManifest(
        artifact_files=tuple(_parse_artifact_file(item) for item in raw_files),
        onnx_opset=_validated_backends(value["backends"]),
        model_id=value["model_id"],
        taxonomy=value["taxonomy"],
        preprocess=value["preprocess"],
        schema_version=value["schema_version"],
    )
    if canonical_json_bytes(manifest.to_json()) != raw:
        raise ArtifactError("artifact_manifest_not_canonical")
    return manifest


def load_artifact_manifest(path: Path) -> ArtifactManifest:
    """Load one bounded canonical in-bundle manifest."""
    descriptor = _open_regular_file(path)
    try:
        with os.fdopen(descriptor, "rb") as handle:
            raw = handle.read(_MAX_MANIFEST_BYTES + 1)
    except OSError as exc:
        raise ArtifactError("artifact_manifest_unavailable") from exc
    if len(raw) > _MAX_MANIFEST_BYTES:
        raise ArtifactError("artifact_manifest_too_large")
    return parse_artifact_manifest(raw)


def verify_artifact_manifest(directory: Path) -> ArtifactManifest:
    """Verify the manifest against the exact eight payload bytes."""
    manifest = load_artifact_manifest(directory / "artifact-manifest.json")
    observed = _payload_files(directory)
    if tuple(manifest.artifact_files) != observed:
        raise ArtifactError("artifact_manifest_digest_mismatch")
    return manifest


def parse_artifact_receipt(raw: bytes) -> ArtifactReceipt:
    """Parse canonical trusted-receipt bytes with duplicate-key rejection."""
    if len(raw) > _MAX_RECEIPT_BYTES:
        raise ArtifactError("artifact_receipt_too_large")
    try:
        value = json.loads(
            raw,
            object_pairs_hook=_unique_object,
            parse_constant=_reject_json_constant,
        )
    except (TypeError, ValueError, UnicodeError, json.JSONDecodeError) as exc:
        raise ArtifactError("artifact_receipt_json_invalid") from exc
    if not isinstance(value, dict):
        raise ArtifactError("artifact_receipt_invalid")
    _exact_keys(
        value,
        {
            "artifact_files",
            "artifact_sha256",
            "model_id",
            "model_revision",
            "preprocess",
            "schema_version",
            "taxonomy",
        },
    )
    raw_files = value["artifact_files"]
    if not isinstance(raw_files, list):
        raise ArtifactError("artifact_files_invalid")
    receipt = ArtifactReceipt(
        model_id=value["model_id"],
        model_revision=value["model_revision"],
        artifact_sha256=value["artifact_sha256"],
        artifact_files=tuple(_parse_artifact_file(item) for item in raw_files),
        taxonomy=value["taxonomy"],
        preprocess=value["preprocess"],
        schema_version=value["schema_version"],
    )
    if canonical_json_bytes(receipt.to_json()) != raw:
        raise ArtifactError("artifact_receipt_not_canonical")
    return receipt


def load_trusted_receipt(path: Path) -> ArtifactReceipt:
    """Load one operator-pinned local receipt without contacting the Hub."""
    descriptor = _open_regular_file(path)
    try:
        with os.fdopen(descriptor, "rb") as handle:
            raw = handle.read(_MAX_RECEIPT_BYTES + 1)
    except OSError as exc:
        raise ArtifactError("artifact_receipt_unavailable") from exc
    if len(raw) > _MAX_RECEIPT_BYTES:
        raise ArtifactError("artifact_receipt_too_large")
    return parse_artifact_receipt(raw)


__all__ = [
    "ARTIFACT_BUNDLE_FILES",
    "ARTIFACT_MANIFEST_PAYLOAD_FILES",
    "ARTIFACT_MANIFEST_SCHEMA",
    "ARTIFACT_RECEIPT_SCHEMA",
    "ARTIFACT_ROOT_SCHEMA",
    "LEGACY_MODEL_ID",
    "MODEL_ID",
    "SUPPORTED_ONNX_OPSET",
    "ArtifactFile",
    "ArtifactManifest",
    "ArtifactReceipt",
    "artifact_root_sha256",
    "build_artifact_manifest",
    "build_artifact_receipt",
    "canonical_json_bytes",
    "load_artifact_manifest",
    "load_trusted_receipt",
    "parse_artifact_manifest",
    "parse_artifact_receipt",
    "verify_artifact_bundle",
    "verify_artifact_manifest",
]