"""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"]