nextcraft/apps/ai-service/tests/sandbox/test_manager.py

"""SandboxManager tests (REQ-3-001, REQ-3-002).

The manager's create/list/get return `SandboxHandleInfo` rows (handle fields
+ learner_id); `backend.exec` and workdir helpers still take the raw handle,
so handle-mutating tests reach the registry via `manager._handles`.
"""

from __future__ import annotations

import json
import os
from datetime import UTC, datetime, timedelta
from pathlib import Path

import pytest

from ai_service.config import Settings
from ai_service.sandbox import (
    PoolFullError,
    SandboxHandle,
    SandboxManager,
    SandboxNotFoundError,
    UnshareBackend,
)
from ai_service.sandbox.backend import ExecResult, SandboxSpec
from ai_service.sandbox.manager import PID_MARKER
from ai_service.sandbox.workdir import create_layout, spec_for
from tests.sandbox.test_isolation import requires_userns


class FakeBackend:
    """Structural SandboxBackend: lays out the workdir, spawns nothing."""

    def __init__(self) -> None:
        self.destroyed_ids: list[str] = []

    async def spawn(self, spec: SandboxSpec) -> SandboxHandle:
        create_layout(spec)
        return SandboxHandle(
            id=spec.sandbox_id, pid=None, workdir=spec.workdir,
            created_at=datetime.now(UTC),
        )

    async def exec(self, handle: SandboxHandle, cmd: list[str]) -> ExecResult:
        raise NotImplementedError("manager tests never exec")

    async def snapshot(self, handle: SandboxHandle) -> Path:
        from ai_service.sandbox.workdir import snapshot as workdir_snapshot

        return workdir_snapshot(handle.workdir)

    async def destroy(self, handle: SandboxHandle) -> None:
        self.destroyed_ids.append(handle.id)
        handle.pid = None


@pytest.fixture()
def settings(tmp_path: Path) -> Settings:
    return Settings(
        provider="mock",
        sandbox_dir=tmp_path / "sandboxes",
        sandbox_max_concurrent=5,
        sandbox_timeout_s=900.0,
        sandbox_max_workdir_mb=512,
    )


@pytest.fixture()
def backend() -> FakeBackend:
    return FakeBackend()


@pytest.fixture()
def manager(settings: Settings, backend: FakeBackend) -> SandboxManager:
    return SandboxManager(backend=backend, settings=settings)


# -- create / guard -----------------------------------------------------------


async def test_create_registers_handle(manager: SandboxManager) -> None:
    info = await manager.create("learner-1")
    assert info.id.startswith("sbx-")
    assert info.learner_id == "learner-1"
    assert info.pid is None
    assert (info.workdir / "workspace").is_dir()
    assert manager.active_count == 1
    marker = json.loads((info.workdir / PID_MARKER).read_text())
    assert marker["pid"] == os.getpid()
    assert marker["learner_id"] == "learner-1"


async def test_pool_full_guard_raises_at_max_concurrent(
    manager: SandboxManager, settings: Settings
) -> None:
    for _ in range(settings.sandbox_max_concurrent):
        await manager.create("learner-1")
    assert manager.active_count == settings.sandbox_max_concurrent
    with pytest.raises(PoolFullError):  # D-032: no queue; API maps to 503
        await manager.create("learner-1")
    assert manager.active_count == settings.sandbox_max_concurrent


async def test_capacity_frees_after_destroy(manager: SandboxManager) -> None:
    handles = [await manager.create("learner-1") for _ in range(5)]
    await manager.destroy(handles[0].id)
    refill = await manager.create("learner-2")
    assert refill.id != handles[0].id
    assert manager.active_count == 5


# -- list / get / snapshot / destroy -----------------------------------------


async def test_list_and_get_roundtrip(manager: SandboxManager) -> None:
    a = await manager.create("learner-1")
    b = await manager.create("learner-2")
    listed = await manager.list()
    assert {h.id for h in listed} == {a.id, b.id}
    assert {h.learner_id for h in listed} == {"learner-1", "learner-2"}
    got = await manager.get(a.id)
    assert got.id == a.id
    assert got.learner_id == "learner-1"
    with pytest.raises(SandboxNotFoundError):
        await manager.get("sbx-nope")


async def test_snapshot_copies_workspace(manager: SandboxManager) -> None:
    info = await manager.create("learner-1")
    (info.workdir / "workspace" / "solution.py").write_text("print(42)\n")
    snap = await manager.snapshot(info.id)
    assert snap.parent == info.workdir / "snapshots"
    assert (snap / "solution.py").read_text() == "print(42)\n"


async def test_destroy_keeps_workdir_with_snapshots(
    manager: SandboxManager, backend: FakeBackend
) -> None:
    info = await manager.create("learner-1")
    (info.workdir / "workspace" / "keep.txt").write_text("state")
    await manager.snapshot(info.id)
    await manager.destroy(info.id)
    assert info.id in backend.destroyed_ids
    assert manager.active_count == 0
    assert info.workdir.is_dir()  # snapshots survive destroy (restore path)
    assert list((info.workdir / "snapshots").iterdir())
    with pytest.raises(SandboxNotFoundError):
        await manager.get(info.id)


async def test_destroy_purge_removes_workdir(
    manager: SandboxManager, backend: FakeBackend
) -> None:
    info = await manager.create("learner-1")
    await manager.destroy(info.id, purge_workdir=True)
    assert info.id in backend.destroyed_ids
    assert not info.workdir.exists()


async def test_destroy_unknown_id_is_idempotent(manager: SandboxManager) -> None:
    await manager.destroy("sbx-ghost")  # must not raise


async def test_destroy_all_empties_registry_keeps_workdirs(
    manager: SandboxManager, backend: FakeBackend
) -> None:
    a = await manager.create("learner-1")
    b = await manager.create("learner-2")
    await manager.destroy_all()
    assert manager.active_count == 0
    assert backend.destroyed_ids == [a.id, b.id]
    assert a.workdir.is_dir() and b.workdir.is_dir()  # snapshots survive
    await manager.destroy_all()  # idempotent no-op on an empty registry
    assert backend.destroyed_ids == [a.id, b.id]


# -- reap_expired: wall-clock + G-2 sweep -------------------------------------


async def test_reap_expired_destroys_timed_out_sandbox(
    manager: SandboxManager, backend: FakeBackend
) -> None:
    info = await manager.create("learner-1")
    # created_at on the returned info is a snapshot; mutate the registry row.
    manager._handles[info.id].created_at = datetime.now(UTC) - timedelta(seconds=901)
    fresh = await manager.create("learner-2")
    destroyed = await manager.reap_expired()
    assert destroyed == [info.id]
    assert info.id in backend.destroyed_ids
    assert (await manager.list())[0].id == fresh.id


async def test_reap_expired_within_timeout_keeps_sandbox(
    manager: SandboxManager, backend: FakeBackend
) -> None:
    await manager.create("learner-1")
    assert await manager.reap_expired() == []
    assert backend.destroyed_ids == []
    assert manager.active_count == 1


async def test_workdir_size_sweep_destroys_over_cap_and_records_signal(
    settings: Settings, backend: FakeBackend, caplog: pytest.LogCaptureFixture
) -> None:
    settings.sandbox_max_workdir_mb = 1  # 1 MiB cap so the test stays tiny
    manager = SandboxManager(backend=backend, settings=settings)
    info = await manager.create("learner-greedy")
    # Over the cap, spread across many files (G-2 is the aggregate guard;
    # RLIMIT_FSIZE alone does not catch this).
    for i in range(9):
        (info.workdir / "workspace" / f"chunk-{i}.bin").write_bytes(b"x" * 256 * 1024)
    with caplog.at_level("WARNING"):
        destroyed = await manager.reap_expired()
    assert destroyed == [info.id]
    assert info.id in backend.destroyed_ids
    events = manager.integrity_events
    assert len(events) == 1
    assert events[0].kind == "workdir_size_cap"
    assert events[0].sandbox_id == info.id
    assert events[0].learner_id == "learner-greedy"
    assert any("G-2" in rec.message for rec in caplog.records)
    # Snapshot-then-destroy: evidence preserved on disk after the reap.
    snapshots = list((info.workdir / "snapshots").iterdir())
    assert len(snapshots) == 1
    assert (snapshots[0] / "chunk-0.bin").is_file()


async def test_sweep_skips_under_cap_sandbox(manager: SandboxManager) -> None:
    info = await manager.create("learner-1")
    (info.workdir / "workspace" / "small.txt").write_text("ok")
    assert await manager.reap_expired() == []
    assert manager.active_count == 1


# -- startup reaper (a-1) ------------------------------------------------------


async def _orphan_workdir(settings: Settings, pid: int, sandbox_id: str) -> Path:
    spec = spec_for(sandbox_id, "learner-orphan", settings)
    create_layout(spec)
    (spec.workdir / PID_MARKER).write_text(
        json.dumps({"sandbox_id": sandbox_id, "learner_id": "learner-orphan", "pid": pid})
    )
    return spec.workdir


async def test_start_reaps_dead_pid_workdirs(
    settings: Settings, backend: FakeBackend, caplog: pytest.LogCaptureFixture
) -> None:
    orphan = await _orphan_workdir(settings, pid=2**22 + 12345, sandbox_id="sbx-dead")
    assert orphan.is_dir()

    manager = SandboxManager(backend=backend, settings=settings)
    with caplog.at_level("WARNING"):
        await manager.start()
    assert not orphan.exists()
    assert any("startup reaper" in rec.message for rec in caplog.records)
    events = manager.integrity_events
    assert [e.kind for e in events] == ["orphan_reaped"]
    assert events[0].sandbox_id == "sbx-dead"


async def test_start_keeps_live_pid_workdirs(
    settings: Settings, backend: FakeBackend, caplog: pytest.LogCaptureFixture
) -> None:
    live = await _orphan_workdir(settings, pid=os.getpid(), sandbox_id="sbx-alive")
    manager = SandboxManager(backend=backend, settings=settings)
    with caplog.at_level("WARNING"):
        await manager.start()
    assert live.is_dir()
    assert manager.integrity_events == []
    assert not any("startup reaper" in rec.message for rec in caplog.records)
    await manager.start()  # idempotent: a second boot hook changes nothing
    assert manager.integrity_events == []


async def test_start_reaps_markerless_and_malformed_workdirs(
    settings: Settings, backend: FakeBackend
) -> None:
    markerless = settings.sandbox_dir / "sbx-markerless"
    markerless.mkdir(parents=True)
    malformed = settings.sandbox_dir / "sbx-badjson"
    malformed.mkdir()
    (malformed / PID_MARKER).write_text("{not json")
    manager = SandboxManager(backend=backend, settings=settings)
    await manager.start()
    assert not markerless.exists()
    assert not malformed.exists()


# -- real-backend smoke (probe-guarded) ----------------------------------------


@requires_userns
async def test_real_unshare_backend_create_exec_destroy(tmp_path: Path) -> None:
    """End-to-end on a capable host: create → exec echo → destroy."""
    # NOT under /tmp: the in-namespace tmpfs shadows host /tmp, hiding a
    # workspace rooted there — anchor in the repo like test_isolation.py.
    sandbox_root = Path(__file__).resolve().parents[1] / "sandboxes" / tmp_path.name
    settings = Settings(provider="mock", sandbox_dir=sandbox_root)
    backend = UnshareBackend()
    manager = SandboxManager(backend=backend, settings=settings)
    info = await manager.create("learner-smoke")
    # exec takes the raw backend handle; the registry holds it for info.id.
    handle = manager._handles[info.id]
    result = await backend.exec(handle, ["echo", "hello-from-manager"])
    assert result.returncode == 0 and result.stdout.strip() == "hello-from-manager"
    await manager.destroy(info.id)
    assert manager.active_count == 0
    assert info.workdir.is_dir()  # workdir kept for snapshot restore