mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 302caaf902 into e9b1c11098
This commit is contained in:
Hudson
GitHub
commit
249a2c9ffd
7 changed files with 922 additions and 27 deletions
|
|
@ -351,6 +351,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)
|
||||
|
||||
</details>
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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 ("<hex> ./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"))
|
||||
|
|
|
|||
437
litellm/harness/sandbox/prized.py
Normal file
437
litellm/harness/sandbox/prized.py
Normal file
|
|
@ -0,0 +1,437 @@
|
|||
"""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.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"
|
||||
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: 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()
|
||||
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:
|
||||
"""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")
|
||||
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()
|
||||
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:
|
||||
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 __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,
|
||||
)
|
||||
|
|
@ -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 ("<hex> ./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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
449
tests/unit/harness/sandbox/test_prized.py
Normal file
449
tests/unit/harness/sandbox/test_prized.py
Normal file
|
|
@ -0,0 +1,449 @@
|
|||
import asyncio
|
||||
import importlib
|
||||
import json
|
||||
import re
|
||||
import shlex
|
||||
from collections.abc import Awaitable, 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
|
||||
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 | 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,
|
||||
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 | 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"}}
|
||||
|
||||
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":
|
||||
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")
|
||||
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)
|
||||
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"):
|
||||
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 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)
|
||||
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)
|
||||
|
||||
|
||||
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="<html>bad gateway</html>"), "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()
|
||||
Loading…
Add table
Reference in a new issue