Nexum-Expanded/runtime/src/nexum_runtime/tooling/artifacts.py

"""Content-addressed artifacts for tool outputs and generated files."""

from __future__ import annotations

import hashlib
import hmac
import json
import os
import secrets
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from .sandbox import control_root, ensure_control_root


@dataclass(frozen=True)
class ArtifactRecord:
    artifact_id: str
    sha256: str
    media_type: str
    bytes: int
    created_unix: int
    source: str
    relative_path: str

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


def _atomic_bytes(path: Path, data: bytes) -> None:
    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    os.chmod(path.parent, 0o700)
    temporary = path.with_name(f".{path.name}.{secrets.token_hex(8)}.tmp")
    temporary.write_bytes(data)
    temporary.chmod(0o600)
    with temporary.open("rb") as handle:
        os.fsync(handle.fileno())
    os.replace(temporary, path)
    directory_fd = os.open(path.parent, os.O_RDONLY)
    try:
        os.fsync(directory_fd)
    finally:
        os.close(directory_fd)


class ArtifactStore:
    """Store complete outputs while allowing model-selected focused reads."""

    def __init__(self, workspace: str | Path) -> None:
        workspace_root = ensure_control_root(workspace)
        self.root = control_root(workspace_root) / "artifacts"
        self.root.mkdir(parents=True, exist_ok=True, mode=0o700)
        os.chmod(self.root, 0o700)

    @staticmethod
    def _session_digest(
        *, session_id: str = "", session_id_sha256: str = ""
    ) -> str:
        computed = (
            hashlib.sha256(session_id.encode("utf-8")).hexdigest()
            if session_id
            else ""
        )
        if session_id_sha256:
            if len(session_id_sha256) != 64 or any(
                character not in "0123456789abcdef"
                for character in session_id_sha256
            ):
                raise ValueError("session digest is invalid")
            if computed and not hmac.compare_digest(computed, session_id_sha256):
                raise ValueError("session id and digest do not match")
            return session_id_sha256
        return computed

    def _reference_path(self, session_digest: str, artifact_id: str) -> Path:
        self._data_path(artifact_id)
        return self.root / "refs" / session_digest / f"{artifact_id}.json"

    def _associate(
        self,
        record: ArtifactRecord,
        *,
        source: str,
        session_id: str = "",
        session_id_sha256: str = "",
    ) -> None:
        digest = self._session_digest(
            session_id=session_id,
            session_id_sha256=session_id_sha256,
        )
        if not digest:
            return
        payload = {
            "artifact_id": record.artifact_id,
            "created_unix": int(time.time()),
            "source": source,
        }
        _atomic_bytes(
            self._reference_path(digest, record.artifact_id),
            (json.dumps(payload, sort_keys=True) + "\n").encode("utf-8"),
        )

    def _require_session(
        self,
        artifact_id: str,
        *,
        session_id: str = "",
        session_id_sha256: str = "",
    ) -> None:
        digest = self._session_digest(
            session_id=session_id,
            session_id_sha256=session_id_sha256,
        )
        if digest and not self._reference_path(digest, artifact_id).is_file():
            raise PermissionError("artifact does not belong to this session")

    def _data_path(self, artifact_id: str) -> Path:
        if not artifact_id.startswith("art_") or len(artifact_id) != 36:
            raise ValueError("artifact id is invalid")
        suffix = artifact_id.removeprefix("art_")
        if any(char not in "0123456789abcdef" for char in suffix):
            raise ValueError("artifact id is invalid")
        return self.root / f"{artifact_id}.bin"

    def _metadata_path(self, artifact_id: str) -> Path:
        return self._data_path(artifact_id).with_suffix(".json")

    def put(
        self,
        data: bytes,
        *,
        media_type: str,
        source: str,
        session_id: str = "",
        session_id_sha256: str = "",
    ) -> ArtifactRecord:
        digest = hashlib.sha256(data).hexdigest()
        artifact_id = "art_" + digest[:32]
        data_path = self._data_path(artifact_id)
        metadata_path = self._metadata_path(artifact_id)
        if not data_path.exists():
            _atomic_bytes(data_path, data)
        if metadata_path.is_file():
            record = self.get(artifact_id)
        else:
            record = ArtifactRecord(
                artifact_id=artifact_id,
                sha256=digest,
                media_type=media_type,
                bytes=len(data),
                created_unix=int(time.time()),
                source=source,
                relative_path=(
                    Path(".nexum")
                    / "artifacts"
                    / data_path.relative_to(self.root)
                ).as_posix(),
            )
            _atomic_bytes(
                metadata_path,
                (json.dumps(record.to_dict(), sort_keys=True) + "\n").encode(
                    "utf-8"
                ),
            )
        self._associate(
            record,
            source=source,
            session_id=session_id,
            session_id_sha256=session_id_sha256,
        )
        return record

    def put_text(
        self,
        text: str,
        *,
        media_type: str = "text/plain; charset=utf-8",
        source: str,
        session_id: str = "",
        session_id_sha256: str = "",
    ) -> ArtifactRecord:
        return self.put(
            text.encode("utf-8"),
            media_type=media_type,
            source=source,
            session_id=session_id,
            session_id_sha256=session_id_sha256,
        )

    def get(self, artifact_id: str) -> ArtifactRecord:
        payload = json.loads(self._metadata_path(artifact_id).read_text("utf-8"))
        if not isinstance(payload, dict):
            raise RuntimeError("artifact metadata is invalid")
        record = ArtifactRecord(**payload)
        data = self._data_path(artifact_id).read_bytes()
        if not hashlib.sha256(data).hexdigest() == record.sha256:
            raise RuntimeError("artifact integrity check failed")
        return record

    def read(
        self,
        artifact_id: str,
        *,
        offset: int = 0,
        length: int | None = None,
        session_id: str = "",
        session_id_sha256: str = "",
    ) -> tuple[ArtifactRecord, bytes]:
        self._require_session(
            artifact_id,
            session_id=session_id,
            session_id_sha256=session_id_sha256,
        )
        record = self.get(artifact_id)
        if offset < 0 or (length is not None and length < 0):
            raise ValueError("artifact offset and length must be non-negative")
        data = self._data_path(artifact_id).read_bytes()
        return record, data[offset:] if length is None else data[offset : offset + length]

    def list(
        self,
        *,
        session_id: str = "",
        session_id_sha256: str = "",
    ) -> tuple[ArtifactRecord, ...]:
        digest = self._session_digest(
            session_id=session_id,
            session_id_sha256=session_id_sha256,
        )
        if digest:
            paths = sorted((self.root / "refs" / digest).glob("art_*.json"))
            artifact_ids = [path.stem for path in paths]
        else:
            artifact_ids = [path.stem for path in sorted(self.root.glob("art_*.json"))]
        records: list[ArtifactRecord] = []
        for artifact_id in artifact_ids:
            try:
                records.append(self.get(artifact_id))
            except (OSError, TypeError, ValueError, RuntimeError):
                continue
        return tuple(records)


__all__ = ["ArtifactRecord", "ArtifactStore"]