nextcraft/apps/ai-service/ai_service/sandbox/manager.py

"""SandboxManager — lifecycle, concurrency guard, reapers (REQ-3-001, REQ-3-002).

Responsibilities (D-032, G-2, a-1):

- create/list/get/snapshot/destroy over a `SandboxBackend` port. create/list/
  get return `SandboxHandleInfo` rows — the handle fields plus the owning
  `learner_id` — so the API layer never re-asks "who owns this id?".
- Capacity guard (D-032): `create` raises `PoolFullError` when the active
  count reaches `settings.sandbox_max_concurrent`. No queue — the API layer
  maps this to 503.
- Wall-clock reaper: `reap_expired()` destroys sandboxes older than
  `settings.sandbox_timeout_s`. Run it on an async timer owned by the caller
  (app lifespan wires the loop; the manager owns only the pass).
- Workdir-size sweep (G-2): the same timer pass also measures each sandbox's
  `workspace/` tree; anything over `settings.sandbox_max_workdir_mb` is
  snapshotted (evidence preserved), destroyed, and recorded as an integrity
  signal. SOFT CAP, best-effort, NOT kernel-enforced — without cgroup
  delegation or sudo there is no hard per-sandbox disk quota on this host.
  RLIMIT_FSIZE bounds a single file; this sweep bounds aggregate growth
  between passes.
- Startup reaper (a-1): `start()` scans `settings.sandbox_dir` for workdirs
  whose recorded pid is dead (marker file `sandbox.json` beside workspace/)
  and reaps them, logging a warning. Handles are IN-MEMORY and process-local
  (D-019 precedent): on process restart every handle is orphaned, so boot
  must recover disk state.

Registry: plain dict guarded by an `asyncio.Lock`, process-local, explicitly
NOT a store. Swapping in persistence (D-027) must not change this interface.

Resource limits: enforced at exec time by the spawner's in-namespace shim
(RLIMIT_AS / RLIMIT_CPU / RLIMIT_FSIZE — see UnshareBackend), never here;
the manager's enforcement surface is lifecycle (capacity, wall-clock, disk
sweep).

BOUNDARY: this module NEVER imports `api/` or `agents/`.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import shutil
import uuid
from collections.abc import Callable
from datetime import UTC, datetime
from pathlib import Path

from pydantic import BaseModel, ConfigDict, Field

from ..config import Settings
from . import workdir as workdir_mod
from .backend import SandboxBackend, SandboxHandle

logger = logging.getLogger(__name__)

#: Marker file beside workspace/ recording the owning process + metadata.
#: It is what the startup reaper uses after this process's in-memory
#: registry is lost (crash/restart → orphan detection, a-1).
PID_MARKER = "sandbox.json"


class PoolFullError(RuntimeError):
    """D-032: active sandbox count reached `settings.sandbox_max_concurrent`.

    The API layer maps this to 503. There is deliberately NO queue.
    """


class SandboxNotFoundError(KeyError):
    """No live sandbox with that id in this process's registry."""


class SandboxIntegrityEvent(BaseModel):
    """One manager-observed integrity signal (G-2).

    Recorded in-process (`SandboxManager.integrity_events`, for the proctor
    pipeline to drain) AND logged at WARNING (durable trail) — the same
    dual-sink pattern a DB-backed store will keep behind D-027.
    """

    model_config = ConfigDict(frozen=True)

    kind: str = Field(description="e.g. 'workdir_size_cap' (G-2), 'orphan_reaped' (a-1)")
    sandbox_id: str
    learner_id: str
    detail: str
    observed_at: datetime = Field(default_factory=lambda: datetime.now(UTC))


class SandboxHandleInfo(BaseModel):
    """A `SandboxHandle` plus its owning `learner_id` (manager return row).

    Handles alone don't carry the learner — the registry side-table does —
    and every API read/list needs it, so the manager joins the two ONCE here
    instead of exposing `_learner_ids` internals.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    id: str
    learner_id: str
    workdir: Path
    created_at: datetime
    pid: int | None = None


class SandboxManager:
    """Lifecycle owner for learner sandboxes.

    Dependencies are injected (D-017 style): the backend port, settings, and
    a wall clock. Single-process only; the registry is in-memory.
    """

    def __init__(
        self,
        *,
        backend: SandboxBackend,
        settings: Settings,
        clock: Callable[[], datetime] | None = None,
    ) -> None:
        self._backend = backend
        self._settings = settings
        self._clock = clock or (lambda: datetime.now(UTC))
        self._handles: dict[str, SandboxHandle] = {}
        self._learner_ids: dict[str, str] = {}  # sandbox_id -> learner_id
        self._lock = asyncio.Lock()
        self._integrity_events: list[SandboxIntegrityEvent] = []
        self._started = False

    # -- introspection ------------------------------------------------------

    @property
    def active_count(self) -> int:
        return len(self._handles)

    @property
    def integrity_events(self) -> list[SandboxIntegrityEvent]:
        """Drainable view of recorded integrity signals (G-2, a-1)."""
        return list(self._integrity_events)

    # -- lifecycle ----------------------------------------------------------

    async def create(
        self, learner_id: str, task_id: str | None = None
    ) -> SandboxHandleInfo:
        """Spawn a sandbox for `learner_id`, or raise `PoolFullError` (D-032).

        `task_id` (REQ-3-003): when set, the sandbox is telemetry-wired — the
        backend copies `scripts/sandbox-agent.py` into the workdir and starts
        the stdlib capture agent inside the sandbox with the `NC_*` env baked
        here (identity + WS ingest URL). The agent's lifecycle is tied to the
        sandbox: `destroy()` reaps it (agent → inner → helper). When `task_id`
        is None the sandbox is a pure shell sandbox (no capture).
        """
        async with self._lock:
            if len(self._handles) >= self._settings.sandbox_max_concurrent:
                raise PoolFullError(
                    f"sandbox pool full "
                    f"({len(self._handles)}/{self._settings.sandbox_max_concurrent}); "
                    "no queue (D-032) — retry later"
                )
            sandbox_id = f"sbx-{uuid.uuid4().hex[:12]}"
            spec = workdir_mod.spec_for(sandbox_id, learner_id, self._settings)
            if task_id is not None:
                spec = spec.model_copy(
                    update={
                        "capture_env": self._capture_env(sandbox_id, learner_id, task_id)
                    }
                )
            handle = await self._backend.spawn(spec)
            self._handles[handle.id] = handle
            self._learner_ids[handle.id] = learner_id
            self._write_pid_marker(handle, learner_id)
            logger.info(
                "sandbox created: id=%s learner=%s task=%s",
                handle.id,
                learner_id,
                task_id or "-",
            )
            return self._info_for(handle)

    def _capture_env(self, sandbox_id: str, learner_id: str, task_id: str) -> dict[str, str]:
        """Env baked for the in-sandbox capture agent (REQ-3-003).

        The agent joins the sandbox mount namespace but NOT its (offline)
        network namespace, so it reaches this service over loopback
        (`telemetry_ingest_host`, A-004 port).
        """
        from urllib.parse import urlencode

        query = urlencode(
            {
                "learner_id": learner_id,
                "task_id": task_id,
                "sandbox_id": sandbox_id,
            }
        )
        ingest_url = (
            f"ws://{self._settings.telemetry_ingest_host}:{self._settings.port}"
            f"/v1/telemetry/ingest?{query}"
        )
        return {
            "NC_LEARNER_ID": learner_id,
            "NC_TASK_ID": task_id,
            "NC_SANDBOX_ID": sandbox_id,
            "NC_INGEST_URL": ingest_url,
        }

    async def list(self) -> list[SandboxHandleInfo]:
        """All live sandboxes (idle + busy; the backend has no busy flag)."""
        async with self._lock:
            return [self._info_for(h) for h in self._handles.values()]

    async def get(self, sandbox_id: str) -> SandboxHandleInfo:
        async with self._lock:
            handle = self._handles.get(sandbox_id)
            learner_id = self._learner_ids.get(sandbox_id, "unknown")
        if handle is None:
            raise SandboxNotFoundError(sandbox_id)
        return SandboxHandleInfo(
            id=handle.id,
            learner_id=learner_id,
            workdir=handle.workdir,
            created_at=handle.created_at,
            pid=handle.pid,
        )

    async def snapshot(self, sandbox_id: str) -> Path:
        """Copy the workspace into `<workdir>/snapshots/<utc-ts>/`; return it."""
        async with self._lock:
            handle = self._handles.get(sandbox_id)
        if handle is None:
            raise SandboxNotFoundError(sandbox_id)
        return await self._backend.snapshot(handle)

    async def destroy(self, sandbox_id: str, *, purge_workdir: bool = False) -> None:
        """Tear down one sandbox (idempotent).

        `purge_workdir=False` keeps the workdir on disk — snapshots must
        survive destroy so a learner's last state can be restored (this is
        also why UnshareBackend.destroy intentionally leaves the tree alone).
        `purge_workdir=True` removes the whole workdir.
        """
        async with self._lock:
            handle = self._handles.pop(sandbox_id, None)
            learner_id = self._learner_ids.pop(sandbox_id, "unknown")
        if handle is not None:
            await self._backend.destroy(handle)
            logger.info(
                "sandbox destroyed: id=%s learner=%s purge=%s",
                sandbox_id,
                learner_id,
                purge_workdir,
            )
            root = handle.workdir
        else:
            # Idempotent destroy of an unknown id: resolve the on-disk root so
            # an explicit purge still works (e.g. cleanup of orphan leftovers).
            root = workdir_mod.resolve_sandbox_dir(self._settings) / sandbox_id
        if purge_workdir:
            shutil.rmtree(root, ignore_errors=True)

    # -- reapers ------------------------------------------------------------

    async def reap_expired(self) -> list[str]:
        """One reaper pass: wall-clock timeout + G-2 workdir-size sweep.

        Destroys sandboxes older than `settings.sandbox_timeout_s`, then sweeps
        every remaining sandbox whose `workspace/` exceeds
        `settings.sandbox_max_workdir_mb` (snapshot → destroy → integrity
        signal). Returns the ids destroyed this pass. Invoke on an async timer
        (the app lifespan owns the loop interval); both checks deliberately
        share one pass so the periodic work is O(live sandboxes) once.
        """
        now = self._clock()
        destroyed: list[str] = []
        timeout_s = float(self._settings.sandbox_timeout_s)
        cap_bytes = int(self._settings.sandbox_max_workdir_mb) * 1024 * 1024

        async with self._lock:
            rows = [
                (handle, self._learner_ids.get(handle.id, "unknown"), handle.created_at)
                for handle in self._handles.values()
            ]

        for handle, learner_id, created_at in rows:
            age_s = (now - created_at).total_seconds()
            if age_s > timeout_s:
                await self.destroy(handle.id)
                destroyed.append(handle.id)
                logger.warning(
                    "sandbox reaped (timeout): id=%s age=%.0fs > %.0fs",
                    handle.id,
                    age_s,
                    timeout_s,
                )
                continue  # already gone; no size sweep needed on a dead handle
            size = _tree_size_bytes(workdir_mod.workspace_path_from_workdir(handle.workdir))
            if size > cap_bytes:
                await self._reap_oversized(handle, learner_id, size, cap_bytes)
                destroyed.append(handle.id)
        return destroyed

    async def start(self) -> None:
        """Boot hook (a-1): reap on-disk orphans left by a previous process.

        The handle registry is in-memory and process-local (D-019 precedent):
        after a restart nothing here remembers old sandboxes, so we scan
        `settings.sandbox_dir` for workdirs whose pid marker names a dead
        process and purge them, logging a warning. Idempotent; safe to call
        once per process lifetime.
        """
        if self._started:
            return
        self._started = True
        root = workdir_mod.resolve_sandbox_dir(self._settings)
        if not root.is_dir():
            return
        for entry in sorted(root.iterdir()):
            if not entry.is_dir():
                continue
            marker = entry / PID_MARKER
            pid = _read_marker_pid(marker)
            if pid is not None and _pid_alive(pid):
                continue  # live sandbox owned by another live process — leave it
            logger.warning(
                "startup reaper (a-1): reaping orphaned workdir %s "
                "(recorded pid %s is dead or marker missing)",
                entry,
                pid,
            )
            shutil.rmtree(entry, ignore_errors=True)
            self._record_integrity(
                SandboxIntegrityEvent(
                    kind="orphan_reaped",
                    sandbox_id=entry.name,
                    learner_id="unknown",
                    detail=f"workdir {entry} reaped at boot; recorded pid={pid} dead",
                )
            )

    async def destroy_all(self) -> None:
        """Shutdown hook: destroy every live sandbox (no orphans on exit).

        Workdirs (and their snapshots) are kept on disk — destroy semantics
        here match `destroy(purge_workdir=False)`; the next boot's startup
        reaper (a-1) decides what to clean based on pid markers.
        """
        async with self._lock:
            handles = list(self._handles.values())
        for handle in handles:
            await self.destroy(handle.id)

    # -- internals ------------------------------------------------------------

    def _info_for(self, handle: SandboxHandle) -> SandboxHandleInfo:
        # Caller holds the lock (create/list) — the side-table read is atomic.
        return SandboxHandleInfo(
            id=handle.id,
            learner_id=self._learner_ids.get(handle.id, "unknown"),
            workdir=handle.workdir,
            created_at=handle.created_at,
            pid=handle.pid,
        )

    def _write_pid_marker(self, handle: SandboxHandle, learner_id: str) -> None:
        marker = handle.workdir / PID_MARKER
        try:
            marker.write_text(
                json.dumps(
                    {
                        "sandbox_id": handle.id,
                        "learner_id": learner_id,
                        "pid": os.getpid(),
                        "created_at": handle.created_at.isoformat(),
                    }
                )
            )
        except OSError:  # marker is advisory; spawning must not fail on it
            logger.warning("could not write pid marker %s", marker)

    async def _reap_oversized(
        self,
        handle: SandboxHandle,
        learner_id: str,
        size_bytes: int,
        cap_bytes: int,
    ) -> None:
        """G-2 sweep step: snapshot evidence → destroy → record the signal."""
        snapshot_path: Path | None = None
        try:
            snapshot_path = await self._backend.snapshot(handle)
        except (OSError, RuntimeError):
            logger.exception(
                "G-2 sweep: snapshot failed for over-cap sandbox %s; destroying anyway",
                handle.id,
            )
        await self.destroy(handle.id)
        event = SandboxIntegrityEvent(
            kind="workdir_size_cap",
            sandbox_id=handle.id,
            learner_id=learner_id,
            detail=(
                f"workspace {size_bytes}B exceeded soft cap {cap_bytes}B; "
                f"snapshot={snapshot_path} then destroyed (G-2, best-effort, "
                "NOT kernel-enforced)"
            ),
        )
        self._record_integrity(event)
        logger.warning(
            "G-2 workdir sweep: sandbox %s (learner=%s) destroyed over soft disk cap",
            handle.id,
            learner_id,
        )

    def _record_integrity(self, event: SandboxIntegrityEvent) -> None:
        self._integrity_events.append(event)


# -- module helpers ---------------------------------------------------------


def _tree_size_bytes(root: Path) -> int:
    """Total bytes under `root` (best-effort; unreadable entries count 0)."""
    if not root.is_dir():
        return 0
    total = 0
    for dirpath, _dirnames, filenames in os.walk(root):
        for name in filenames:
            try:
                total += (Path(dirpath) / name).lstat().st_size
            except OSError:
                continue
    return total


def _read_marker_pid(marker: Path) -> int | None:
    try:
        data = json.loads(marker.read_text())
    except (OSError, json.JSONDecodeError):
        return None
    pid = data.get("pid")
    return pid if isinstance(pid, int) else None


def _pid_alive(pid: int) -> bool:
    """True if `pid` exists on this host (signal 0 probe; no signal sent)."""
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        return True  # exists, owned by another user
    return True