core/hosted_runtime.py

"""Runtime isolation and quota protection for the public Hugging Face Space.

Each Streamlit browser session receives its own SQLite database and trace tree.
An atomic lock limits expensive negotiation pipelines to one at a time while
allowing completed pipelines to wait independently at their human gates.
"""
from __future__ import annotations

import json
import os
import shutil
import time
import uuid
from dataclasses import dataclass
from pathlib import Path

from core.config import PROJECT_ROOT, settings


def is_hosted() -> bool:
    return os.environ.get("AUTOSOURCE_HOSTED", "").lower() in {"1", "true", "yes"}


@dataclass(frozen=True)
class SessionWorkspace:
    session_id: str
    root: Path
    db: Path
    traces: Path

    def child_env(self) -> dict[str, str]:
        env = os.environ.copy()
        env.update({
            "AUTOSOURCE_DB_PATH": str(self.db),
            "AUTOSOURCE_TRACES_DIR": str(self.traces),
            "PYTHONUNBUFFERED": "1",
            "PYTHONIOENCODING": "utf-8",
        })
        return env


def session_base() -> Path:
    value = os.environ.get("AUTOSOURCE_SESSION_ROOT")
    root = Path(value) if value else PROJECT_ROOT / ".runtime" / "sessions"
    root.mkdir(parents=True, exist_ok=True)
    return root.resolve()


def get_workspace(session_id: str) -> SessionWorkspace:
    safe_id = "".join(c for c in session_id if c.isalnum() or c in "-_")[:64]
    if not safe_id:
        raise ValueError("invalid session id")
    root = session_base() / safe_id
    traces = root / "traces"
    traces.mkdir(parents=True, exist_ok=True)
    os.utime(root, None)
    return SessionWorkspace(safe_id, root, root / "warehouse.db", traces)


def new_session_id() -> str:
    return uuid.uuid4().hex


def cleanup_expired_sessions(max_age_s: int | None = None) -> int:
    """Remove abandoned ephemeral workspaces; never touches the bundled replay."""
    max_age_s = max_age_s or int(settings().get("hosted", {}).get(
        "session_ttl_s", 21600))
    cutoff = time.time() - max_age_s
    removed = 0
    for path in session_base().iterdir():
        try:
            if path.is_dir() and path.stat().st_mtime < cutoff:
                shutil.rmtree(path)
                removed += 1
        except OSError:
            continue
    return removed


def lock_path() -> Path:
    value = os.environ.get("AUTOSOURCE_LIVE_LOCK")
    return Path(value) if value else session_base().parent / "live-pipeline.lock"


def _read_lock() -> dict | None:
    path = lock_path()
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (FileNotFoundError, json.JSONDecodeError, OSError):
        return None


def live_slot_status() -> dict | None:
    record = _read_lock()
    if not record:
        return None
    max_age = int(settings().get("hosted", {}).get("pipeline_lock_ttl_s", 900))
    if time.time() - float(record.get("created_at", 0)) > max_age:
        try:
            lock_path().unlink()
        except OSError:
            pass
        return None
    return record


def acquire_live_slot(session_id: str, run_id: str) -> str | None:
    """Atomically reserve the shared provider slot. Returns an ownership token."""
    live_slot_status()  # clears a stale lock before the atomic create
    token = uuid.uuid4().hex
    record = {
        "token": token,
        "session_id": session_id,
        "run_id": run_id,
        "created_at": time.time(),
    }
    path = lock_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    try:
        fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    except FileExistsError:
        return None
    with os.fdopen(fd, "w", encoding="utf-8") as stream:
        json.dump(record, stream)
    return token


def release_live_slot(token: str) -> bool:
    """Release only when the caller owns the lock; prevents cross-session unlocks."""
    record = _read_lock()
    if not record or record.get("token") != token:
        return False
    try:
        lock_path().unlink()
        return True
    except FileNotFoundError:
        return True
    except OSError:
        return False