File size: 8,780 Bytes
a495d22 c3b49d6 a495d22 8a0b154 a495d22 8a0b154 a495d22 c3b49d6 a495d22 c3b49d6 a495d22 c3b49d6 a495d22 c3b49d6 a495d22 c3b49d6 a495d22 c3b49d6 a495d22 c3b49d6 a495d22 c3b49d6 a495d22 c3b49d6 a495d22 c3b49d6 a495d22 c3b49d6 a495d22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 | """
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)
|