zcbot/core/artifacts.py

"""Artifact message references and working-dir scoped path resolution.

Files remain the content source of truth. Version-2 refs carry a stable database identity
plus a task-relative path snapshot; version-1 path-only refs remain readable.
"""
from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
from uuid import UUID

ARTIFACT_REF_VERSION = 1
ARTIFACT_REF_CURRENT_VERSION = 2
MAX_ARTIFACTS_PER_MESSAGE = 10
_CONTAINER_ROOT = Path("/workspace")
ARTIFACT_TRASH_DIR = ".zcbot_artifact_trash"


class ArtifactPathError(ValueError):
    """A proposed artifact path is invalid or outside the current working_dir."""


@dataclass(frozen=True)
class ArtifactRef:
    path: str
    label: str = ""
    artifact_id: Optional[UUID] = None
    scope: str = "working_dir"
    version: int = ARTIFACT_REF_VERSION

    def as_dict(self) -> dict:
        out = {
            "version": self.version,
            "scope": self.scope,
            "path": self.path,
        }
        if self.artifact_id:
            out["artifact_id"] = str(self.artifact_id)
        if self.label:
            out["label"] = self.label
        return out


class ToolExecutionResult(str):
    """String-compatible rich result for tools that publish deliverables.

    A few internal scripts and tests call tools directly and historically received a
    plain string.  Subclassing ``str`` preserves that contract while executors can
    still consume the structured artifact metadata.
    """

    content: str
    artifacts: tuple[ArtifactRef, ...]

    def __new__(
        cls,
        content: str,
        artifacts: Iterable[ArtifactRef] = (),
    ) -> "ToolExecutionResult":
        obj = super().__new__(cls, content)
        obj.content = content
        obj.artifacts = tuple(artifacts)
        return obj


def _relative_parts(path: Path) -> tuple[str, ...]:
    return tuple(part for part in path.parts if part not in ("", "."))


def resolve_artifact_path(
    raw_path: str,
    *,
    working_dir: Path,
    user_root: Path,
    require_file: bool = True,
    allow_legacy_user_relative: bool = True,
) -> tuple[Path, str]:
    """Resolve legacy/canonical input and return (absolute, task-relative POSIX path).

    Canonical input is relative to working_dir (``reports/a.pdf``).  For compatibility,
    user-root-relative paths (``<wd>/reports/a.pdf``), container absolute paths under
    ``/workspace`` and host absolute paths inside working_dir are accepted too.
    """
    raw = str(raw_path or "").strip().replace("\\", "/")
    if not raw or "\x00" in raw:
        raise ArtifactPathError("artifact path is empty or contains NUL")

    wd = Path(working_dir).resolve()
    root = Path(user_root).resolve()
    try:
        wd_rel = wd.relative_to(root)
    except ValueError as exc:
        raise ArtifactPathError("working_dir is outside user_root") from exc

    explicit_task_relative = raw.startswith("./")
    p = Path(raw[2:] if explicit_task_relative else raw)
    if raw == "/workspace" or raw.startswith("/workspace/"):
        rest = raw[len("/workspace"):].lstrip("/")
        candidate = root / Path(rest)
    elif p.is_absolute():
        candidate = p
    else:
        parts = _relative_parts(p)
        wd_parts = _relative_parts(wd_rel)
        if (
            allow_legacy_user_relative
            and not explicit_task_relative
            and wd_parts
            and parts[:len(wd_parts)] == wd_parts
        ):
            candidate = root.joinpath(*parts)
        else:
            candidate = wd.joinpath(*parts)

    resolved = candidate.resolve()
    try:
        rel = resolved.relative_to(wd)
    except ValueError as exc:
        raise ArtifactPathError("artifact path escapes working_dir") from exc
    if rel == Path("."):
        raise ArtifactPathError("artifact path must reference a file")
    if require_file and not resolved.is_file():
        raise ArtifactPathError(f"artifact file not found: {rel.as_posix()}")
    return resolved, rel.as_posix()


def normalize_artifact_refs(refs: Iterable[ArtifactRef]) -> list[dict]:
    """Deduplicate validated refs while preserving order and enforcing the UI limit."""
    out: list[dict] = []
    seen: set[tuple[str, str]] = set()
    for ref in refs:
        if ref.scope != "working_dir" or ref.version not in (
            ARTIFACT_REF_VERSION,
            ARTIFACT_REF_CURRENT_VERSION,
        ):
            continue
        key = (ref.scope, ref.path)
        if key in seen:
            continue
        seen.add(key)
        out.append(ref.as_dict())
        if len(out) >= MAX_ARTIFACTS_PER_MESSAGE:
            break
    return out


def artifact_ref_for_file(
    path: Path,
    *,
    working_dir: Path,
    user_root: Path,
    label: Optional[str] = None,
) -> ArtifactRef:
    _, rel = resolve_artifact_path(
        str(path), working_dir=working_dir, user_root=user_root, require_file=True,
    )
    return ArtifactRef(path=rel, label=(label or "").strip())