tools_api/storage.py

"""
Per-run temp storage for tools_api.

Each tool request creates a fresh dir under ARTIFACTS_ROOT/tools/<run_id>/.
Files are reaped after TTL by _reap_old_runs(). Kept independent of the main
job-tracker so a tool failure can't corrupt or block pipeline state.
"""
from __future__ import annotations

import shutil
import time
import uuid
from pathlib import Path
from typing import Optional

# Pull ARTIFACTS_ROOT from server.py without importing the heavy modules
# (server.py imports torch/whisper/etc. at top level — we already loaded it
# at app startup, so this is just a name lookup).
from server import ARTIFACTS_ROOT

TOOLS_ROOT = ARTIFACTS_ROOT / "tools"
TOOLS_ROOT.mkdir(parents=True, exist_ok=True)

# Tool runs are reaped 1h after creation (shorter than pipeline jobs since
# users typically download immediately).
RUN_TTL_SECONDS = 60 * 60


def new_run_dir() -> tuple[str, Path]:
    """Allocate a fresh per-request directory. Returns (run_id, path)."""
    run_id = uuid.uuid4().hex[:16]
    path = TOOLS_ROOT / run_id
    path.mkdir(parents=True, exist_ok=True)
    return run_id, path


def run_dir(run_id: str) -> Optional[Path]:
    """Resolve a run_id to its directory, or None if missing/invalid."""
    if not run_id or "/" in run_id or ".." in run_id:
        return None
    candidate = TOOLS_ROOT / run_id
    if not candidate.exists() or not candidate.is_dir():
        return None
    return candidate


def file_url(run_id: str, filename: str) -> str:
    """Construct the public download URL for an artifact."""
    return f"/api/tools/file/{run_id}/{filename}"


def safe_filename(name: str, fallback: str = "file") -> str:
    """Strip path separators and dangerous chars from a user-supplied name."""
    if not name:
        return fallback
    base = Path(name).name
    return base or fallback


def reap_old_runs() -> int:
    """Delete tool run dirs older than RUN_TTL_SECONDS. Returns count removed."""
    if not TOOLS_ROOT.exists():
        return 0
    cutoff = time.time() - RUN_TTL_SECONDS
    removed = 0
    for child in TOOLS_ROOT.iterdir():
        try:
            if child.is_dir() and child.stat().st_mtime < cutoff:
                shutil.rmtree(child, ignore_errors=True)
                removed += 1
        except OSError:
            continue
    return removed