release/v1/scripts/compact_logbook_run_attachments.py

#!/usr/bin/env python3
"""Compact source/config snapshots embedded by ``trackio logbook run``.

Trackio run cells preserve the command, timing, captured output, and copies of
input source/config files. Repeating a large runner on several pages can make
an otherwise small logbook difficult for agents to read. This utility replaces
only titled four-backtick Python/JSON attachments in code cells that also have
a bash command fence. It leaves every cell header, command, output, artifact,
figure, and non-code cell byte-for-byte unchanged.

The replacement retains the attachment title, language, UTF-8 byte count, and
SHA-256 of the exact fenced payload. Writes use a temporary file in the page's
directory followed by ``os.replace``; running the utility again is a no-op.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import stat
import tempfile
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Sequence


CELL_HEADER_RE = re.compile(
    r"^---\r?\n<!-- trackio-cell\r?\n(?P<metadata>.*?)\r?\n-->\r?\n",
    flags=re.MULTILINE | re.DOTALL,
)
ATTACHMENT_RE = re.compile(
    r"^````(?P<language>python|json) title=(?P<title>[^`\r\n]+)\r?\n"
    r"(?P<payload>.*?)"
    r"^````[ \t]*(?P<closing_newline>\r?\n|$)",
    flags=re.MULTILINE | re.DOTALL,
)


@dataclass(frozen=True)
class Cell:
    """One parsed Trackio cell, including its unmodified serialized header."""

    header: str
    metadata: dict[str, Any]
    body: str


@dataclass(frozen=True)
class AttachmentSummary:
    """Audit data retained for a compacted attachment."""

    title: str
    language: str
    byte_count: int
    sha256: str


def _parse_page(document: str, path: Path) -> tuple[str, list[Cell]]:
    matches = list(CELL_HEADER_RE.finditer(document))
    if not matches:
        return document, []

    cells: list[Cell] = []
    for index, match in enumerate(matches):
        try:
            metadata = json.loads(match.group("metadata"))
        except json.JSONDecodeError as error:
            raise ValueError(
                f"invalid Trackio cell metadata in {path}: {error}"
            ) from error
        if not isinstance(metadata, dict):
            raise ValueError(f"non-object Trackio cell metadata in {path}")

        body_end = matches[index + 1].start() if index + 1 < len(matches) else len(document)
        cells.append(
            Cell(
                header=match.group(0),
                metadata=metadata,
                body=document[match.end() : body_end],
            )
        )

    return document[: matches[0].start()], cells


def _serialize_page(prefix: str, cells: Sequence[Cell]) -> str:
    return prefix + "".join(cell.header + cell.body for cell in cells)


def _is_source_or_config(language: str, title: str) -> bool:
    suffix = Path(title).suffix.lower()
    return (language == "python" and suffix == ".py") or (
        language == "json" and suffix == ".json"
    )


def _compact_code_body(body: str) -> tuple[str, list[AttachmentSummary]]:
    # A bash fence distinguishes a generated run cell from a titled code block
    # that someone may have intentionally authored in a Markdown cell body.
    if "````bash" not in body:
        return body, []

    summaries: list[AttachmentSummary] = []

    def replace(match: re.Match[str]) -> str:
        language = match.group("language")
        title = match.group("title").strip()
        if not _is_source_or_config(language, title):
            return match.group(0)

        payload_bytes = match.group("payload").encode("utf-8")
        summary = AttachmentSummary(
            title=title,
            language=language,
            byte_count=len(payload_bytes),
            sha256=hashlib.sha256(payload_bytes).hexdigest(),
        )
        summaries.append(summary)
        note = (
            "> Trackio run attachment compacted: "
            f"`{summary.title}` ({summary.language}, "
            f"{summary.byte_count} UTF-8 bytes; embedded-payload SHA-256 "
            f"`{summary.sha256}`). Command, timing, and captured output are "
            "preserved in this cell."
        )
        return note + match.group("closing_newline")

    return ATTACHMENT_RE.sub(replace, body), summaries


def _compact_page(document: str, path: Path) -> tuple[str, list[AttachmentSummary]]:
    prefix, before_cells = _parse_page(document, path)
    summaries: list[AttachmentSummary] = []
    after_cells: list[Cell] = []

    for cell in before_cells:
        if cell.metadata.get("type") == "code":
            body, compacted = _compact_code_body(cell.body)
            summaries.extend(compacted)
        else:
            body = cell.body
        after_cells.append(Cell(header=cell.header, metadata=cell.metadata, body=body))

    compacted_document = _serialize_page(prefix, after_cells)
    after_prefix, reparsed_cells = _parse_page(compacted_document, path)

    if after_prefix != prefix or len(reparsed_cells) != len(before_cells):
        raise AssertionError(f"Trackio page structure changed while compacting {path}")
    for before, after in zip(before_cells, reparsed_cells, strict=True):
        if before.header != after.header or before.metadata != after.metadata:
            raise AssertionError(f"Trackio cell metadata changed while compacting {path}")
        if before.metadata.get("type") != "code" and before.body != after.body:
            raise AssertionError(f"non-code cell changed while compacting {path}")

    # Fail closed if the transformation is not idempotent.
    second_pass_cells: list[Cell] = []
    for cell in reparsed_cells:
        body, repeated = (
            _compact_code_body(cell.body)
            if cell.metadata.get("type") == "code"
            else (cell.body, [])
        )
        if repeated:
            raise AssertionError(f"compaction is not idempotent for {path}")
        second_pass_cells.append(Cell(cell.header, cell.metadata, body))
    if _serialize_page(after_prefix, second_pass_cells) != compacted_document:
        raise AssertionError(f"second pass changed {path}")

    return compacted_document, summaries


def _atomic_write(path: Path, content: bytes) -> None:
    original_mode = stat.S_IMODE(path.stat().st_mode)
    descriptor, temporary_name = tempfile.mkstemp(
        dir=path.parent, prefix=f".{path.name}.", suffix=".tmp"
    )
    temporary_path = Path(temporary_name)
    try:
        with os.fdopen(descriptor, "wb") as handle:
            handle.write(content)
            handle.flush()
            os.fsync(handle.fileno())
        os.chmod(temporary_path, original_mode)
        os.replace(temporary_path, path)
        directory_descriptor = os.open(path.parent, os.O_RDONLY)
        try:
            os.fsync(directory_descriptor)
        finally:
            os.close(directory_descriptor)
    finally:
        if temporary_path.exists():
            temporary_path.unlink()


def _page_paths(logbook_root: Path) -> list[Path]:
    manifest_path = logbook_root / "logbook.json"
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    declared: list[Path] = []

    def visit(node: dict[str, Any]) -> None:
        declared.append(logbook_root / node["file"])
        for child in node.get("children", []):
            visit(child)

    visit(manifest["root"])
    missing = [path for path in declared if not path.is_file()]
    if missing:
        raise FileNotFoundError(f"declared Trackio pages are missing: {missing}")

    undeclared = set((logbook_root / "pages").rglob("*.md")) - set(declared)
    if undeclared:
        raise ValueError(f"undeclared Trackio pages found: {sorted(undeclared)}")
    return declared


def compact_logbook(logbook_root: Path, *, write: bool) -> dict[str, Any]:
    """Compact every declared page and return a machine-readable audit summary."""

    root = logbook_root.resolve()
    pages: list[dict[str, Any]] = []
    total_attachments = 0

    for path in _page_paths(root):
        before_bytes = path.read_bytes()
        before_document = before_bytes.decode("utf-8")
        _, before_cells = _parse_page(before_document, path)
        after_document, summaries = _compact_page(before_document, path)
        after_bytes = after_document.encode("utf-8")
        _, after_cells = _parse_page(after_document, path)

        if write and before_bytes != after_bytes:
            _atomic_write(path, after_bytes)

        before_types = Counter(str(cell.metadata.get("type")) for cell in before_cells)
        after_types = Counter(str(cell.metadata.get("type")) for cell in after_cells)
        if before_types != after_types:
            raise AssertionError(f"cell type counts changed for {path}")

        total_attachments += len(summaries)
        pages.append(
            {
                "path": str(path.relative_to(root)),
                "before_bytes": len(before_bytes),
                "after_bytes": len(after_bytes),
                "bytes_removed": len(before_bytes) - len(after_bytes),
                "cells_before": len(before_cells),
                "cells_after": len(after_cells),
                "cell_types": dict(sorted(before_types.items())),
                "attachments_compacted": len(summaries),
                "attachments": [summary.__dict__ for summary in summaries],
            }
        )

    return {
        "logbook_root": str(root),
        "mode": "write" if write else "dry-run",
        "pages": pages,
        "totals": {
            "before_bytes": sum(page["before_bytes"] for page in pages),
            "after_bytes": sum(page["after_bytes"] for page in pages),
            "bytes_removed": sum(page["bytes_removed"] for page in pages),
            "cells_before": sum(page["cells_before"] for page in pages),
            "cells_after": sum(page["cells_after"] for page in pages),
            "attachments_compacted": total_attachments,
        },
    }


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--logbook-root",
        type=Path,
        default=Path(".trackio/logbook"),
        help="Trackio logbook directory containing logbook.json (default: %(default)s)",
    )
    parser.add_argument(
        "--write",
        action="store_true",
        help="Atomically replace changed page files; otherwise only report a dry run",
    )
    return parser.parse_args()


def main() -> None:
    args = _parse_args()
    summary = compact_logbook(args.logbook_root, write=args.write)
    print(json.dumps(summary, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()