src/managers/execution/sandbox/executor.py

"""
SandboxedExecutor — the client half of the ADR-0007 Phase 1 sandbox executor.

Implements the `Executor` protocol (`base.Executor`) so it drops in behind
`get_executor()` with NO agent-loop change: the agent already depends only on
`reset_environment` / `send_functions` / `send_variables` / `execute` / `__call__`.

Lifecycle (session-granular, per ADR-0007 — state persists across execute calls):
  * lazy start: the FIRST call that needs the kernel starts ONE kernel for the
    session via the pluggable launcher (SubprocessLauncher for dev/test,
    ContainerLauncher for prod), then waits for `/health`.
  * proxy: every call is a small HTTP round-trip; only captured stdout (a string)
    and JSON control messages cross the boundary.
  * teardown: `close()` / context-exit / session end tears the kernel down and
    reaps the process/container.

`send_functions` receives LIVE tool-wrapper objects in-process; those can't cross
the boundary, so the executor extracts their NAMES (+ docstrings as schema) and
sends only that — inside the kernel the names resolve to MCP-dispatching stubs
(see `mcp_bridge`). Non-tool callables that AREN'T MCP tools are simply named to
the kernel too; wiring arbitrary local callables into the sandbox is out of scope
for Phase 1 (the agent's injected functions are the MCP tool wrappers).

`requests` is imported lazily inside the HTTP helpers (boto3-style discipline) so
importing this module never fails when it's absent.
"""

from __future__ import annotations

import json
import os
import time
import urllib.error
import urllib.request
from typing import Any

from .launchers import ContainerLauncher, Launcher, SubprocessLauncher

__all__ = ["SandboxedExecutor", "get_sandboxed_executor"]

# How long to wait for the kernel's /health after launch (seconds), and poll gap.
_HEALTH_TIMEOUT = float(os.environ.get("SANDBOX_HEALTH_TIMEOUT", "30"))
_HEALTH_POLL = 0.2


def _http_post(url: str, payload: dict, timeout: float = 300.0) -> dict:
    """POST JSON to `url`, return parsed JSON. stdlib-only so no requests dep.

    Kept on urllib (not requests) deliberately: the CLIENT must import cleanly in
    the dep-light agent process; requests is only needed INSIDE the kernel for MCP
    dispatch (and imported lazily there).
    """
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url, data=data, headers={"Content-Type": "application/json"}, method="POST"
    )
    # localhost-only sandbox kernel URL; scheme is fixed http, not user-controlled.
    with urllib.request.urlopen(req, timeout=timeout) as resp:  # nosec B310
        return json.loads(resp.read().decode("utf-8"))


def _http_get(url: str, timeout: float = 5.0) -> dict:
    req = urllib.request.Request(url, method="GET")
    # localhost-only sandbox kernel URL; scheme is fixed http, not user-controlled.
    with urllib.request.urlopen(req, timeout=timeout) as resp:  # nosec B310
        return json.loads(resp.read().decode("utf-8"))


def _tool_entries_from_functions(functions: dict[str, Any]) -> list[dict]:
    """Reduce live tool-wrapper objects to boundary-crossing {name, description}.

    Only the NAME (and docstring, as a description hint) crosses; the actual
    callable stays in the agent process's tool namespace and its computation runs
    in the MCP server.
    """
    entries: list[dict] = []
    for name, fn in (functions or {}).items():
        doc = (getattr(fn, "__doc__", None) or "").strip()
        entries.append({"name": name, "description": doc})
    return entries


class SandboxedExecutor:
    """Executor that runs generated code inside a per-session sandbox kernel.

    Conforms to `base.Executor`. See module docstring for lifecycle + boundary
    semantics.
    """

    def __init__(
        self,
        launcher: Launcher | None = None,
        session_id: str | None = None,
        mcp_url: str | None = None,
    ) -> None:
        self.session_id = (
            session_id or os.environ.get("SANDBOX_SESSION_ID") or f"sess-{os.getpid()}"
        )
        self.mcp_url = mcp_url if mcp_url is not None else os.environ.get("SANDBOX_MCP_URL")
        self._launcher = launcher or self._default_launcher()
        self._base_url: str | None = None
        self._pending_tools: list[dict] = []  # tools sent before the kernel started

    # -- launcher selection ------------------------------------------------- #
    def _default_launcher(self) -> Launcher:
        """Pick a launcher from env (SANDBOX_LAUNCHER=subprocess|container).

        Default = container (the prod target). Dev/test override to `subprocess`
        so the whole executor is validatable without Docker.
        """
        kind = os.environ.get("SANDBOX_LAUNCHER", "container").strip().lower()
        if kind == "subprocess":
            return SubprocessLauncher(self.session_id, mcp_url=self.mcp_url)
        if kind == "container":
            return ContainerLauncher(self.session_id, mcp_url=self.mcp_url)
        raise ValueError(f"Unknown SANDBOX_LAUNCHER={kind!r}. Valid: ['subprocess', 'container'].")

    # -- lifecycle ---------------------------------------------------------- #
    def _ensure_started(self) -> str:
        if self._base_url is not None:
            return self._base_url
        base_url = self._launcher.start()
        self._wait_for_health(base_url)
        self._base_url = base_url
        # Flush any tool registrations queued before the kernel existed.
        if self._pending_tools:
            _http_post(f"{base_url}/send_functions", {"tools": self._pending_tools})
            self._pending_tools = []
        return base_url

    def _wait_for_health(self, base_url: str) -> None:
        deadline = time.monotonic() + _HEALTH_TIMEOUT
        last_err: str | None = None
        while time.monotonic() < deadline:
            if not self._launcher.is_alive():
                raise RuntimeError(
                    f"Sandbox kernel process exited before becoming healthy "
                    f"(session={self.session_id}). Last error: {last_err}"
                )
            try:
                body = _http_get(f"{base_url}/health")
                if body.get("status") == "ok":
                    return
            except (urllib.error.URLError, ConnectionError, OSError) as e:
                last_err = str(e)
            time.sleep(_HEALTH_POLL)
        raise TimeoutError(
            f"Sandbox kernel at {base_url} not healthy within {_HEALTH_TIMEOUT}s "
            f"(session={self.session_id}). Last error: {last_err}"
        )

    def close(self) -> None:
        """Tear down the kernel and reap the process/container (idempotent)."""
        try:
            self._launcher.close()
        finally:
            self._base_url = None

    def __enter__(self) -> SandboxedExecutor:
        self._ensure_started()
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    def __del__(self):  # best-effort reap if the caller forgot to close()
        try:
            self.close()
        except Exception:  # pragma: no cover
            pass

    # -- Executor protocol -------------------------------------------------- #
    def reset_environment(self) -> None:
        base_url = self._ensure_started()
        _http_post(f"{base_url}/reset", {})

    def send_functions(self, functions: dict[str, Any]) -> None:
        entries = _tool_entries_from_functions(functions)
        if self._base_url is None:
            # Queue until the kernel is up (agent injects tools before first exec).
            self._pending_tools = entries
            return
        _http_post(f"{self._base_url}/send_functions", {"tools": entries})

    def send_variables(self, variables: dict[str, Any]) -> None:
        # Only JSON-safe variables cross the boundary; complex objects (AnnData,
        # DataFrames) live in the kernel namespace via executed code, not here.
        safe: dict[str, Any] = {}
        for k, v in (variables or {}).items():
            try:
                json.dumps(v)
                safe[k] = v
            except (TypeError, ValueError):
                continue
        if not safe:
            return
        base_url = self._ensure_started()
        _http_post(f"{base_url}/send_variables", {"variables": safe})

    def execute(self, code: str) -> str:
        base_url = self._ensure_started()
        body = _http_post(f"{base_url}/execute", {"code": code})
        return body.get("stdout", "")

    def __call__(self, code: str) -> str:
        return self.execute(code)


def get_sandboxed_executor(**kwargs) -> SandboxedExecutor:
    """Construct a `SandboxedExecutor` with env-driven defaults (factory helper)."""
    return SandboxedExecutor(**kwargs)