ltx/probe_artifacts.py

from __future__ import annotations

import json
import os
import zipfile
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Iterable


@dataclass(frozen=True)
class ProbeMember:
    """One file included in a returned Probe bundle."""

    role: str
    path: Path
    arcname: str | None = None
    required: bool = True


def _normalized_members(members: Iterable[ProbeMember]) -> list[ProbeMember]:
    out: list[ProbeMember] = []
    seen_arcnames: set[str] = set()
    for member in members:
        path = Path(member.path)
        if not path.exists():
            if member.required:
                raise FileNotFoundError(f"Required Probe member is missing: {path}")
            continue
        if not path.is_file():
            raise ValueError(f"Probe member must be a file: {path}")
        arcname = str(member.arcname or path.name).strip()
        if not arcname or Path(arcname).is_absolute() or ".." in Path(arcname).parts:
            raise ValueError(f"Unsafe Probe archive name: {arcname!r}")
        if arcname in seen_arcnames:
            raise ValueError(f"Duplicate Probe archive name: {arcname}")
        seen_arcnames.add(arcname)
        out.append(ProbeMember(role=str(member.role), path=path, arcname=arcname, required=member.required))
    return out


def build_probe_zip(
    output_path: str | Path,
    *,
    candidate_id: str,
    request_id: str,
    status: str,
    members: Iterable[ProbeMember],
    sha256_file: Callable[[str | Path], str],
    metadata: dict[str, Any] | None = None,
) -> Path:
    """Build one self-describing Probe ZIP without leaking absolute runtime paths."""

    output = Path(output_path)
    output.parent.mkdir(parents=True, exist_ok=True)
    normalized = _normalized_members(members)

    manifest_files = []
    for member in normalized:
        assert member.arcname is not None
        manifest_files.append(
            {
                "role": member.role,
                "name": member.arcname,
                "size_bytes": int(member.path.stat().st_size),
                "sha256": sha256_file(member.path),
            }
        )

    manifest = {
        "schema_version": "ltx25-probe-bundle-v1",
        "candidate_id": str(candidate_id),
        "request_id": str(request_id),
        "status": str(status).upper(),
        "created_utc": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "files": manifest_files,
        "metadata": dict(metadata or {}),
    }

    tmp = output.with_name(f".{output.name}.{os.getpid()}.tmp")
    try:
        with zipfile.ZipFile(tmp, mode="w", compression=zipfile.ZIP_DEFLATED, compresslevel=6) as archive:
            archive.writestr("probe_manifest.json", json.dumps(manifest, indent=2, sort_keys=True) + "\n")
            for member in normalized:
                assert member.arcname is not None
                archive.write(member.path, arcname=member.arcname)
        tmp.replace(output)
    finally:
        try:
            tmp.unlink(missing_ok=True)
        except Exception:
            pass
    return output


def build_request_probe(
    paths: Any,
    *,
    candidate_id: str,
    status: str,
    sha256_file: Callable[[str | Path], str],
    include_run_info: bool = True,
    include_diagnostics: bool = True,
    include_log: bool = True,
    extra_members: Iterable[ProbeMember] = (),
    metadata: dict[str, Any] | None = None,
) -> Path:
    """Bundle the ordinary request evidence files owned by one RequestPaths object."""

    members: list[ProbeMember] = []
    if include_run_info:
        members.append(ProbeMember("run_info", Path(paths.run_info), required=False))
    if include_diagnostics:
        members.append(ProbeMember("diagnostics", Path(paths.diagnostics), required=False))
    if include_log:
        members.append(ProbeMember("request_log", Path(paths.log), required=False))
    members.extend(extra_members)
    return build_probe_zip(
        paths.probe,
        candidate_id=candidate_id,
        request_id=paths.request_id,
        status=status,
        members=members,
        sha256_file=sha256_file,
        metadata=metadata,
    )


def append_probe_path(existing: Any, new_path: str | Path, *, limit: int | None = None) -> list[str]:
    """Append one Probe bundle path while preserving existing gr.File-style history."""

    values: list[str] = []
    if isinstance(existing, (list, tuple)):
        values.extend(str(item) for item in existing if item)
    elif existing:
        values.append(str(existing))

    new_value = str(new_path)
    values = [item for item in values if item != new_value]
    values.append(new_value)
    if limit is not None and int(limit) > 0:
        values = values[-int(limit) :]
    return values


def list_session_probe_paths(worker_root: str | Path, session_id: str, *, limit: int | None = None) -> list[str]:
    """Return existing request Probe ZIPs for one UUID session, oldest-to-newest."""

    session_root = Path(worker_root) / str(session_id or "").strip().lower()
    if not session_root.is_dir():
        return []
    found = [path for path in session_root.glob("*/ltx25_probe_*.zip") if path.is_file()]
    found.sort(key=lambda path: (path.stat().st_mtime_ns, path.name))
    values = [str(path) for path in found]
    if limit is not None and int(limit) > 0:
        values = values[-int(limit) :]
    return values