ai-agent/backend/files/assets.py

"""Serving media: project assets, artefacts, screenshots, data/ and
scratchpad files, and the composer's uploads."""

import pathlib
import uuid as uuidlib

from fastapi import APIRouter, File, HTTPException, UploadFile
from fastapi.responses import FileResponse

import projects as projects_mod
import schemas
from core.config import (
    _UNSAFE_NAME_RE,
    DATA_DIR,
    IMAGE_EXTS,
    MAX_UPLOAD_BYTES,
    SCRATCHPAD_DIR,
    SCREENSHOT_EXTS,
    SCREENSHOTS_DIR,
    UPLOADS_DIR,
    UPLOADS_REPO_PREFIX,
)
from core.http import _r
from projects import artefacts as artefacts_mod
from users.caller import CallerDep

router = APIRouter()


# Raster formats that are safe to render inline on the app origin. Anything
# else served from user/repo-supplied bytes (SVG can carry scripts, HTML is
# HTML) goes out as a download so it can never script against the API's cookies.
_INLINE_IMAGE_EXTS = {".jpeg", ".jpg", ".png", ".webp", ".gif", ".bmp", ".avif"}

# Video containers safe to play inline on the app origin: a media file can't
# script against the API's cookies the way an SVG/HTML can, so it need not go
# out as an attachment. FileResponse honours Range requests (Starlette), so the
# <video> element can seek.
_INLINE_VIDEO_EXTS = {".mp4", ".webm", ".mov", ".m4v", ".ogv"}


def _safe_file_response(p: pathlib.Path) -> FileResponse:
    headers = {"X-Content-Type-Options": "nosniff"}
    ext = p.suffix.lower()
    if ext not in _INLINE_IMAGE_EXTS and ext not in _INLINE_VIDEO_EXTS:
        headers["Content-Disposition"] = f'attachment; filename="{p.name}"'
    return FileResponse(p, headers=headers)


@router.get("/api/project-asset")
def project_asset(caller: CallerDep, slug: str, path: str):
    """Serve a whitelisted image asset (e.g. og-image) from a project dir."""
    if not caller.can(slug, "view"):
        raise HTTPException(404, "asset not found")
    p = projects_mod.project_asset_path(slug, path)
    if p is None:
        raise HTTPException(404, "asset not found")
    return _safe_file_response(p)


@router.get("/api/artefact")
def artefact_asset(caller: CallerDep, kind: str, slug: str, path: str):
    """Serve one generated artefact from a project/service's `.ai/artefacts/`
    tree (``path`` is ``<date>/<sessionId>/<file>``, relative to that tree)."""
    if not caller.is_owner and (kind != "project" or not caller.can(slug, "view")):
        raise HTTPException(404, "artefact not found")
    p = artefacts_mod.artefact_path(kind, slug, path)
    if p is None:
        raise HTTPException(404, "artefact not found")
    return _safe_file_response(p)


@router.get("/api/screenshot")
def screenshot_asset(path: str):
    """Serve a captured screenshot image from docs/screenshots/ (read-only).

    ``path`` is the image's location relative to the screenshots root
    (``<project>/<name>_<ts>.jpeg``). A leading ``docs/screenshots/`` prefix,
    if present, is stripped so the raw path printed by the skill also works.
    """
    rel = path.replace("\\", "/").lstrip("/")
    for pref in ("docs/screenshots/", "screenshots/"):
        if rel.startswith(pref):
            rel = rel[len(pref):]
    rel_p = pathlib.Path(rel)
    if not rel or rel_p.is_absolute() or ".." in rel_p.parts:
        raise HTTPException(400, "bad path")
    p = (SCREENSHOTS_DIR / rel_p).resolve()
    try:
        p.relative_to(SCREENSHOTS_DIR)
    except ValueError:
        raise HTTPException(400, "bad path")
    if not (p.is_file() and p.suffix.lower() in SCREENSHOT_EXTS):
        raise HTTPException(404, "screenshot not found")
    return _safe_file_response(p)


@router.get("/api/data-asset")
def data_asset(path: str):
    """Serve an image or video artifact from the repo's root `data/` tree (read-only).

    ``path`` is the file's location relative to ``data/`` (a leading ``data/``
    prefix, if present, is stripped so the raw repo-relative path also works).
    Used to preview media the agent ``Read`` from under ``data/`` inline in the
    conversation viewer — see the CLAUDE.md convention to write artifacts there.
    """
    rel = path.replace("\\", "/").lstrip("/")
    if rel.startswith("data/"):
        rel = rel[len("data/"):]
    rel_p = pathlib.Path(rel)
    if not rel or rel_p.is_absolute() or ".." in rel_p.parts:
        raise HTTPException(400, "bad path")
    p = (DATA_DIR / rel_p).resolve()
    try:
        p.relative_to(DATA_DIR)
    except ValueError:
        raise HTTPException(400, "bad path")
    if not (p.is_file() and p.suffix.lower() in (IMAGE_EXTS | _INLINE_VIDEO_EXTS)):
        raise HTTPException(404, "data asset not found")
    return _safe_file_response(p)


@router.get("/api/scratchpad-asset")
def scratchpad_asset(path: str):
    """Serve an image or video from Claude Code's session scratchpad tree.

    ``path`` is relative to the scratchpad root — ``<encoded-cwd>/<sessionId>/
    scratchpad/<file>`` — so a run's throwaway media previews inline in the
    thread without having to be copied into a project's ``.ai/artefacts/``.
    The mount is read-only and only media extensions are served.
    """
    rel = path.replace("\\", "/").lstrip("/")
    rel_p = pathlib.Path(rel)
    if not rel or rel_p.is_absolute() or ".." in rel_p.parts:
        raise HTTPException(400, "bad path")
    p = (SCRATCHPAD_DIR / rel_p).resolve()
    try:
        p.relative_to(SCRATCHPAD_DIR)
    except ValueError:
        raise HTTPException(400, "bad path")
    if not (p.is_file() and p.suffix.lower() in (IMAGE_EXTS | _INLINE_VIDEO_EXTS)):
        raise HTTPException(404, "scratchpad asset not found")
    return _safe_file_response(p)


# ── composer file uploads ────────────────────────────────────────────────────
# The composer (SpawnBox / ResumeBox) uploads files here; each batch lands in a
# random subdir under UPLOADS_DIR. The endpoint returns, per file, the
# repo-relative path to inject into the prompt (so the host-side `claude -p`
# session can Read it) and a serve URL the frontend renders inline in the thread.
def _safe_upload_name(name: str) -> str:
    base = pathlib.PurePosixPath((name or "").replace("\\", "/")).name
    base = _UNSAFE_NAME_RE.sub("_", base).strip("._")
    return (base or "file")[:120]


def _upload_rel(path: str) -> pathlib.Path:
    """Validate a client-supplied uploads path and return it relative to
    UPLOADS_DIR, tolerating a leading repo prefix (the injected prompt form)."""
    rel = path.replace("\\", "/").lstrip("/")
    prefix = UPLOADS_REPO_PREFIX + "/"
    if rel.startswith(prefix):
        rel = rel[len(prefix):]
    rel_p = pathlib.Path(rel)
    if not rel or rel_p.is_absolute() or ".." in rel_p.parts:
        raise HTTPException(400, "bad path")
    p = (UPLOADS_DIR / rel_p).resolve()
    try:
        p.relative_to(UPLOADS_DIR)
    except ValueError:
        raise HTTPException(400, "bad path")
    return p


@router.post("/api/upload", responses=_r(schemas.UploadResult))
async def upload_files(files: list[UploadFile] = File(...)):
    """Store one or more composer attachments and describe where they landed.

    Each upload is written to ``UPLOADS_DIR/<batch>/<safe-name>``. The returned
    ``repoPath`` is what the composer injects into the prompt (the session reads
    it from the repo working dir); ``url`` is the backend serve route the thread
    viewer renders the file from."""
    if not files:
        raise HTTPException(400, "no files")
    batch = uuidlib.uuid4().hex[:12]
    dest = UPLOADS_DIR / batch
    dest.mkdir(parents=True, exist_ok=True)
    out = []
    for f in files:
        name = _safe_upload_name(f.filename or "file")
        target = dest / name
        i = 1
        while target.exists():
            target = dest / f"{target.stem}-{i}{pathlib.Path(name).suffix}"
            i += 1
        size = 0
        try:
            with open(target, "wb") as w:
                while chunk := await f.read(1 << 20):
                    size += len(chunk)
                    if size > MAX_UPLOAD_BYTES:
                        raise HTTPException(
                            413, f"{name} exceeds {MAX_UPLOAD_BYTES // (1<<20)} MB")
                    w.write(chunk)
        except HTTPException:
            target.unlink(missing_ok=True)
            raise
        rel = f"{batch}/{target.name}"
        out.append({
            "name": target.name,
            "size": size,
            "contentType": f.content_type or "application/octet-stream",
            "repoPath": f"{UPLOADS_REPO_PREFIX}/{rel}",
            "url": f"/api/upload-file?path={rel}",
        })
    return {"files": out}


@router.get("/api/upload-file")
def upload_file_asset(path: str):
    """Serve a previously uploaded composer attachment (raw bytes, inline)."""
    p = _upload_rel(path)
    if not p.is_file():
        raise HTTPException(404, "upload not found")
    return _safe_file_response(p)