mostlyright-data 0.9.0__py3-none-any.whl

mostlyright/data_harness/visual_run/materialize.py

@@ -0,0 +1,486 @@
+"""Deterministic nbformat materialization from reduced visual-run state."""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import secrets
+import stat
+import sys
+from collections.abc import Callable
+from ctypes import CDLL, c_char_p, c_int, get_errno
+from pathlib import Path
+from typing import Any
+
+from mostlyright.data_harness.canonical import canonical_json_bytes, sha256_bytes
+from mostlyright.data_harness.visual_run.contracts import (
+    DOCUMENT_HEADER_CELL_ID,
+    LEGACY_DOCUMENT_HEADER_CELL_ID_PREFIX,
+)
+from mostlyright.data_harness.visual_run.observations import scrub_for_display
+from mostlyright.data_harness.visual_run.reducer import CellState, ReducedRun
+
+NOTEBOOK_FORMAT = 4
+NOTEBOOK_MINOR = 5
+MAX_MATERIALIZED_NOTEBOOK_BYTES = 8 * 1024 * 1024
+DOCUMENT_HEADER_SOURCE = (
+    "# Research record\n\n"
+    "This document explains the source choice, dataset rules, build steps, and checks.\n"
+)
+_MARKDOWN_FENCE = re.compile(r"^ {0,3}(`{3,}|~{3,})\s*([^`]*)$")
+_MARKDOWN_H1 = re.compile(r"^#\s+")
+
+
+def _demote_legacy_h1(source: str) -> str:
+    """Keep historical artifact truth while projecting agent H1 headings as H2 sections."""
+
+    lines = source.splitlines(keepends=True)
+    fence_character: str | None = None
+    projected: list[str] = []
+    for line in lines:
+        body = line.rstrip("\r\n")
+        if fence_character is not None:
+            closing = re.fullmatch(rf" {{0,3}}{re.escape(fence_character)}{{3,}}\s*", body)
+            if closing:
+                fence_character = None
+            projected.append(line)
+            continue
+        fence = _MARKDOWN_FENCE.match(body)
+        if fence:
+            fence_character = fence.group(1)[0]
+            projected.append(line)
+            continue
+        projected.append("#" + line if _MARKDOWN_H1.match(body) else line)
+    return "".join(projected)
+
+
+def _document_header_cell() -> dict[str, Any]:
+    """Return the fixed opening cell shared by every research notebook projection."""
+
+    return {
+        "cell_type": "markdown",
+        "id": DOCUMENT_HEADER_CELL_ID,
+        "metadata": {"mostlyright": {"document_header": True}},
+        "source": DOCUMENT_HEADER_SOURCE.splitlines(keepends=True),
+    }
+
+
+def _artifact_output_note(stage: str) -> str:
+    """Describe a stored output without exposing its internal artifact identifier."""
+
+    label = " ".join(stage.replace("_", "-").split("-")).strip() or "step"
+    return f"{label.capitalize()} output is stored with this run.\n"
+
+
+def _cell(cell: CellState, artifact_reader: Callable[[Any], bytes] | None) -> dict[str, Any]:
+    mostlyright: dict[str, Any] = {
+        "object_version": cell.version,
+        "phase": cell.phase,
+        "stage": cell.stage,
+        "status": cell.status,
+    }
+    if cell.status in {"writing", "executing"}:
+        mostlyright["activity"] = {
+            "schema_version": "mostlyright.research-cell-activity.v1",
+            "state": cell.status,
+        }
+    metadata = {"mostlyright": mostlyright}
+    source = []
+    if cell.source is not None and artifact_reader is not None:
+        source = artifact_reader(cell.source).decode("utf-8").splitlines(keepends=True)
+    if cell.cell_type == "markdown":
+        source = _demote_legacy_h1("".join(source)).splitlines(keepends=True)
+        return {
+            "cell_type": "markdown",
+            "id": cell.cell_id,
+            "metadata": metadata,
+            "source": source,
+        }
+    outputs = []
+    for reference in cell.outputs:
+        output: Any = None
+        if artifact_reader is not None:
+            try:
+                output = json.loads(artifact_reader(reference))
+            except (UnicodeDecodeError, json.JSONDecodeError):
+                output = None
+        if not isinstance(output, dict) or output.get("output_type") not in {
+            "stream",
+            "display_data",
+            "execute_result",
+            "error",
+        }:
+            output = {
+                "data": {
+                    "application/vnd.mostlyright.artifact-reference.v1+json": (
+                        reference.to_mapping()
+                    ),
+                    "text/plain": [_artifact_output_note(cell.stage)],
+                },
+                "metadata": {},
+                "output_type": "display_data",
+            }
+        outputs.append(output)
+    return {
+        "cell_type": "code",
+        # Jupyter's live prompt is represented by a null execution count plus the exact
+        # persisted ``executing`` activity.  The requested count becomes visible only when
+        # the cell completes, so a running cell renders as ``In [*]`` rather than falsely
+        # claiming the execution already finished.
+        "execution_count": None if cell.status == "executing" else cell.execution_count,
+        "id": cell.cell_id,
+        "metadata": metadata,
+        "outputs": outputs,
+        "source": source,
+    }
+
+
+def _observation_cell(
+    index: int,
+    reference: Any,
+    artifact_reader: Callable[[Any], bytes] | None,
+) -> dict[str, Any]:
+    value: Any = None
+    if artifact_reader is not None:
+        try:
+            value = json.loads(artifact_reader(reference))
+        except (UnicodeDecodeError, json.JSONDecodeError):
+            value = None
+    if isinstance(value, dict):
+        value, redacted = scrub_for_display(value)
+        assert isinstance(value, dict)
+        source_observation = (
+            reference.artifact_schema == "mostlyright.source-observation.v1"
+            and value.get("schema_version") == "mostlyright.source-observation.v1"
+        )
+        display = value.get("display", reference.display)
+        if not source_observation and display in {
+            "aggregate_only",
+            "schema_only",
+            "summary",
+            "redacted",
+        }:
+            value["sample"] = []
+            value["operation_parameters"] = {}
+            value["source_coordinates"] = []
+            coverage = value.get("coverage")
+            if isinstance(coverage, dict):
+                coverage["representative_window"] = None
+        if redacted and not source_observation:
+            value["display"] = "schema_only"
+            value["sample"] = []
+            value["operation_parameters"] = {}
+            value["source_coordinates"] = []
+            coverage = value.get("coverage")
+            if isinstance(coverage, dict):
+                coverage["representative_window"] = None
+        mime = (
+            "application/vnd.mostlyright.source-observation+json"
+            if source_observation
+            else "application/vnd.mostlyright.stage-observation+json"
+        )
+        label = "source" if source_observation else value.get("operation", "stage")
+        rows = value.get("row_count") if source_observation else value.get("output_rows")
+        data = {
+            mime: value,
+            "text/plain": [
+                f"Mostly Right {label} observation "
+                f"({rows if rows is not None else 'unknown'} rows)\n"
+            ],
+        }
+    else:
+        data = {
+            "application/vnd.mostlyright.artifact-reference.v1+json": {
+                "sha256": reference.sha256,
+                "display": reference.display,
+            },
+            "text/plain": ["Mostly Right observation artifact\n"],
+        }
+    return {
+        "cell_type": "code",
+        "execution_count": None,
+        "id": f"mr-observation-{index:04d}-{reference.sha256[:24]}",
+        "metadata": {
+            "mostlyright": {
+                "artifact_sha256": reference.sha256,
+                "passive": True,
+                "schema_version": "mostlyright.passive-observation-cell.v1",
+            }
+        },
+        "outputs": [{"data": data, "metadata": {}, "output_type": "display_data"}],
+        "source": [],
+    }
+
+
+def notebook_value(
+    state: ReducedRun, artifact_reader: Callable[[Any], bytes] | None = None
+) -> dict[str, Any]:
+    if any(key != cell.cell_id for key, cell in state.cells.items()):
+        raise ValueError("visual-run state has a cell id that differs from its mapping key")
+    # Refuse before calling an artifact reader: references are durable, authenticated size facts,
+    # so they provide the only way to prevent a huge source/output from being allocated first and
+    # rejected only after serialization.
+    artifact_bytes = sum(
+        reference.byte_size
+        for cell in state.cells.values()
+        for reference in ((cell.source,) if cell.source is not None else ()) + cell.outputs
+    ) + sum(reference.byte_size for reference in state.observations)
+    if artifact_bytes > MAX_MATERIALIZED_NOTEBOOK_BYTES:
+        raise ValueError("visual-run artifacts exceed the materialized notebook byte limit")
+    authored_ids = {cell.cell_id for cell in state.cells.values()}
+    legacy_header_id = LEGACY_DOCUMENT_HEADER_CELL_ID_PREFIX
+    suffix = 2
+    while legacy_header_id in authored_ids:
+        legacy_header_id = f"{LEGACY_DOCUMENT_HEADER_CELL_ID_PREFIX}-{suffix}"
+        suffix += 1
+    cells = [_document_header_cell()]
+    for cell in state.cells.values():
+        projected = _cell(cell, artifact_reader)
+        if projected["id"] == DOCUMENT_HEADER_CELL_ID:
+            projected["id"] = legacy_header_id
+        cells.append(projected)
+    cells.extend(
+        _observation_cell(index, reference, artifact_reader)
+        for index, reference in enumerate(state.observations, 1)
+    )
+    return {
+        "cells": cells,
+        "metadata": {
+            "kernelspec": {
+                "display_name": "Python 3",
+                "language": "python",
+                "name": "python3",
+            },
+            "language_info": {"name": "python", "version": "3.13"},
+            "mostlyright": {
+                "cursor": state.cursor,
+                "lifecycle": state.lifecycle,
+                "run_id": state.run_id,
+                "schema_version": "mostlyright.research-notebook.v2",
+            },
+        },
+        "nbformat": NOTEBOOK_FORMAT,
+        "nbformat_minor": NOTEBOOK_MINOR,
+    }
+
+
+def notebook_bytes(
+    state: ReducedRun, artifact_reader: Callable[[Any], bytes] | None = None
+) -> bytes:
+    content = canonical_json_bytes(notebook_value(state, artifact_reader))
+    if len(content) > MAX_MATERIALIZED_NOTEBOOK_BYTES:
+        raise ValueError("generated notebook exceeds the visual materializer byte limit")
+    return content
+
+
+def notebook_digest(
+    state: ReducedRun, artifact_reader: Callable[[Any], bytes] | None = None
+) -> str:
+    return sha256_bytes(notebook_bytes(state, artifact_reader))
+
+
+def materialize_notebook(
+    path: Path,
+    state: ReducedRun,
+    artifact_reader: Callable[[Any], bytes] | None = None,
+    *,
+    directory_fd: int | None = None,
+) -> str:
+    path = Path(path)
+    content = notebook_bytes(state, artifact_reader)
+    owned_directory = directory_fd is None
+    if directory_fd is None:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        directory_fd = os.open(
+            path.parent,
+            os.O_RDONLY | os.O_DIRECTORY | getattr(os, "O_NOFOLLOW", 0),
+        )
+    try:
+        existing_raw, existing_identity = _read_existing(directory_fd, path.name)
+        if existing_raw is not None:
+            try:
+                existing = json.loads(existing_raw)
+            except (UnicodeError, json.JSONDecodeError) as exc:
+                raise ValueError("existing notebook is not a valid derived notebook") from exc
+            mostlyright = existing.get("metadata", {}).get("mostlyright", {})
+            marker = mostlyright.get("schema_version")
+            seeded = mostlyright.get("phase") == "ready_for_research" and all(
+                cell.get("metadata", {}).get("mostlyright", {}).get("seed") is True
+                for cell in existing.get("cells", ())
+            )
+            if marker != "mostlyright.research-notebook.v2" and not seeded:
+                raise ValueError(
+                    "refusing to replace a notebook not owned by the visual materializer"
+                )
+            if existing_raw == content:
+                return sha256_bytes(content)
+        temporary = f".{path.name}.{secrets.token_hex(12)}.tmp"
+        descriptor = os.open(
+            temporary,
+            os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_CLOEXEC | os.O_NOFOLLOW,
+            0o600,
+            dir_fd=directory_fd,
+        )
+        try:
+            written = 0
+            while written < len(content):
+                progress = os.write(descriptor, content[written:])
+                if progress <= 0:
+                    raise OSError("notebook write made no progress")
+                written += progress
+            os.fsync(descriptor)
+        finally:
+            os.close(descriptor)
+        _install_if_unchanged(
+            directory_fd,
+            temporary,
+            path.name,
+            existing_identity,
+            sha256_bytes(existing_raw) if existing_raw is not None else None,
+        )
+        os.fsync(directory_fd)
+        return sha256_bytes(content)
+    finally:
+        if owned_directory:
+            os.close(directory_fd)
+
+
+def _file_identity(observed: os.stat_result) -> tuple[int, int, int, int, int, int, int]:
+    return (
+        observed.st_dev,
+        observed.st_ino,
+        observed.st_mode,
+        observed.st_nlink,
+        observed.st_size,
+        observed.st_mtime_ns,
+        observed.st_ctime_ns,
+    )
+
+
+def _read_existing(
+    directory_fd: int, name: str
+) -> tuple[bytes | None, tuple[int, int, int, int, int, int, int] | None]:
+    try:
+        descriptor = os.open(
+            name,
+            os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW,
+            dir_fd=directory_fd,
+        )
+    except FileNotFoundError:
+        return None, None
+    try:
+        observed = os.fstat(descriptor)
+        if not stat.S_ISREG(observed.st_mode) or observed.st_nlink != 1:
+            raise ValueError("existing notebook must be a regular single-link file")
+        if observed.st_size > MAX_MATERIALIZED_NOTEBOOK_BYTES:
+            raise ValueError("existing notebook exceeds the visual materializer byte limit")
+        chunks: list[bytes] = []
+        remaining = MAX_MATERIALIZED_NOTEBOOK_BYTES + 1
+        while remaining and (chunk := os.read(descriptor, min(64 * 1024, remaining))):
+            chunks.append(chunk)
+            remaining -= len(chunk)
+        raw = b"".join(chunks)
+        after = os.fstat(descriptor)
+        if len(raw) > MAX_MATERIALIZED_NOTEBOOK_BYTES or _file_identity(after) != _file_identity(
+            observed
+        ):
+            raise ValueError("existing notebook changed while it was being read")
+        return raw, _file_identity(observed)
+    finally:
+        os.close(descriptor)
+
+
+def _require_destination_identity(
+    directory_fd: int,
+    name: str,
+    expected: tuple[int, int, int, int, int, int, int] | None,
+    expected_sha256: str | None,
+) -> None:
+    observed_raw, observed_identity = _read_existing(directory_fd, name)
+    # Linux updates ctime during RENAME_EXCHANGE, so it cannot be compared after the atomic swap.
+    # Device/inode/mode/link identity plus the exact bounded content digest proves that the
+    # displaced object is the one we admitted, even if an attacker tried to spoof timestamps.
+    if (
+        observed_identity is None
+        or expected is None
+        or observed_identity[:4] != expected[:4]
+        or observed_raw is None
+        or expected_sha256 is None
+        or sha256_bytes(observed_raw) != expected_sha256
+    ):
+        raise ValueError("existing notebook changed before atomic materialization")
+
+
+def _rename_with_flags(directory_fd: int, source: str, destination: str, flag: int) -> None:
+    libc = CDLL(None, use_errno=True)
+    encoded_source = os.fsencode(source)
+    encoded_destination = os.fsencode(destination)
+    if sys.platform == "darwin":
+        result = libc.renameatx_np(
+            c_int(directory_fd),
+            c_char_p(encoded_source),
+            c_int(directory_fd),
+            c_char_p(encoded_destination),
+            c_int(flag),
+        )
+    else:
+        result = libc.renameat2(
+            c_int(directory_fd),
+            c_char_p(encoded_source),
+            c_int(directory_fd),
+            c_char_p(encoded_destination),
+            c_int(flag),
+        )
+    if result != 0:
+        error = get_errno()
+        raise OSError(error, os.strerror(error), destination)
+
+
+def _install_if_unchanged(
+    directory_fd: int,
+    temporary: str,
+    destination: str,
+    expected: tuple[int, int, int, int, int, int, int] | None,
+    expected_sha256: str | None,
+) -> None:
+    """Atomically install without ever overwriting an unchecked destination object."""
+
+    if expected is None:
+        # Darwin RENAME_EXCL and Linux RENAME_NOREPLACE have different flag values.
+        try:
+            _rename_with_flags(
+                directory_fd,
+                temporary,
+                destination,
+                0x00000004 if sys.platform == "darwin" else 1,
+            )
+        except BaseException:
+            # A failed no-replace leaves our generated file at the temporary name.
+            os.unlink(temporary, dir_fd=directory_fd)
+            raise
+        return
+    # Exchange is the only portable kernel primitive that prevents a substitution between a
+    # userspace identity check and installation.  The displaced object remains recoverable under
+    # our random temporary name until its identity is proven.
+    try:
+        _rename_with_flags(directory_fd, temporary, destination, 0x00000002)
+    except BaseException:
+        # A failed exchange leaves our generated file at the temporary name.
+        os.unlink(temporary, dir_fd=directory_fd)
+        raise
+    try:
+        _require_destination_identity(directory_fd, temporary, expected, expected_sha256)
+    except BaseException as identity_error:
+        # Restore both names atomically; never unlink the displaced, unchecked object.
+        try:
+            _rename_with_flags(directory_fd, temporary, destination, 0x00000002)
+        except BaseException as rollback_error:
+            # Both objects remain recoverable.  In particular, the unchecked displaced object is
+            # never deleted merely because restoration encountered an operating-system failure.
+            raise RuntimeError(
+                "atomic notebook replacement could not be restored"
+            ) from rollback_error
+        os.unlink(temporary, dir_fd=directory_fd)
+        raise identity_error
+    os.unlink(temporary, dir_fd=directory_fd)