mostlyright-data 0.9.0__py3-none-any.whl

mostlyright/data_harness/visual_run/store.py

@@ -0,0 +1,584 @@
+"""Durable local storage for visual-run events and immutable artifacts."""
+
+from __future__ import annotations
+
+import fcntl
+import hashlib
+import json
+import os
+import secrets
+import stat
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import BinaryIO
+
+from mostlyright.data_harness.visual_run.contracts import (
+    DOCUMENT_HEADER_CELL_ID,
+    LEGACY_DOCUMENT_HEADER_CELL_ID_PREFIX,
+    ArtifactReference,
+    VisualRunDraft,
+    VisualRunError,
+    VisualRunEvent,
+    canonical_draft_bytes,
+    canonical_event_bytes,
+    parse_event,
+)
+
+MAX_LOG_BYTES = 64 * 1024 * 1024
+MAX_EVENTS = 100_000
+LOG_NAME = "events.jsonl"
+LOCK_NAME = ".append.lock"
+ARTIFACT_DIR = "artifacts"
+
+
+def _command_bytes(draft: VisualRunDraft) -> bytes:
+    """Canonical command identity, excluding only the generated observation timestamp."""
+
+    value = draft.to_mapping()
+    value.pop("observed_at")
+    return (
+        json.dumps(
+            value, ensure_ascii=False, allow_nan=False, sort_keys=True, separators=(",", ":")
+        )
+        + "\n"
+    ).encode("utf-8")
+
+
+@dataclass(frozen=True)
+class EventPage:
+    events: tuple[VisualRunEvent, ...]
+    cursor: int
+    truncated: bool
+
+
+class AuthorityLease:
+    """A held authority lock with one commit staged after final identity validation."""
+
+    def __init__(self, validate: Callable[[], None]) -> None:
+        self._validate = validate
+        self._commit: Callable[[], None] | None = None
+
+    def __call__(self) -> None:
+        self._validate()
+
+    def commit_after_validation(self, commit: Callable[[], None]) -> None:
+        if self._commit is not None:
+            raise RuntimeError("authority lease accepts one staged commit")
+        self._commit = commit
+
+    def _finish(self) -> None:
+        self._validate()
+        if self._commit is not None:
+            self._commit()
+
+
+def _error(code: str, message: str) -> VisualRunError:
+    return VisualRunError(code, message)
+
+
+def _regular_descriptor(descriptor: int, name: str) -> os.stat_result:
+    info = os.fstat(descriptor)
+    if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1:
+        raise _error("VISUAL_RUN_STORE_IDENTITY", f"{name} must be one regular file")
+    return info
+
+
+def _read_descriptor(descriptor: int, *, maximum: int, code: str, message: str) -> bytes:
+    os.lseek(descriptor, 0, os.SEEK_SET)
+    chunks: list[bytes] = []
+    total = 0
+    while True:
+        chunk = os.read(descriptor, min(64 * 1024, maximum + 1 - total))
+        if not chunk:
+            return b"".join(chunks)
+        chunks.append(chunk)
+        total += len(chunk)
+        if total > maximum:
+            raise _error(code, message)
+
+
+@contextmanager
+def _locked(directory_fd: int) -> Iterator[BinaryIO]:
+    descriptor = os.open(
+        LOCK_NAME,
+        os.O_RDWR | os.O_CREAT | os.O_CLOEXEC | os.O_NOFOLLOW,
+        0o600,
+        dir_fd=directory_fd,
+    )
+    _regular_descriptor(descriptor, LOCK_NAME)
+    stream = os.fdopen(descriptor, "r+b", closefd=True)
+    try:
+        fcntl.flock(stream.fileno(), fcntl.LOCK_EX)
+        yield stream
+    finally:
+        fcntl.flock(stream.fileno(), fcntl.LOCK_UN)
+        stream.close()
+
+
+@contextmanager
+def visual_authority_lock(directory_fd: int, *, exclusive: bool) -> Iterator[AuthorityLease]:
+    """Serialize first visual authority installation with legacy viewer admission.
+
+    The already-pinned research directory inode is the lock object. Unlike a child lock pathname,
+    that descriptor-relative synchronization boundary cannot be substituted between validation
+    and commit. Visual writers take it exclusively while installing ``visual-run``; viewers take
+    it shared across the final absent observation and one staged legacy journal commit.
+    """
+
+    descriptor = os.open(
+        ".",
+        os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW,
+        dir_fd=directory_fd,
+    )
+    try:
+        info = os.fstat(descriptor)
+        pinned = os.fstat(directory_fd)
+        if not stat.S_ISDIR(info.st_mode) or (info.st_dev, info.st_ino) != (
+            pinned.st_dev,
+            pinned.st_ino,
+        ):
+            raise _error(
+                "VISUAL_RUN_STORE_IDENTITY",
+                "visual authority must lock the pinned research directory",
+            )
+        fcntl.flock(descriptor, fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH)
+
+        def validate() -> None:
+            try:
+                current = os.fstat(descriptor)
+                pinned_now = os.fstat(directory_fd)
+            except OSError as exc:
+                raise _error(
+                    "VISUAL_RUN_STORE_IDENTITY",
+                    "pinned research authority became unavailable during the locked operation",
+                ) from exc
+            if (
+                (current.st_dev, current.st_ino) != (info.st_dev, info.st_ino)
+                or (pinned_now.st_dev, pinned_now.st_ino) != (info.st_dev, info.st_ino)
+                or not stat.S_ISDIR(current.st_mode)
+            ):
+                raise _error(
+                    "VISUAL_RUN_STORE_IDENTITY",
+                    "pinned research authority changed during the locked operation",
+                )
+
+        try:
+            lease = AuthorityLease(validate)
+            lease()
+            yield lease
+            lease._finish()
+        finally:
+            fcntl.flock(descriptor, fcntl.LOCK_UN)
+    finally:
+        os.close(descriptor)
+
+
+class VisualRunStore:
+    """One run-bound append-only event log.
+
+    The root is supplied by the trusted viewer/coordinator. The store never accepts a path from an
+    event. Sequence allocation and idempotency checks happen under one kernel file lock.
+    """
+
+    def __init__(
+        self,
+        root: Path,
+        run_id: str,
+        *,
+        parent_fd: int | None = None,
+        root_fd: int | None = None,
+    ) -> None:
+        self.root = Path(root)
+        self.run_id = run_id
+        self._root_fd = -1
+        self._artifact_fd = -1
+        if parent_fd is not None and root_fd is not None:
+            raise ValueError("visual store accepts one pinned directory boundary")
+        if root_fd is not None:
+            self._root_fd = os.dup(root_fd)
+        elif parent_fd is not None:
+            if self.root.name != str(self.root) or self.root.name in {"", ".", ".."}:
+                raise ValueError("descriptor-relative visual store root must be one name")
+            with visual_authority_lock(parent_fd, exclusive=True) as validate_authority:
+                validate_authority()
+                try:
+                    os.mkdir(self.root.name, mode=0o700, dir_fd=parent_fd)
+                    os.fsync(parent_fd)
+                except FileExistsError:
+                    pass
+                validate_authority()
+                self._root_fd = os.open(
+                    self.root.name,
+                    os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW,
+                    dir_fd=parent_fd,
+                )
+        else:
+            parent = self.root.parent
+            parent.mkdir(mode=0o700, parents=True, exist_ok=True)
+            parent_fd = os.open(
+                parent,
+                os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW,
+            )
+            try:
+                with visual_authority_lock(parent_fd, exclusive=True) as validate_authority:
+                    validate_authority()
+                    try:
+                        os.mkdir(self.root.name, mode=0o700, dir_fd=parent_fd)
+                        os.fsync(parent_fd)
+                    except FileExistsError:
+                        pass
+                    validate_authority()
+                    self._root_fd = os.open(
+                        self.root.name,
+                        os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW,
+                        dir_fd=parent_fd,
+                    )
+            finally:
+                os.close(parent_fd)
+            root_info = os.fstat(self._root_fd)
+            path_info = self.root.stat(follow_symlinks=False)
+            if (root_info.st_dev, root_info.st_ino) != (path_info.st_dev, path_info.st_ino):
+                root_descriptor, self._root_fd = self._root_fd, -1
+                os.close(root_descriptor)
+                raise _error("VISUAL_RUN_STORE_IDENTITY", "store root changed while it was opened")
+        if not stat.S_ISDIR(os.fstat(self._root_fd).st_mode):
+            root_descriptor, self._root_fd = self._root_fd, -1
+            os.close(root_descriptor)
+            raise _error("VISUAL_RUN_STORE_IDENTITY", "store root must be a real directory")
+        self.artifact_root = self.root / ARTIFACT_DIR
+        try:
+            os.mkdir(ARTIFACT_DIR, mode=0o700, dir_fd=self._root_fd)
+            os.fsync(self._root_fd)
+        except FileExistsError:
+            pass
+        try:
+            self._artifact_fd = os.open(
+                ARTIFACT_DIR,
+                os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW,
+                dir_fd=self._root_fd,
+            )
+        except OSError as exc:
+            root_descriptor, self._root_fd = self._root_fd, -1
+            os.close(root_descriptor)
+            raise _error(
+                "VISUAL_RUN_STORE_IDENTITY", "artifact root must be a real directory"
+            ) from exc
+        if not stat.S_ISDIR(os.fstat(self._artifact_fd).st_mode):
+            artifact_descriptor, self._artifact_fd = self._artifact_fd, -1
+            root_descriptor, self._root_fd = self._root_fd, -1
+            os.close(artifact_descriptor)
+            os.close(root_descriptor)
+            raise _error("VISUAL_RUN_STORE_IDENTITY", "artifact root must be a real directory")
+        self.log_path = self.root / LOG_NAME
+        self.lock_path = self.root / LOCK_NAME
+
+    def close(self) -> None:
+        """Release the pinned directory descriptors; repeated calls are harmless."""
+
+        artifact_fd = getattr(self, "_artifact_fd", -1)
+        root_fd = getattr(self, "_root_fd", -1)
+        self._artifact_fd = -1
+        self._root_fd = -1
+        if artifact_fd >= 0:
+            os.close(artifact_fd)
+        if root_fd >= 0:
+            os.close(root_fd)
+
+    def __del__(self) -> None:
+        try:
+            self.close()
+        except OSError:
+            pass
+
+    def _open_log(self, flags: int, *, create: bool = False) -> tuple[int, bool]:
+        common = flags | os.O_CLOEXEC | os.O_NOFOLLOW
+        if not create:
+            descriptor = os.open(LOG_NAME, common, dir_fd=self._root_fd)
+            try:
+                _regular_descriptor(descriptor, LOG_NAME)
+            except BaseException:
+                os.close(descriptor)
+                raise
+            return descriptor, False
+        try:
+            descriptor = os.open(
+                LOG_NAME,
+                common | os.O_CREAT | os.O_EXCL,
+                0o600,
+                dir_fd=self._root_fd,
+            )
+            try:
+                _regular_descriptor(descriptor, LOG_NAME)
+            except BaseException:
+                os.close(descriptor)
+                raise
+            return descriptor, True
+        except FileExistsError:
+            descriptor = os.open(LOG_NAME, common, dir_fd=self._root_fd)
+            try:
+                _regular_descriptor(descriptor, LOG_NAME)
+            except BaseException:
+                os.close(descriptor)
+                raise
+            return descriptor, False
+
+    def _parse_events(self, raw: bytes) -> tuple[VisualRunEvent, ...]:
+        if raw and not raw.endswith(b"\n"):
+            raise _error("VISUAL_RUN_LOG_TRUNCATED", "visual run log has a partial final record")
+        lines = raw.splitlines()
+        if len(lines) > MAX_EVENTS:
+            raise _error("VISUAL_RUN_LOG_LIMIT", "visual run log exceeds its event ceiling")
+        events: list[VisualRunEvent] = []
+        seen_keys: dict[str, bytes] = {}
+        for index, line in enumerate(lines, 1):
+            event = parse_event(line)
+            if event.run_id != self.run_id:
+                raise _error("VISUAL_RUN_RUN_MISMATCH", "event belongs to another run")
+            if event.sequence != index:
+                raise _error("VISUAL_RUN_SEQUENCE_GAP", "event sequence is not contiguous")
+            encoded = canonical_draft_bytes(VisualRunDraft.from_event(event))
+            previous = seen_keys.setdefault(event.idempotency_key, encoded)
+            if previous != encoded:
+                raise _error("VISUAL_RUN_IDEMPOTENCY_CONFLICT", "idempotency key was reused")
+            events.append(event)
+        return tuple(events)
+
+    def _events_from_descriptor(self, descriptor: int) -> tuple[VisualRunEvent, ...]:
+        _regular_descriptor(descriptor, LOG_NAME)
+        raw = _read_descriptor(
+            descriptor,
+            maximum=MAX_LOG_BYTES,
+            code="VISUAL_RUN_LOG_LIMIT",
+            message="visual run log exceeds its byte ceiling",
+        )
+        return self._parse_events(raw)
+
+    def _read_events(self) -> tuple[VisualRunEvent, ...]:
+        try:
+            descriptor, _ = self._open_log(os.O_RDONLY)
+        except FileNotFoundError:
+            return ()
+        try:
+            return self._events_from_descriptor(descriptor)
+        finally:
+            os.close(descriptor)
+
+    def events(self) -> tuple[VisualRunEvent, ...]:
+        with _locked(self._root_fd):
+            return self._read_events()
+
+    def event_for_key(self, key: str) -> VisualRunEvent | None:
+        """Return the durable result for an idempotent command key, when present."""
+
+        with _locked(self._root_fd):
+            return next(
+                (event for event in self._read_events() if event.idempotency_key == key), None
+            )
+
+    def append(self, draft: VisualRunDraft) -> VisualRunEvent:
+        if draft.run_id != self.run_id:
+            raise _error("VISUAL_RUN_RUN_MISMATCH", "draft belongs to another run")
+        # Validate the complete draft byte bound before comparing its command identity.
+        canonical_draft_bytes(draft)
+        command_bytes = _command_bytes(draft)
+        with _locked(self._root_fd):
+            descriptor, created = self._open_log(os.O_RDWR | os.O_APPEND, create=True)
+            try:
+                existing = self._events_from_descriptor(descriptor)
+                for event in existing:
+                    if event.idempotency_key != draft.idempotency_key:
+                        continue
+                    if _command_bytes(VisualRunDraft.from_event(event)) == command_bytes:
+                        return event
+                    raise _error("VISUAL_RUN_IDEMPOTENCY_CONFLICT", "idempotency key was reused")
+                projection_reserved = draft.object_id == DOCUMENT_HEADER_CELL_ID or (
+                    draft.object_id == LEGACY_DOCUMENT_HEADER_CELL_ID_PREFIX
+                    or draft.object_id.startswith(f"{LEGACY_DOCUMENT_HEADER_CELL_ID_PREFIX}-")
+                )
+                if draft.event_type == "cell.created" and projection_reserved:
+                    raise _error(
+                        "VISUAL_RUN_RESERVED_CELL_ID",
+                        "cell id is reserved for projected document structure",
+                    )
+                if len(existing) >= MAX_EVENTS:
+                    raise _error("VISUAL_RUN_LOG_LIMIT", "visual run log exceeds its event ceiling")
+                event = draft.with_sequence(len(existing) + 1)
+                # Syntax-valid is not enough: invalid transitions never enter the durable log.
+                from mostlyright.data_harness.visual_run.materialize import notebook_bytes
+                from mostlyright.data_harness.visual_run.reducer import apply_event, reduce_events
+
+                before = reduce_events(existing)
+                prospective = apply_event(before, event)
+                prospective_cell = prospective.cells.get(event.object_id)
+                if (
+                    event.actor == "agent"
+                    and event.event_type
+                    in {"cell.source_extended", "cell.source_revised", "cell.completed"}
+                    and prospective_cell is not None
+                    and prospective_cell.cell_type == "markdown"
+                ):
+                    from mostlyright.data_harness.visual_run.authoring import (
+                        validate_agent_authoring_event,
+                    )
+
+                    validate_agent_authoring_event(before, event, prospective, self.read_artifact)
+                try:
+                    notebook_bytes(prospective, self.read_artifact)
+                except ValueError as exc:
+                    raise _error(
+                        "VISUAL_RUN_PROJECTION_LIMIT",
+                        "event would exceed the materialized research notebook limit",
+                    ) from exc
+                encoded = canonical_event_bytes(event)
+                current_size = os.fstat(descriptor).st_size
+                if current_size + len(encoded) > MAX_LOG_BYTES:
+                    raise _error("VISUAL_RUN_LOG_LIMIT", "visual run log exceeds its byte ceiling")
+                written = 0
+                while written < len(encoded):
+                    progress = os.write(descriptor, encoded[written:])
+                    if progress <= 0:
+                        raise _error("VISUAL_RUN_LOG_WRITE", "event log write made no progress")
+                    written += progress
+                os.fsync(descriptor)
+                if created:
+                    os.fsync(self._root_fd)
+                return event
+            finally:
+                os.close(descriptor)
+
+    def page_after(self, cursor: int, *, limit: int = 256) -> EventPage:
+        if isinstance(cursor, bool) or not isinstance(cursor, int) or cursor < 0:
+            raise _error("VISUAL_RUN_CURSOR_INVALID", "cursor must be a non-negative integer")
+        if isinstance(limit, bool) or not isinstance(limit, int) or not 1 <= limit <= 1024:
+            raise _error("VISUAL_RUN_CURSOR_INVALID", "limit is outside the closed range")
+        events = self.events()
+        if cursor > len(events):
+            raise _error("VISUAL_RUN_CURSOR_FUTURE", "cursor is beyond the durable log")
+        selected = events[cursor : cursor + limit]
+        next_cursor = cursor + len(selected)
+        return EventPage(selected, next_cursor, next_cursor < len(events))
+
+    def install_artifact(self, reference: ArtifactReference, content: bytes) -> Path:
+        if not isinstance(content, bytes):
+            raise _error("VISUAL_RUN_ARTIFACT_INVALID", "artifact content must be bytes")
+        if len(content) != reference.byte_size:
+            raise _error("VISUAL_RUN_ARTIFACT_SIZE", "artifact byte size differs from reference")
+        if hashlib.sha256(content).hexdigest() != reference.sha256:
+            raise _error("VISUAL_RUN_ARTIFACT_DIGEST", "artifact digest differs from reference")
+        target = self.artifact_root / reference.sha256
+        with _locked(self._root_fd):
+            try:
+                existing_fd = os.open(
+                    reference.sha256,
+                    os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW,
+                    dir_fd=self._artifact_fd,
+                )
+            except FileNotFoundError:
+                existing_fd = -1
+            if existing_fd >= 0:
+                try:
+                    _regular_descriptor(existing_fd, reference.sha256)
+                    existing = _read_descriptor(
+                        existing_fd,
+                        maximum=reference.byte_size,
+                        code="VISUAL_RUN_ARTIFACT_SIZE",
+                        message="installed artifact size differs",
+                    )
+                finally:
+                    os.close(existing_fd)
+                if existing != content:
+                    raise _error("VISUAL_RUN_ARTIFACT_CONFLICT", "artifact digest path conflicts")
+                return target
+            temporary = f".{reference.sha256}.{secrets.token_hex(12)}.tmp"
+            descriptor = os.open(
+                temporary,
+                os.O_RDWR | os.O_CREAT | os.O_EXCL | os.O_CLOEXEC | os.O_NOFOLLOW,
+                0o400,
+                dir_fd=self._artifact_fd,
+            )
+            try:
+                try:
+                    written = 0
+                    while written < len(content):
+                        progress = os.write(descriptor, content[written:])
+                        if progress <= 0:
+                            raise _error(
+                                "VISUAL_RUN_ARTIFACT_WRITE", "artifact write made no progress"
+                            )
+                        written += progress
+                    os.fsync(descriptor)
+                    installed = _read_descriptor(
+                        descriptor,
+                        maximum=reference.byte_size,
+                        code="VISUAL_RUN_ARTIFACT_SIZE",
+                        message="staged artifact size differs",
+                    )
+                    if installed != content:
+                        raise _error("VISUAL_RUN_ARTIFACT_DIGEST", "staged artifact differs")
+                finally:
+                    os.close(descriptor)
+            except BaseException:
+                try:
+                    os.unlink(temporary, dir_fd=self._artifact_fd)
+                except FileNotFoundError:
+                    pass
+                raise
+            try:
+                os.replace(
+                    temporary,
+                    reference.sha256,
+                    src_dir_fd=self._artifact_fd,
+                    dst_dir_fd=self._artifact_fd,
+                )
+                os.fsync(self._artifact_fd)
+            except BaseException:
+                try:
+                    os.unlink(temporary, dir_fd=self._artifact_fd)
+                except FileNotFoundError:
+                    pass
+                raise
+        return target
+
+    def read_artifact(self, reference: ArtifactReference) -> bytes:
+        try:
+            descriptor = os.open(
+                reference.sha256,
+                os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW,
+                dir_fd=self._artifact_fd,
+            )
+        except FileNotFoundError as exc:
+            raise _error("VISUAL_RUN_ARTIFACT_MISSING", "artifact is not installed") from exc
+        try:
+            _regular_descriptor(descriptor, reference.sha256)
+            content = _read_descriptor(
+                descriptor,
+                maximum=reference.byte_size,
+                code="VISUAL_RUN_ARTIFACT_SIZE",
+                message="installed artifact size differs",
+            )
+        finally:
+            os.close(descriptor)
+        if len(content) != reference.byte_size:
+            raise _error("VISUAL_RUN_ARTIFACT_SIZE", "installed artifact size differs")
+        if hashlib.sha256(content).hexdigest() != reference.sha256:
+            raise _error("VISUAL_RUN_ARTIFACT_DIGEST", "installed artifact digest differs")
+        return content
+
+
+def snapshot_json(page: EventPage) -> bytes:
+    return (
+        json.dumps(
+            {
+                "cursor": page.cursor,
+                "events": [event.to_mapping() for event in page.events],
+                "truncated": page.truncated,
+            },
+            ensure_ascii=False,
+            allow_nan=False,
+            sort_keys=True,
+            separators=(",", ":"),
+        )
+        + "\n"
+    ).encode("utf-8")