From a5f81109577182bf5e3908022403f45f8dc7469f Mon Sep 17 00:00:00 2001 From: Hudson Griffith <28578738+HudsonGri@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:28:22 +0000 Subject: [PATCH 1/2] feat(harness): add sandbox.prized(), a remote microVM sandbox for litellm.agent() --- README.md | 2 + litellm/harness/sandbox/__init__.py | 3 + litellm/harness/sandbox/docker.py | 28 +- litellm/harness/sandbox/prized.py | 441 ++++++++++++++++++++++ litellm/harness/sandbox/snapshot.py | 26 ++ litellm/sandbox/__init__.py | 4 +- tests/unit/harness/sandbox/test_prized.py | 153 ++++++++ 7 files changed, 630 insertions(+), 27 deletions(-) create mode 100644 litellm/harness/sandbox/prized.py create mode 100644 tests/unit/harness/sandbox/test_prized.py diff --git a/README.md b/README.md index 4feb6e00ffa..1daf623e40f 100644 --- a/README.md +++ b/README.md @@ -307,6 +307,8 @@ print(result.text, result.cost, [f.path for f in result.files]) Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex`, `opencode` or `pi`), or `deepagents langchain-litellm` for Deep Agents. +To keep the agent's commands and file edits off your machine, pass `sandbox=sandbox.docker(image)`, or `sandbox=sandbox.prized()` for a remote [Prized](https://prized.dev/docs/sandboxes?via=litellm) microVM (set `PRIZED_API_KEY`, works with Deep Agents). + [**Docs: Agent Harnesses**](https://docs.litellm.ai/docs/harness) diff --git a/litellm/harness/sandbox/__init__.py b/litellm/harness/sandbox/__init__.py index 434917d625a..eb000a32955 100644 --- a/litellm/harness/sandbox/__init__.py +++ b/litellm/harness/sandbox/__init__.py @@ -3,6 +3,7 @@ from litellm.harness.sandbox.base import CompletedRun, Process, Sandbox from litellm.harness.sandbox.docker import DockerSandbox, docker from litellm.harness.sandbox.local import LocalSandbox, local +from litellm.harness.sandbox.prized import PrizedSandbox, prized from litellm.harness.sandbox.snapshot import ( build_file_changes, capture_text_contents, @@ -14,6 +15,7 @@ __all__ = ( "CompletedRun", "DockerSandbox", "LocalSandbox", + "PrizedSandbox", "Process", "Sandbox", "build_file_changes", @@ -21,5 +23,6 @@ __all__ = ( "diff_snapshots", "docker", "local", + "prized", "snapshot_local", ) diff --git a/litellm/harness/sandbox/docker.py b/litellm/harness/sandbox/docker.py index ac8f357200b..f5d5be62142 100644 --- a/litellm/harness/sandbox/docker.py +++ b/litellm/harness/sandbox/docker.py @@ -10,40 +10,16 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import Final -from litellm.constants import HARNESS_SNAPSHOT_SKIP_DIRS from litellm.harness.errors import SandboxError from litellm.harness.sandbox.base import CompletedRun from litellm.harness.sandbox.local import SubprocessHandle, collect_output -from litellm.harness.sandbox.snapshot import HARNESS_SNAPSHOT_MAX_FILE_BYTES +from litellm.harness.sandbox.snapshot import parse_sha256sum, remote_snapshot_script DOCKER_HOST_ALIAS: Final = "host.docker.internal" -_SHA256_HEX_LEN: Final = 64 _WRITE_SCRIPT: Final = 'mkdir -p "$(dirname "$1")" && cat > "$1"' _WHICH_SCRIPT: Final = 'command -v "$1"' -def _snapshot_script() -> str: - prune = " -o ".join(f"-name '{name}'" for name in sorted(HARNESS_SNAPSHOT_SKIP_DIRS)) - return ( - 'cd "$1" && find . -type d \\( ' - + prune - + " \\) -prune -o -type f -size -" - + f"{HARNESS_SNAPSHOT_MAX_FILE_BYTES + 1}c" - + " -exec sha256sum {} +" - ) - - -def parse_sha256sum(output: str) -> Mapping[str, str]: - """Parse `sha256sum` lines (" ./rel/path") into {rel/path: hex}.""" - return MappingProxyType( - { - line[_SHA256_HEX_LEN + 2 :].removeprefix("./"): line[:_SHA256_HEX_LEN] - for line in output.splitlines() - if len(line) > _SHA256_HEX_LEN + 2 - } - ) - - class DockerSandbox: """Sandbox backed by a long-lived `sleep infinity` container.""" @@ -243,7 +219,7 @@ class DockerSandbox: return path async def snapshot(self) -> Mapping[str, str]: - code, out, err = await self._exec_capture(("sh", "-c", _snapshot_script(), "sh", self.workdir)) + code, out, err = await self._exec_capture(("sh", "-c", remote_snapshot_script(), "sh", self.workdir)) if code != 0: raise SandboxError(f"snapshot failed: {err.decode(errors='replace').strip()}") return parse_sha256sum(out.decode("utf-8", errors="replace")) diff --git a/litellm/harness/sandbox/prized.py b/litellm/harness/sandbox/prized.py new file mode 100644 index 00000000000..667a5f97ae9 --- /dev/null +++ b/litellm/harness/sandbox/prized.py @@ -0,0 +1,441 @@ +"""PrizedSandbox: run harness commands and files in a Prized sandbox (a Firecracker microVM) over its HTTP API.""" + +from __future__ import annotations + +import asyncio +import math +import posixpath +import shlex +import uuid +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final + +import httpx +from pydantic import BaseModel, ConfigDict, Field, ValidationError + +from litellm.harness.errors import CapabilityUnsupported, SandboxError +from litellm.harness.sandbox.base import CompletedRun +from litellm.harness.sandbox.snapshot import parse_sha256sum, remote_snapshot_script +from litellm.secret_managers.main import get_secret_str + +PRIZED_API_BASE: Final = "https://api.prized.dev/api/v1" +PRIZED_DEFAULT_TEMPLATE: Final = "prized-base-v1" +PRIZED_HOME: Final = "/home/prized" +PRIZED_WORKDIR: Final = "/home/prized/work" +PRIZED_EXEC_MAX_SECONDS: Final = 3600 +_START_TIMEOUT_SECONDS: Final = 180.0 +_READY_TIMEOUT_SECONDS: Final = 30.0 +_POLL_SECONDS: Final = 0.5 +_ANSWER_SLACK_SECONDS: Final = 30.0 +_STOP_TIMEOUT_SECONDS: Final = 15.0 +_RETRYABLE_EDGE_STATUSES: Final = frozenset({502, 503}) +_RUN_DIR: Final = "/tmp/litellm-harness-runs" +_WHICH_SCRIPT: Final = 'command -v "$1"' +_STOP_SCRIPT: Final = ( + 'f="$1"; i=0; while [ ! -s "$f" ] && [ $i -lt 20 ]; do sleep 0.1; i=$((i+1)); done; ' + '[ -s "$f" ] || exit 0; pg=$(cat "$f"); kill -TERM -- "-$pg" 2>/dev/null; i=0; ' + 'while kill -0 -- "-$pg" 2>/dev/null && [ $i -lt 20 ]; do sleep 0.1; i=$((i+1)); done; ' + 'kill -KILL -- "-$pg" 2>/dev/null; rm -f "$f"' +) + + +class _SandboxDoc(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore", populate_by_name=True) + + id: str + state: str + toolbox_base: str | None = Field(default=None, alias="toolboxBase") + error_reason: str | None = Field(default=None, alias="errorReason") + + +class _SandboxEnvelope(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + sandbox: _SandboxDoc + + +class _ExecAnswer(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore", populate_by_name=True) + + stdout: str = "" + stderr: str = "" + exit_code: int = Field(alias="exitCode") + timed_out: bool = Field(default=False, alias="timedOut") + + +class _PrizedProcess: + """A command run to completion in the sandbox; its output arrives on the readers when it exits.""" + + def __init__(self, run: asyncio.Task[CompletedRun]) -> None: + self.stdin: asyncio.StreamWriter | None = None + self.stdout: Final = asyncio.StreamReader() + self.stderr: Final = asyncio.StreamReader() + self._run: Final = run + self.done: Final = run + run.add_done_callback(self._deliver) + + def _deliver(self, run: asyncio.Task[CompletedRun]) -> None: + if not run.cancelled() and run.exception() is None: + self.stdout.feed_data(run.result().stdout.encode()) + self.stderr.feed_data(run.result().stderr.encode()) + self.stdout.feed_eof() + self.stderr.feed_eof() + + async def wait(self) -> int: + try: + return (await asyncio.shield(self._run)).exit_code + except asyncio.CancelledError: + if self._run.cancelled(): + return -9 + raise + + async def kill(self) -> None: + self._run.cancel() + await asyncio.gather(self._run, return_exceptions=True) + + +class PrizedSandbox: + """Sandbox in a Prized microVM, created lazily on first use and deleted on close(). + + Pass `sandbox_id` to attach to a sandbox you already have instead; close() then leaves it running. + Commands run under `bash` as the user `prized`, with stdin at EOF and output returned when they exit. + Files move through the sandbox's file API, so `workdir` and every path must be under /home/prized. + The sandbox cannot reach this machine, so CLI harnesses that call back to a local model endpoint + (Claude Code, Codex, OpenCode) are not supported; Deep Agents and Tool Loop are. + """ + + is_container = True + + def __init__( + self, + *, + template: str = PRIZED_DEFAULT_TEMPLATE, + sandbox_id: str | None = None, + workdir: str = PRIZED_WORKDIR, + env: Mapping[str, str] | None = None, + labels: Mapping[str, str] | None = None, + vcpu: int | None = None, + memory_mb: int | None = None, + max_runtime_min: int | None = None, + api_key: str | None = None, + api_base: str | None = None, + client: httpx.AsyncClient | None = None, + ) -> None: + normalized: Final = posixpath.normpath(workdir) + if not posixpath.isabs(normalized) or not _under_home(normalized): + raise SandboxError(f"prized workdir must be an absolute path under {PRIZED_HOME}: {workdir}") + self.workdir: str = normalized + self.template = template + self.env: Mapping[str, str] = MappingProxyType(dict(env or ())) + self.labels: Mapping[str, str] = MappingProxyType(dict(labels or ())) + self.sandbox_id: str | None = sandbox_id + self._owns_sandbox = sandbox_id is None + self._resources: Final = MappingProxyType( + {k: v for k, v in (("vcpu", vcpu), ("memoryMb", memory_mb)) if v is not None} + ) + self._max_runtime_min = max_runtime_min + self._api_key = api_key + self._api_base = (api_base or PRIZED_API_BASE).rstrip("/") + self._client = client + self._owns_client = False + self._toolbox: str | None = None + self._start_lock: Final = asyncio.Lock() + self._processes: Final[set[_PrizedProcess]] = set() # mutable-ok: live-process registry, killed on close() + self._closed = False + + def __repr__(self) -> str: + return f"PrizedSandbox({self.sandbox_id or self.template!r}, workdir={self.workdir!r})" + + def _http(self) -> httpx.AsyncClient: + if self._client is None: + self._client = httpx.AsyncClient() + self._owns_client = True + return self._client + + def _headers(self) -> Mapping[str, str]: + key: Final = self._api_key or get_secret_str("PRIZED_API_KEY") + if not key: + raise SandboxError("Prized API key not set. Set PRIZED_API_KEY or pass api_key=...") + return {"Authorization": f"Bearer {key}"} + + def sandbox_path(self, path: str) -> str: + """Absolute sandbox path; relative paths resolve against workdir.""" + return posixpath.normpath(path if posixpath.isabs(path) else posixpath.join(self.workdir, path)) + + async def start(self) -> str: + """Create (or attach to) the sandbox if needed, wait until it answers commands, and return its id.""" + if self._closed: + raise SandboxError("sandbox is closed") + async with self._start_lock: + if self.sandbox_id is not None and self._toolbox is not None: + return self.sandbox_id + doc: Final = await (self._get(self.sandbox_id) if self.sandbox_id else self._create()) + self.sandbox_id = doc.id + running: Final = await self._wait_running(doc) + if running.toolbox_base is None: + raise SandboxError(f"prized sandbox {doc.id} has no command endpoint") + self._toolbox = running.toolbox_base.rstrip("/") + await self._wait_ready() + if self._owns_sandbox: + await self._run_line(f"mkdir -p -- {shlex.quote(self.workdir)}", cwd=PRIZED_HOME, timeout=60) + return doc.id + + async def _create(self) -> _SandboxDoc: + body: Final = { + "template": self.template, + "labels": {"created-by": "litellm-harness", **self.labels}, + **({"envVars": dict(self.env)} if self.env else {}), + **({"resources": dict(self._resources)} if self._resources else {}), + **({"maxRuntimeMin": self._max_runtime_min} if self._max_runtime_min else {}), + } + response: Final = await self._http().post( + f"{self._api_base}/sandboxes", + params={"wait": "running"}, + json=body, + headers={**self._headers(), "Idempotency-Key": f"litellm-{uuid.uuid4().hex}"}, + timeout=_START_TIMEOUT_SECONDS, + ) + return _sandbox_doc(response, "create") + + async def _get(self, sandbox_id: str) -> _SandboxDoc: + response: Final = await self._http().get( + f"{self._api_base}/sandboxes/{sandbox_id}", headers=self._headers(), timeout=_READY_TIMEOUT_SECONDS + ) + return _sandbox_doc(response, f"get {sandbox_id}") + + async def _wait_running(self, doc: _SandboxDoc) -> _SandboxDoc: + deadline: Final = asyncio.get_running_loop().time() + _START_TIMEOUT_SECONDS + current = doc # rebind-ok: polled until the sandbox settles + while current.state != "running": + if current.state in ("deleting", "deleted", "error"): + raise SandboxError(f"prized sandbox {doc.id} is {current.state}: {current.error_reason or ''}".strip()) + if current.state == "paused": + return current + if asyncio.get_running_loop().time() > deadline: + raise SandboxError(f"prized sandbox {doc.id} did not start within {_START_TIMEOUT_SECONDS:g}s") + await asyncio.sleep(_POLL_SECONDS) + current = await self._get(doc.id) + return current + + async def _wait_ready(self) -> None: + """The edge can answer 502 for a moment after a sandbox reaches running; a harmless `true` waits it out.""" + deadline: Final = asyncio.get_running_loop().time() + _READY_TIMEOUT_SECONDS + while True: + response = await self._post_execute({"command": "true", "timeout": 10}, timeout=40.0) + if response.status_code not in _RETRYABLE_EDGE_STATUSES: + _exec_answer(response) + return + if asyncio.get_running_loop().time() > deadline: + raise SandboxError(f"prized sandbox {self.sandbox_id} did not answer commands ({response.status_code})") + await asyncio.sleep(_POLL_SECONDS) + + async def _post_execute(self, body: Mapping[str, object], *, timeout: float) -> httpx.Response: + return await self._http().post( + f"{self._toolbox}/process/execute", json=dict(body), headers=self._headers(), timeout=timeout + ) + + async def _run_line( + self, line: str, *, cwd: str, env: Mapping[str, str] | None = None, timeout: int = PRIZED_EXEC_MAX_SECONDS + ) -> _ExecAnswer: + body: Final = {"command": line, "cwd": cwd, "timeout": timeout, **({"env": dict(env)} if env else {})} + return _exec_answer(await self._post_execute(body, timeout=timeout + _ANSWER_SLACK_SECONDS)) + + async def exec( + self, + cmd: Sequence[str], + *, + env: Mapping[str, str] | None = None, + cwd: str | None = None, + ) -> _PrizedProcess: + if not cmd: + raise SandboxError("exec() needs a non-empty command") + await self.start() + process: Final = _PrizedProcess(asyncio.ensure_future(self.run(cmd, env=env, cwd=cwd))) + self._processes.add(process) + process.done.add_done_callback(lambda _: self._processes.discard(process)) + return process + + async def run( + self, + cmd: Sequence[str], + *, + env: Mapping[str, str] | None = None, + cwd: str | None = None, + timeout: float | None = None, + ) -> CompletedRun: + if not cmd: + raise SandboxError("run() needs a non-empty command") + await self.start() + limit: Final = PRIZED_EXEC_MAX_SECONDS if timeout is None else max(1, math.ceil(timeout)) + if limit > PRIZED_EXEC_MAX_SECONDS: + raise SandboxError(f"prized runs one command for at most {PRIZED_EXEC_MAX_SECONDS}s") + pid_file: Final = f"{_RUN_DIR}/{uuid.uuid4().hex}.pid" + line: Final = ( + f'mkdir -p {_RUN_DIR} && printf %s "$$" > {pid_file}; ' + f"{shlex.join(cmd)}; status=$?; rm -f {pid_file}; exit $status" + ) + try: + answer: Final = await self._run_line( + line, cwd=self.sandbox_path(cwd or self.workdir), env={**self.env, **(env or {})}, timeout=limit + ) + except BaseException: + await asyncio.shield(self._stop(pid_file)) + raise + if answer.timed_out: + raise SandboxError(f"command timed out after {timeout}s: {cmd[0]}") + return CompletedRun(stdout=answer.stdout, stderr=answer.stderr, exit_code=answer.exit_code) + + async def _stop(self, pid_file: str) -> None: + """Kill the process group of a command whose call was cancelled or lost; best effort.""" + try: + await asyncio.wait_for( + self._run_line(shlex.join(("sh", "-c", _STOP_SCRIPT, "sh", pid_file)), cwd=PRIZED_HOME, timeout=10), + _STOP_TIMEOUT_SECONDS, + ) + except (SandboxError, httpx.HTTPError, asyncio.TimeoutError): + return + + async def read(self, path: str) -> bytes: + await self.start() + target: Final = self.sandbox_path(path) + response: Final = await self._http().get( + f"{self._toolbox}/files/download", params={"path": target}, headers=self._headers(), timeout=600.0 + ) + if response.status_code != 200: + raise SandboxError(f"could not read {target}: {_edge_error(response)}") + return response.content + + async def write(self, path: str, data: bytes) -> None: + await self.start() + target: Final = self.sandbox_path(path) + response: Final = await self._http().post( + f"{self._toolbox}/files/upload", + params={"path": target}, + content=data, + headers={**self._headers(), "Content-Type": "application/octet-stream"}, + timeout=600.0, + ) + if response.status_code != 200: + raise SandboxError(f"could not write {target}: {_edge_error(response)}") + + def host_url(self, port: int) -> str: + raise CapabilityUnsupported( + "a Prized sandbox cannot reach this machine, so CLI harnesses (Claude Code, Codex, OpenCode) " + "cannot call the local model endpoint from it; use Harness.DEEPAGENTS or Harness.TOOL_LOOP" + ) + + async def which(self, binary: str) -> str | None: + found: Final = await self.run(("sh", "-lc", _WHICH_SCRIPT, "sh", binary), timeout=30) + path: Final = found.stdout.strip() + return path if found.exit_code == 0 and path else None + + async def tempdir(self) -> str: + """A fresh directory under /home/prized, so the file API reaches it.""" + made: Final = await self.run(("mktemp", "-d", "-p", PRIZED_HOME, "litellm-harness-XXXXXX"), timeout=30) + path: Final = made.stdout.strip() + if made.exit_code != 0 or not path: + raise SandboxError(f"mktemp -d failed: {made.stderr.strip()}") + return path + + async def snapshot(self) -> Mapping[str, str]: + listed: Final = await self.run(("sh", "-c", remote_snapshot_script(), "sh", self.workdir), timeout=300) + if listed.exit_code != 0: + raise SandboxError(f"snapshot failed: {listed.stderr.strip()}") + return parse_sha256sum(listed.stdout) + + async def close(self) -> None: + """Stop running commands and delete the sandbox this object created; an attached one keeps running.""" + if self._closed: + return + self._closed = True + await asyncio.gather(*(process.kill() for process in self._processes), return_exceptions=True) + self._processes.clear() + try: + if self._owns_sandbox and self.sandbox_id is not None: + await self._delete(self.sandbox_id) + finally: + await self._close_client() + + async def _delete(self, sandbox_id: str) -> None: + try: + response: Final = await self._http().delete( + f"{self._api_base}/sandboxes/{sandbox_id}", headers=self._headers(), timeout=30.0 + ) + except httpx.HTTPError as exc: + raise SandboxError(f"could not delete prized sandbox {sandbox_id}: {exc}") from exc + if response.status_code >= 400 and response.status_code != 404: + raise SandboxError(f"could not delete prized sandbox {sandbox_id}: {_edge_error(response)}") + + async def _close_client(self) -> None: + if self._owns_client and self._client is not None: + await self._client.aclose() + + async def __aenter__(self) -> PrizedSandbox: + await self.start() + return self + + async def __aexit__(self, *exc_info: object) -> None: + await self.close() + + +def _under_home(path: str) -> bool: + return path == PRIZED_HOME or path.startswith(PRIZED_HOME + "/") + + +class _ErrorBody(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + error: str = "" + + +def _edge_error(response: httpx.Response) -> str: + try: + code: Final = _ErrorBody.model_validate_json(response.content).error + except ValidationError: + return f"HTTP {response.status_code}" + return f"HTTP {response.status_code} {code}".strip() + + +def _sandbox_doc(response: httpx.Response, action: str) -> _SandboxDoc: + if response.status_code >= 400: + raise SandboxError(f"prized sandbox {action} failed: {_edge_error(response)}") + try: + return _SandboxEnvelope.model_validate_json(response.content).sandbox + except ValidationError as exc: + raise SandboxError(f"prized sandbox {action} returned an unexpected body") from exc + + +def _exec_answer(response: httpx.Response) -> _ExecAnswer: + if response.status_code != 200: + raise SandboxError(f"prized command failed: {_edge_error(response)}") + try: + return _ExecAnswer.model_validate_json(response.content) + except ValidationError as exc: + raise SandboxError("prized command returned an unexpected body") from exc + + +def prized( + template: str = PRIZED_DEFAULT_TEMPLATE, + *, + sandbox_id: str | None = None, + workdir: str = PRIZED_WORKDIR, + env: Mapping[str, str] | None = None, + labels: Mapping[str, str] | None = None, + vcpu: int | None = None, + memory_mb: int | None = None, + max_runtime_min: int | None = None, + api_key: str | None = None, +) -> PrizedSandbox: + """Sandbox in a new Prized microVM from `template` (or the existing `sandbox_id`), started lazily on first use.""" + return PrizedSandbox( + template=template, + sandbox_id=sandbox_id, + workdir=workdir, + env=env, + labels=labels, + vcpu=vcpu, + memory_mb=memory_mb, + max_runtime_min=max_runtime_min, + api_key=api_key, + ) diff --git a/litellm/harness/sandbox/snapshot.py b/litellm/harness/sandbox/snapshot.py index 61407c6c063..004fe06edd2 100644 --- a/litellm/harness/sandbox/snapshot.py +++ b/litellm/harness/sandbox/snapshot.py @@ -32,6 +32,32 @@ _HASH_CHUNK_BYTES: Final = 1024 * 1024 _NO_NEWLINE_MARKER: Final = "\\ No newline at end of file\n" +_SHA256_HEX_LEN: Final = 64 + + +def remote_snapshot_script() -> str: + """`sh -c` script (workdir as $1) printing `sha256sum` lines for the files a snapshot covers.""" + prune: Final = " -o ".join(f"-name '{name}'" for name in sorted(HARNESS_SNAPSHOT_SKIP_DIRS)) + return ( + 'cd "$1" && find . -type d \\( ' + + prune + + " \\) -prune -o -type f -size -" + + f"{HARNESS_SNAPSHOT_MAX_FILE_BYTES + 1}c" + + " -exec sha256sum {} +" + ) + + +def parse_sha256sum(output: str) -> Mapping[str, str]: + """Parse `sha256sum` lines (" ./rel/path") into {rel/path: hex}.""" + return MappingProxyType( + { + line[_SHA256_HEX_LEN + 2 :].removeprefix("./"): line[:_SHA256_HEX_LEN] + for line in output.splitlines() + if len(line) > _SHA256_HEX_LEN + 2 + } + ) + + def _hash_file(path: str) -> str: digest = hashlib.sha256() with open(path, "rb") as fh: diff --git a/litellm/sandbox/__init__.py b/litellm/sandbox/__init__.py index bcbb7bc5743..94fbb675ada 100644 --- a/litellm/sandbox/__init__.py +++ b/litellm/sandbox/__init__.py @@ -1,6 +1,6 @@ """litellm.sandbox: code-interpreter providers (see main.py) plus harness sandboxes. -`sandbox.local(path)` and `sandbox.docker(image, ...)` re-export litellm.harness.sandbox. +`sandbox.local(path)`, `sandbox.docker(image, ...)` and `sandbox.prized(...)` re-export litellm.harness.sandbox. They resolve lazily so `import litellm` does not pull in litellm.harness. """ @@ -15,8 +15,10 @@ _HARNESS_EXPORTS: Final = frozenset( { "local", "docker", + "prized", "LocalSandbox", "DockerSandbox", + "PrizedSandbox", "Sandbox", "Process", "CompletedRun", diff --git a/tests/unit/harness/sandbox/test_prized.py b/tests/unit/harness/sandbox/test_prized.py new file mode 100644 index 00000000000..660645894c8 --- /dev/null +++ b/tests/unit/harness/sandbox/test_prized.py @@ -0,0 +1,153 @@ +import asyncio +import json +import shlex +from collections.abc import Callable +from typing import Final + +import httpx +import pytest + +from litellm import sandbox +from litellm.harness.errors import CapabilityUnsupported, SandboxError +from litellm.harness.sandbox import PrizedSandbox, Sandbox + +API: Final = "https://api.prized.dev/api/v1" +EDGE: Final = "https://term.prized.dev/v1/sandbox/sbx_1" +ExecReply = Callable[[dict[str, object]], httpx.Response] + + +class FakePrized: + """In-memory Prized API: one sandbox, a file store, and scripted command answers.""" + + def __init__(self, exec_reply: ExecReply | None = None, *, unready_answers: int = 0) -> None: + self.creates: list[dict[str, object]] = [] + self.commands: list[dict[str, object]] = [] + self.files: dict[str, bytes] = {} + self.deleted: list[str] = [] + self.unready_answers = unready_answers + self.exec_reply = exec_reply or (lambda body: httpx.Response(200, json={"stdout": "", "exitCode": 0})) + + def _sandbox(self, state: str = "running") -> dict[str, object]: + return {"sandbox": {"id": "sbx_1", "state": state, "toolboxBase": EDGE}} + + def handler(self, request: httpx.Request) -> httpx.Response: + assert request.headers["authorization"] == "Bearer dcp_test" + path = request.url.path + if request.method == "POST" and path == "/api/v1/sandboxes": + self.creates.append(json.loads(request.content)) + return httpx.Response(201, json=self._sandbox()) + if request.method == "GET" and path == "/api/v1/sandboxes/sbx_1": + return httpx.Response(200, json=self._sandbox()) + if request.method == "DELETE" and path == "/api/v1/sandboxes/sbx_1": + self.deleted.append("sbx_1") + return httpx.Response(202, json=self._sandbox("deleting")) + if path.endswith("/process/execute"): + body = json.loads(request.content) + if body["command"] == "true" and self.unready_answers: + self.unready_answers -= 1 + return httpx.Response(502, json={"error": "sandbox_unreachable"}) + self.commands.append(body) + return self.exec_reply(body) + if path.endswith("/files/upload"): + self.files[request.url.params["path"]] = request.content + return httpx.Response(200, json={"path": request.url.params["path"], "size": len(request.content)}) + if path.endswith("/files/download"): + data = self.files.get(request.url.params["path"]) + return ( + httpx.Response(404, json={"error": "not_found"}) if data is None else httpx.Response(200, content=data) + ) + return httpx.Response(404, json={"error": "not_found"}) + + +def make_box(fake: FakePrized, **kwargs: object) -> PrizedSandbox: + client = httpx.AsyncClient(transport=httpx.MockTransport(fake.handler)) + return PrizedSandbox(api_key="dcp_test", client=client, **kwargs) # pyright: ignore[reportArgumentType] # test kwargs + + +def user_commands(fake: FakePrized) -> list[str]: + return [ + str(c["command"]) for c in fake.commands if c["command"] != "true" and "mkdir -p --" not in str(c["command"]) + ] + + +def test_prized_satisfies_protocol_and_is_exported(): + box = sandbox.prized() + assert isinstance(box, Sandbox) + assert box.workdir == "/home/prized/work" + + +@pytest.mark.parametrize("workdir", ["/workspace", "work", "/home/prized/../etc"]) +def test_workdir_must_be_under_the_sandbox_home(workdir): + with pytest.raises(SandboxError): + sandbox.prized(workdir=workdir) + + +async def test_concurrent_first_use_creates_one_sandbox_and_runs_argv_literally(): + fake = FakePrized(lambda body: httpx.Response(200, json={"stdout": "out", "stderr": "err", "exitCode": 7})) + box = make_box(fake, env={"A": "1"}) + payload = "a b; $(rm -rf /) 'q'" + results = await asyncio.gather(*(box.run(["printf", "%s", payload], env={"B": "2"}) for _ in range(3))) + assert len(fake.creates) == 1 + assert {(r.exit_code, r.stdout, r.stderr) for r in results} == {(7, "out", "err")} + command = user_commands(fake)[0] + assert shlex.join(["printf", "%s", payload]) in command + assert fake.commands[-1]["env"] == {"A": "1", "B": "2"} + assert fake.commands[-1]["cwd"] == "/home/prized/work" + + +async def test_first_command_waits_out_an_edge_that_is_not_ready_yet(): + fake = FakePrized(unready_answers=1) + box = make_box(fake) + assert (await box.run(["true"])).exit_code == 0 + + +async def test_a_command_past_its_timeout_raises(): + fake = FakePrized(lambda body: httpx.Response(200, json={"exitCode": -1, "timedOut": True})) + box = make_box(fake) + with pytest.raises(SandboxError, match="timed out"): + await box.run(["sleep", "10"], timeout=1) + assert fake.commands[-1]["timeout"] == 1 + + +async def test_files_round_trip_through_the_file_api_with_absolute_paths(): + fake = FakePrized() + box = make_box(fake) + await box.write("src/a.py", b"\x00data\xff") + assert fake.files == {"/home/prized/work/src/a.py": b"\x00data\xff"} + assert await box.read("/home/prized/work/src/a.py") == b"\x00data\xff" + with pytest.raises(SandboxError): + await box.read("missing.txt") + + +async def test_snapshot_maps_relative_paths_to_hashes(): + digest = "a" * 64 + fake = FakePrized(lambda body: httpx.Response(200, json={"stdout": f"{digest} ./src/a.py\n", "exitCode": 0})) + box = make_box(fake) + assert dict(await box.snapshot()) == {"src/a.py": digest} + + +async def test_exec_delivers_output_on_the_readers(): + fake = FakePrized(lambda body: httpx.Response(200, json={"stdout": "hello\n", "stderr": "", "exitCode": 3})) + box = make_box(fake) + process = await box.exec(["echo", "hello"]) + assert await process.wait() == 3 + assert await process.stdout.read() == b"hello\n" + + +async def test_close_deletes_a_created_sandbox_but_not_an_attached_one(): + created = FakePrized() + box = make_box(created) + await box.run(["true"]) + await box.close() + assert created.deleted == ["sbx_1"] + + attached = FakePrized() + other = make_box(attached, sandbox_id="sbx_1") + await other.run(["true"]) + await other.close() + assert attached.creates == [] and attached.deleted == [] + + +def test_cli_harnesses_cannot_reach_this_machine(): + with pytest.raises(CapabilityUnsupported): + sandbox.prized().host_url(4000) From 302caaf902669aaad4f1f252d8b771b8892e036d Mon Sep 17 00:00:00 2001 From: Hudson Griffith <28578738+HudsonGri@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:42:03 +0000 Subject: [PATCH 2/2] fix(harness): use litellm's shared httpx client in sandbox.prized() and test its error paths PrizedSandbox built its own httpx.AsyncClient, which the code-quality check rejects. It now takes get_async_httpx_client(httpxSpecialProvider.Sandbox), the client the e2b and OpenSandbox providers use, looked up on each request so a long-lived sandbox never holds a client the cache has evicted, and close() no longer closes it. An injected client= is still used as before. The new tests cover create and attach errors, polling to running, the start and ready deadlines, failed and cancelled commands (the remote process group is killed), delete failures, which/tempdir/snapshot failures and reading PRIZED_API_KEY, taking prized.py to full line coverage. --- litellm/harness/sandbox/prized.py | 26 +- tests/unit/harness/sandbox/test_prized.py | 312 +++++++++++++++++++++- 2 files changed, 315 insertions(+), 23 deletions(-) diff --git a/litellm/harness/sandbox/prized.py b/litellm/harness/sandbox/prized.py index 667a5f97ae9..9fd1562f7eb 100644 --- a/litellm/harness/sandbox/prized.py +++ b/litellm/harness/sandbox/prized.py @@ -17,7 +17,11 @@ from pydantic import BaseModel, ConfigDict, Field, ValidationError from litellm.harness.errors import CapabilityUnsupported, SandboxError from litellm.harness.sandbox.base import CompletedRun from litellm.harness.sandbox.snapshot import parse_sha256sum, remote_snapshot_script +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # the factory's params dict is untyped upstream +) from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.custom_http import httpxSpecialProvider PRIZED_API_BASE: Final = "https://api.prized.dev/api/v1" PRIZED_DEFAULT_TEMPLATE: Final = "prized-base-v1" @@ -137,8 +141,7 @@ class PrizedSandbox: self._max_runtime_min = max_runtime_min self._api_key = api_key self._api_base = (api_base or PRIZED_API_BASE).rstrip("/") - self._client = client - self._owns_client = False + self._client: Final = client self._toolbox: str | None = None self._start_lock: Final = asyncio.Lock() self._processes: Final[set[_PrizedProcess]] = set() # mutable-ok: live-process registry, killed on close() @@ -148,10 +151,10 @@ class PrizedSandbox: return f"PrizedSandbox({self.sandbox_id or self.template!r}, workdir={self.workdir!r})" def _http(self) -> httpx.AsyncClient: - if self._client is None: - self._client = httpx.AsyncClient() - self._owns_client = True - return self._client + """The injected client, else LiteLLM's shared cached client, looked up per request and never closed here.""" + if self._client is not None: + return self._client + return get_async_httpx_client(llm_provider=httpxSpecialProvider.Sandbox).client def _headers(self) -> Mapping[str, str]: key: Final = self._api_key or get_secret_str("PRIZED_API_KEY") @@ -351,11 +354,8 @@ class PrizedSandbox: self._closed = True await asyncio.gather(*(process.kill() for process in self._processes), return_exceptions=True) self._processes.clear() - try: - if self._owns_sandbox and self.sandbox_id is not None: - await self._delete(self.sandbox_id) - finally: - await self._close_client() + if self._owns_sandbox and self.sandbox_id is not None: + await self._delete(self.sandbox_id) async def _delete(self, sandbox_id: str) -> None: try: @@ -367,10 +367,6 @@ class PrizedSandbox: if response.status_code >= 400 and response.status_code != 404: raise SandboxError(f"could not delete prized sandbox {sandbox_id}: {_edge_error(response)}") - async def _close_client(self) -> None: - if self._owns_client and self._client is not None: - await self._client.aclose() - async def __aenter__(self) -> PrizedSandbox: await self.start() return self diff --git a/tests/unit/harness/sandbox/test_prized.py b/tests/unit/harness/sandbox/test_prized.py index 660645894c8..bcc30032b85 100644 --- a/tests/unit/harness/sandbox/test_prized.py +++ b/tests/unit/harness/sandbox/test_prized.py @@ -1,7 +1,9 @@ import asyncio +import importlib import json +import re import shlex -from collections.abc import Callable +from collections.abc import Awaitable, Callable from typing import Final import httpx @@ -10,27 +12,46 @@ import pytest from litellm import sandbox from litellm.harness.errors import CapabilityUnsupported, SandboxError from litellm.harness.sandbox import PrizedSandbox, Sandbox +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.types.llms.custom_http import httpxSpecialProvider API: Final = "https://api.prized.dev/api/v1" EDGE: Final = "https://term.prized.dev/v1/sandbox/sbx_1" -ExecReply = Callable[[dict[str, object]], httpx.Response] +ExecReply = Callable[[dict[str, object]], httpx.Response | Awaitable[httpx.Response]] +# The module, not the `prized()` factory that litellm.harness.sandbox re-exports under the same name. +prized_module: Final = importlib.import_module("litellm.harness.sandbox.prized") class FakePrized: """In-memory Prized API: one sandbox, a file store, and scripted command answers.""" - def __init__(self, exec_reply: ExecReply | None = None, *, unready_answers: int = 0) -> None: + def __init__( + self, + exec_reply: ExecReply | None = None, + *, + unready_answers: int = 0, + states: tuple[str, ...] = ("running",), + toolbox: str | None = EDGE, + delete_answer: httpx.Response | httpx.HTTPError | None = None, + read_only: bool = False, + ) -> None: self.creates: list[dict[str, object]] = [] self.commands: list[dict[str, object]] = [] self.files: dict[str, bytes] = {} self.deleted: list[str] = [] self.unready_answers = unready_answers self.exec_reply = exec_reply or (lambda body: httpx.Response(200, json={"stdout": "", "exitCode": 0})) + self.states = list(states) + self.toolbox = toolbox + self.delete_answer = delete_answer + self.read_only = read_only - def _sandbox(self, state: str = "running") -> dict[str, object]: - return {"sandbox": {"id": "sbx_1", "state": state, "toolboxBase": EDGE}} + def _sandbox(self, state: str | None = None) -> dict[str, object]: + """The next scripted state; the last one repeats.""" + current = state or (self.states.pop(0) if len(self.states) > 1 else self.states[0]) + return {"sandbox": {"id": "sbx_1", "state": current, "toolboxBase": self.toolbox, "errorReason": "boom"}} - def handler(self, request: httpx.Request) -> httpx.Response: + async def handler(self, request: httpx.Request) -> httpx.Response: assert request.headers["authorization"] == "Bearer dcp_test" path = request.url.path if request.method == "POST" and path == "/api/v1/sandboxes": @@ -40,15 +61,20 @@ class FakePrized: return httpx.Response(200, json=self._sandbox()) if request.method == "DELETE" and path == "/api/v1/sandboxes/sbx_1": self.deleted.append("sbx_1") - return httpx.Response(202, json=self._sandbox("deleting")) + if isinstance(self.delete_answer, httpx.HTTPError): + raise self.delete_answer + return self.delete_answer or httpx.Response(202, json=self._sandbox("deleting")) if path.endswith("/process/execute"): body = json.loads(request.content) if body["command"] == "true" and self.unready_answers: self.unready_answers -= 1 return httpx.Response(502, json={"error": "sandbox_unreachable"}) self.commands.append(body) - return self.exec_reply(body) + reply = self.exec_reply(body) + return reply if isinstance(reply, httpx.Response) else await reply if path.endswith("/files/upload"): + if self.read_only: + return httpx.Response(403, json={"error": "forbidden"}) self.files[request.url.params["path"]] = request.content return httpx.Response(200, json={"path": request.url.params["path"], "size": len(request.content)}) if path.endswith("/files/download"): @@ -64,12 +90,46 @@ def make_box(fake: FakePrized, **kwargs: object) -> PrizedSandbox: return PrizedSandbox(api_key="dcp_test", client=client, **kwargs) # pyright: ignore[reportArgumentType] # test kwargs +def box_answering(response: httpx.Response, **kwargs: object) -> PrizedSandbox: + """A sandbox whose API answers every request with `response`.""" + client = httpx.AsyncClient(transport=httpx.MockTransport(lambda request: response)) + return PrizedSandbox(api_key="dcp_test", client=client, **kwargs) # pyright: ignore[reportArgumentType] # test kwargs + + def user_commands(fake: FakePrized) -> list[str]: return [ str(c["command"]) for c in fake.commands if c["command"] != "true" and "mkdir -p --" not in str(c["command"]) ] +def is_setup(body: dict[str, object]) -> bool: + """The readiness probe and workdir mkdir that start() sends before any user command.""" + command = str(body["command"]) + return command == "true" or command.startswith("mkdir -p --") + + +class Hold: + """Exec reply that answers setup and stop commands at once and holds each user command until released.""" + + def __init__(self) -> None: + self.arrived = asyncio.Event() + self.release = asyncio.Event() + + async def __call__(self, body: dict[str, object]) -> httpx.Response: + if not is_setup(body) and "kill -TERM" not in str(body["command"]): + self.arrived.set() + await self.release.wait() + return httpx.Response(200, json={"exitCode": 0}) + + +@pytest.fixture +def no_waits(monkeypatch): + """Poll without sleeping and give up on the first missed deadline.""" + monkeypatch.setattr(prized_module, "_POLL_SECONDS", 0) + monkeypatch.setattr(prized_module, "_START_TIMEOUT_SECONDS", -1.0) + monkeypatch.setattr(prized_module, "_READY_TIMEOUT_SECONDS", -1.0) + + def test_prized_satisfies_protocol_and_is_exported(): box = sandbox.prized() assert isinstance(box, Sandbox) @@ -151,3 +211,239 @@ async def test_close_deletes_a_created_sandbox_but_not_an_attached_one(): def test_cli_harnesses_cannot_reach_this_machine(): with pytest.raises(CapabilityUnsupported): sandbox.prized().host_url(4000) + + +async def test_without_an_injected_client_it_uses_the_shared_client_and_leaves_it_open(): + fake = FakePrized() + shared = get_async_httpx_client(llm_provider=httpxSpecialProvider.Sandbox) + created = shared.client + shared.client = httpx.AsyncClient(transport=httpx.MockTransport(fake.handler)) + box = PrizedSandbox(api_key="dcp_test") + assert (await box.run(["true"])).exit_code == 0 + await box.close() + assert len(fake.creates) == 1 and fake.deleted == ["sbx_1"] + assert not shared.client.is_closed + await shared.client.aclose() + await created.aclose() + + +async def test_the_api_key_comes_from_prized_api_key_and_a_missing_one_fails_before_any_request(monkeypatch): + monkeypatch.delenv("PRIZED_API_KEY", raising=False) + fake = FakePrized() + keyless = PrizedSandbox(client=httpx.AsyncClient(transport=httpx.MockTransport(fake.handler))) + with pytest.raises(SandboxError, match="PRIZED_API_KEY"): + await keyless.start() + assert fake.creates == [] + + monkeypatch.setenv("PRIZED_API_KEY", "dcp_test") + from_env = PrizedSandbox(client=httpx.AsyncClient(transport=httpx.MockTransport(fake.handler))) + assert await from_env.start() == "sbx_1" + + +async def test_create_sends_the_template_labels_env_and_resources(): + fake = FakePrized() + box = make_box( + fake, template="py-3.12", env={"A": "1"}, labels={"team": "ml"}, vcpu=4, memory_mb=8192, max_runtime_min=30 + ) + assert repr(box) == "PrizedSandbox('py-3.12', workdir='/home/prized/work')" + await box.start() + assert repr(box) == "PrizedSandbox('sbx_1', workdir='/home/prized/work')" + assert fake.creates == [ + { + "template": "py-3.12", + "labels": {"created-by": "litellm-harness", "team": "ml"}, + "envVars": {"A": "1"}, + "resources": {"vcpu": 4, "memoryMb": 8192}, + "maxRuntimeMin": 30, + } + ] + assert fake.commands[-1]["command"] == "mkdir -p -- /home/prized/work" + + +async def test_start_polls_until_a_new_sandbox_is_running(monkeypatch): + monkeypatch.setattr(prized_module, "_POLL_SECONDS", 0) + fake = FakePrized(states=("creating", "starting", "running")) + box = make_box(fake) + assert await box.start() == "sbx_1" + assert fake.states == ["running"] + + +@pytest.mark.parametrize("state", ["error", "deleted"]) +async def test_start_fails_on_a_sandbox_that_will_never_run(state): + box = make_box(FakePrized(states=(state,))) + with pytest.raises(SandboxError, match=f"is {state}: boom"): + await box.start() + + +async def test_start_fails_when_the_sandbox_has_no_command_endpoint(): + box = make_box(FakePrized(states=("paused",), toolbox=None)) + with pytest.raises(SandboxError, match="no command endpoint"): + await box.start() + + +async def test_start_gives_up_on_a_sandbox_that_never_runs(no_waits): + box = make_box(FakePrized(states=("creating",))) + with pytest.raises(SandboxError, match="did not start"): + await box.start() + + +async def test_start_gives_up_on_an_edge_that_never_answers(no_waits): + box = make_box(FakePrized(unready_answers=1_000)) + with pytest.raises(SandboxError, match=re.escape("did not answer commands (502)")): + await box.start() + + +@pytest.mark.parametrize( + ("response", "message"), + [ + (httpx.Response(401, json={"error": "unauthorized"}), "create failed: HTTP 401 unauthorized"), + (httpx.Response(503, text="bad gateway"), "create failed: HTTP 503"), + (httpx.Response(201, json={"sandboxes": []}), "create returned an unexpected body"), + ], +) +async def test_create_errors_name_the_edge_error(response, message): + with pytest.raises(SandboxError, match=message): + await box_answering(response).start() + + +async def test_attaching_to_a_missing_sandbox_names_it(): + box = box_answering(httpx.Response(404, json={"error": "not_found"}), sandbox_id="sbx_9") + with pytest.raises(SandboxError, match="get sbx_9 failed: HTTP 404 not_found"): + await box.start() + + +async def test_run_and_exec_reject_an_empty_command_and_a_timeout_past_the_limit(): + box = make_box(FakePrized()) + with pytest.raises(SandboxError, match=re.escape("run() needs a non-empty command")): + await box.run([]) + with pytest.raises(SandboxError, match=re.escape("exec() needs a non-empty command")): + await box.exec([]) + with pytest.raises(SandboxError, match="at most 3600s"): + await box.run(["sleep", "4000"], timeout=4000) + + +@pytest.mark.parametrize( + ("response", "message"), + [ + (httpx.Response(500, json={"error": "internal"}), "command failed: HTTP 500 internal"), + (httpx.Response(200, json={"stdout": "no exit code"}), "command returned an unexpected body"), + ], +) +async def test_a_failed_command_raises_and_tries_to_stop_its_process(response, message): + fake = FakePrized(lambda body: httpx.Response(200, json={"exitCode": 0}) if is_setup(body) else response) + box = make_box(fake) + with pytest.raises(SandboxError, match=message): + await box.run(["make", "test"]) + assert "kill -TERM" in user_commands(fake)[-1] + + +async def test_close_kills_a_running_command_and_deletes_the_sandbox(): + hold = Hold() + fake = FakePrized(hold) + box = make_box(fake) + process = await box.exec(["sleep", "600"]) + await asyncio.wait_for(hold.arrived.wait(), 5) + await box.close() + assert await process.wait() == -9 + assert await process.stdout.read() == b"" + run_line, stop_line = user_commands(fake) + pid_file = re.search(r"/tmp/litellm-harness-runs/[0-9a-f]+\.pid", run_line) + assert pid_file is not None and pid_file.group(0) in stop_line + assert fake.deleted == ["sbx_1"] + + +async def test_cancelling_a_wait_leaves_the_command_running(): + hold = Hold() + fake = FakePrized(hold) + box = make_box(fake) + process = await box.exec(["sleep", "600"]) + waiter = asyncio.ensure_future(process.wait()) + await asyncio.wait_for(hold.arrived.wait(), 5) + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + assert not process.done.done() + hold.release.set() + assert await process.wait() == 0 + assert len(user_commands(fake)) == 1 + + +async def test_a_closed_sandbox_cannot_start_again(): + fake = FakePrized() + async with make_box(fake) as box: + await box.run(["true"]) + await box.close() + assert fake.deleted == ["sbx_1"] + with pytest.raises(SandboxError, match="closed"): + await box.run(["true"]) + + +@pytest.mark.parametrize( + ("answer", "error"), + [ + (httpx.Response(404, json={"error": "not_found"}), None), + (httpx.Response(500, json={"error": "internal"}), "HTTP 500 internal"), + (httpx.ConnectError("edge down"), "edge down"), + ], +) +async def test_close_ignores_a_sandbox_already_gone_and_reports_other_delete_failures(answer, error): + fake = FakePrized(delete_answer=answer) + box = make_box(fake) + await box.start() + if error is None: + await box.close() + else: + with pytest.raises(SandboxError, match=f"could not delete prized sandbox sbx_1: {error}"): + await box.close() + assert fake.deleted == ["sbx_1"] + + +async def test_a_rejected_write_names_the_edge_error(): + box = make_box(FakePrized(read_only=True)) + with pytest.raises(SandboxError, match=re.escape("could not write /home/prized/work/a.txt: HTTP 403 forbidden")): + await box.write("a.txt", b"data") + + +@pytest.mark.parametrize( + ("stdout", "exit_code", "found"), + [("/usr/bin/git\n", 0, "/usr/bin/git"), ("", 1, None)], +) +async def test_which_returns_the_path_only_when_found(stdout, exit_code, found): + fake = FakePrized(lambda body: httpx.Response(200, json={"stdout": stdout, "exitCode": exit_code})) + box = make_box(fake) + assert await box.which("git") == found + assert shlex.join(["sh", "-lc", 'command -v "$1"', "sh", "git"]) in user_commands(fake)[-1] + + +async def test_tempdir_returns_a_new_directory_under_home_and_raises_when_mktemp_fails(): + fake = FakePrized( + lambda body: httpx.Response(200, json={"stdout": "/home/prized/litellm-harness-x1\n", "exitCode": 0}) + ) + made = make_box(fake) + assert await made.tempdir() == "/home/prized/litellm-harness-x1" + assert "mktemp -d -p /home/prized litellm-harness-XXXXXX" in user_commands(fake)[-1] + refused = make_box( + FakePrized( + lambda body: ( + httpx.Response(200, json={"exitCode": 0}) + if is_setup(body) + else httpx.Response(200, json={"stderr": "no space\n", "exitCode": 1}) + ) + ) + ) + with pytest.raises(SandboxError, match="mktemp -d failed: no space"): + await refused.tempdir() + + +async def test_a_failed_snapshot_raises(): + box = make_box( + FakePrized( + lambda body: ( + httpx.Response(200, json={"exitCode": 0}) + if is_setup(body) + else httpx.Response(200, json={"stderr": "sha256sum: not found\n", "exitCode": 127}) + ) + ) + ) + with pytest.raises(SandboxError, match="snapshot failed: sha256sum: not found"): + await box.snapshot()