asbox 0.1.0__py3-none-any.whl

agent_sandbox/__init__.py

@@ -0,0 +1,29 @@
+"""agent_sandbox: embeddable Python SDK for microVM sandboxes on AWS Lambda MicroVMs.
+
+A Python fork of the microsandbox developer experience. The local ``libkrun``
+runtime is replaced by AWS Lambda MicroVMs (managed Firecracker); the SDK talks
+to the ``lambda-microvms`` control plane and to a guest agent (``agentd``) baked
+into the MicroVM image over its dedicated HTTPS endpoint.
+"""
+
+from agent_sandbox.errors import (
+    AgentError,
+    ControlPlaneError,
+    SandboxError,
+    SandboxTimeoutError,
+)
+from agent_sandbox.models import ExecResult, SandboxConfig, SandboxState
+from agent_sandbox.sandbox import Sandbox
+
+__all__ = [
+    "Sandbox",
+    "SandboxConfig",
+    "SandboxState",
+    "ExecResult",
+    "SandboxError",
+    "ControlPlaneError",
+    "AgentError",
+    "SandboxTimeoutError",
+]
+
+__version__ = "0.1.0"

agent_sandbox/_bundled_guest/Dockerfile

@@ -0,0 +1,28 @@
+# MicroVM image for AWS Lambda MicroVMs.
+#
+# Lambda MicroVMs runs this Dockerfile, starts the app, and snapshots the fully
+# initialized environment. Every launched MicroVM resumes from that snapshot, so
+# agentd is already listening the moment run-microvm returns.
+#
+# Lambda MicroVMs is ARM64-only.
+FROM --platform=linux/arm64 python:3.12-slim
+
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    AGENTD_PORT=8080 \
+    AGENTD_HOOK_PORT=9000
+
+WORKDIR /opt/agentd
+
+COPY requirements.txt ./
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY agentd ./agentd
+
+# 8080: agentd application API (proxied to clients). 9000: lifecycle hooks
+# (platform-only). Both must be exposed so the platform can reach the hook port.
+EXPOSE 8080 9000
+
+# serve.py starts both listeners in one process: agentd on AGENTD_PORT and the
+# lifecycle-hook server on AGENTD_HOOK_PORT. User workloads are launched via exec.
+CMD ["python", "-m", "agentd.serve"]

agent_sandbox/_bundled_guest/agentd/__init__.py

@@ -0,0 +1,8 @@
+"""agentd: the in-VM guest agent.
+
+Runs inside the Lambda MicroVM image and exposes a small HTTP API that the SDK's
+``AgentClient`` calls to execute commands and read/write files. This is the
+Python analogue of microsandbox's ``agentd`` guest agent.
+"""
+
+__version__ = "0.1.0"

agent_sandbox/_bundled_guest/agentd/app.py

@@ -0,0 +1,120 @@
+"""FastAPI application for the guest agent.
+
+Auth note: authentication is handled by the Lambda MicroVMs ingress proxy, which
+validates the ``X-aws-proxy-auth`` token before traffic reaches this process.
+agentd therefore trusts inbound requests and does not re-validate the token.
+"""
+
+from __future__ import annotations
+
+import base64
+import os
+
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+
+from agentd import state
+from agentd.exec import run_command
+
+
+@asynccontextmanager
+async def _lifespan(_app: FastAPI):
+    # Signals the /ready lifecycle hook that agentd has fully started, so the
+    # platform snapshots a serving agentd rather than a half-booted one.
+    state.ready.set()
+    yield
+
+
+app = FastAPI(title="agentd", version="0.1.0", lifespan=_lifespan)
+
+
+class ExecRequest(BaseModel):
+    command: str
+    args: list[str] = Field(default_factory=list)
+    cwd: str | None = None
+    env: dict[str, str] | None = None
+    timeout: float | None = None
+
+
+class ExecResponse(BaseModel):
+    exit_code: int
+    stdout: str
+    stderr: str
+    encoding: str = "base64"
+
+
+class ReadRequest(BaseModel):
+    path: str
+
+
+class ReadResponse(BaseModel):
+    content: str
+    encoding: str = "base64"
+
+
+class WriteRequest(BaseModel):
+    path: str
+    content: str
+    encoding: str = "base64"
+
+
+@app.get("/healthz")
+async def healthz() -> dict[str, str]:
+    return {"status": "ok", "version": app.version}
+
+
+@app.post("/v1/exec", response_model=ExecResponse)
+async def exec_command(req: ExecRequest) -> ExecResponse:
+    try:
+        with state.track():
+            outcome = await run_command(
+                req.command,
+                req.args,
+                cwd=req.cwd,
+                env=req.env,
+                timeout=req.timeout,
+            )
+    except TimeoutError as exc:
+        raise HTTPException(status_code=504, detail="command timed out") from exc
+    except FileNotFoundError as exc:
+        raise HTTPException(status_code=400, detail=f"command not found: {req.command}") from exc
+    return ExecResponse(
+        exit_code=outcome.exit_code,
+        stdout=base64.b64encode(outcome.stdout).decode("ascii"),
+        stderr=base64.b64encode(outcome.stderr).decode("ascii"),
+    )
+
+
+@app.post("/v1/fs/read", response_model=ReadResponse)
+async def read_file(req: ReadRequest) -> ReadResponse:
+    try:
+        with open(req.path, "rb") as fh:
+            data = fh.read()
+    except FileNotFoundError as exc:
+        raise HTTPException(status_code=404, detail=f"no such file: {req.path}") from exc
+    except OSError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+    return ReadResponse(content=base64.b64encode(data).decode("ascii"))
+
+
+@app.post("/v1/fs/write")
+async def write_file(req: WriteRequest) -> dict[str, int]:
+    if req.encoding == "base64":
+        try:
+            data = base64.b64decode(req.content)
+        except (ValueError, TypeError) as exc:
+            raise HTTPException(status_code=400, detail="invalid base64 content") from exc
+    else:
+        data = req.content.encode("utf-8")
+
+    directory = os.path.dirname(req.path)
+    if directory:
+        os.makedirs(directory, exist_ok=True)
+    try:
+        with open(req.path, "wb") as fh:
+            fh.write(data)
+    except OSError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+    return {"bytes_written": len(data)}

agent_sandbox/_bundled_guest/agentd/exec.py

@@ -0,0 +1,55 @@
+"""Command execution inside the guest.
+
+Captures stdout, stderr, and the exit code. A dedicated module (mirroring
+microsandbox's split between the agent loop and exec/session handling) so PTY
+support can be added here later without touching the HTTP layer.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class ExecOutcome:
+    exit_code: int
+    stdout: bytes
+    stderr: bytes
+
+
+async def run_command(
+    command: str,
+    args: list[str] | None = None,
+    *,
+    cwd: str | None = None,
+    env: dict[str, str] | None = None,
+    timeout: float | None = None,
+) -> ExecOutcome:
+    """Run ``command`` with ``args`` and capture its output.
+
+    ``env`` is merged onto the current process environment rather than replacing
+    it, so PATH and friends remain intact.
+    """
+    argv = [command, *(args or [])]
+    merged_env = {**os.environ, **(env or {})}
+
+    proc = await asyncio.create_subprocess_exec(
+        *argv,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE,
+        cwd=cwd,
+        env=merged_env,
+    )
+    try:
+        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
+    except asyncio.TimeoutError:
+        proc.kill()
+        await proc.wait()
+        raise
+    return ExecOutcome(
+        exit_code=proc.returncode if proc.returncode is not None else -1,
+        stdout=stdout or b"",
+        stderr=stderr or b"",
+    )

agent_sandbox/_bundled_guest/agentd/hooks.py

@@ -0,0 +1,102 @@
+"""Lifecycle-hook HTTP server, called by the AWS Lambda MicroVMs platform.
+
+This is the *lifecycle plane*: a separate FastAPI app (served on ``HOOK_PORT`` by
+``serve.py``) that the platform invokes at image build and MicroVM state
+transitions. It is never reached by application clients — the ``forward`` CLI
+guard keeps auth tokens from ever being scoped to the hook port.
+
+Contract: return 200 on success; return 503 from ``/ready`` and ``/validate``
+when the app needs more time (the platform keeps polling until the configured
+timeout). Handlers are kept fast and idempotent per the platform guidance.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import sys
+import threading
+
+_PYTHON = sys.executable or "python3"
+
+from fastapi import FastAPI, Request, Response
+
+from agentd import identity, state
+from agentd.exec import run_command
+
+log = logging.getLogger("agentd.hooks")
+
+hooks_app = FastAPI(title="agentd-hooks", version="0.1.0")
+
+# The platform fixes these paths; only the port is ours to choose.
+PREFIX = "/aws/lambda-microvms/runtime/v1"
+
+# Suspend drain budget: how long /suspend waits for in-flight exec to finish.
+# Must stay under the suspendTimeoutInSeconds configured in the image's --hooks
+# (see agent_sandbox.infra.resources._default_hooks) or the platform gives up
+# first. Overridable so it can track that value without editing code.
+_SUSPEND_DRAIN_SECONDS = float(os.getenv("AGENTD_SUSPEND_DRAIN_SECONDS", "8"))
+
+# /run fires once per boot-from-snapshot; guard against duplicate delivery so a
+# stray call can't rotate identity out from under a running workload.
+_run_lock = threading.Lock()
+_run_done = False
+
+
+@hooks_app.post(f"{PREFIX}/ready")
+async def ready() -> Response:
+    """Build hook: 200 once agentd has started, else 503 so the platform retries."""
+    return Response(status_code=200 if state.ready.is_set() else 503)
+
+
+@hooks_app.post(f"{PREFIX}/validate")
+async def validate() -> Response:
+    """Build hook: smoke-test the snapshot by running a trivial exec in-process."""
+    try:
+        outcome = await run_command(_PYTHON, ["-c", "pass"], timeout=10.0)
+    except Exception as exc:  # noqa: BLE001 - any failure means "not valid yet"
+        log.warning("validate: exec failed: %s", exc)
+        return Response(status_code=503)
+    return Response(status_code=200 if outcome.exit_code == 0 else 503)
+
+
+@hooks_app.post(f"{PREFIX}/run")
+async def run(request: Request) -> Response:
+    """Runtime hook: fires once after run-from-snapshot. Mint per-VM identity."""
+    global _run_done
+    with _run_lock:
+        if _run_done:
+            return Response(status_code=200)  # idempotent: already ran this boot
+    payload = await request.body()
+    identity.regenerate("run", run_payload=payload or None)
+    # Mark done only AFTER a successful regenerate, so a transient failure leaves
+    # the flag clear and the platform's retry still gets a fresh identity.
+    with _run_lock:
+        _run_done = True
+    return Response(status_code=200)
+
+
+@hooks_app.post(f"{PREFIX}/resume")
+async def resume() -> Response:
+    """Runtime hook: fires after SUSPENDED -> RUNNING. Reseed randomness/identity."""
+    identity.regenerate("resume")
+    return Response(status_code=200)
+
+
+@hooks_app.post(f"{PREFIX}/suspend")
+async def suspend() -> Response:
+    """Runtime hook: fires before RUNNING -> SUSPENDED. Drain in-flight exec."""
+    state.wait_drain(_SUSPEND_DRAIN_SECONDS)
+    return Response(status_code=200)
+
+
+@hooks_app.post(f"{PREFIX}/terminate")
+async def terminate() -> Response:
+    """Runtime hook: fires before termination. Flush logs; best-effort cleanup."""
+    for stream in (sys.stdout, sys.stderr):
+        try:
+            stream.flush()
+        except Exception:  # noqa: BLE001 - best-effort
+            pass
+    logging.shutdown()
+    return Response(status_code=200)

agent_sandbox/_bundled_guest/agentd/identity.py

@@ -0,0 +1,66 @@
+"""Per-VM uniqueness state, driven by the ``/run`` and ``/resume`` hooks.
+
+Every MicroVM resumes from the SAME memory+disk snapshot, so anything derived
+from randomness at image-build time is IDENTICAL across VMs (the "snapshot
+uniqueness" pitfall). :func:`regenerate` mints fresh, CSPRNG-backed identity and
+reseeds the stdlib non-CSPRNG so each running VM diverges. It is called on
+``/run`` (first boot from snapshot) and ``/resume``, and is written to a tmpfs
+file so processes launched via ``/v1/exec`` can read it.
+
+Both agentd and the hooks server share this module (same process), so reseeding
+here is visible to agentd's request handlers too.
+"""
+
+from __future__ import annotations
+
+import os
+import random
+import secrets
+import threading
+import time
+
+# tmpfs path (per-VM, not part of the shared on-disk snapshot layout).
+IDENTITY_PATH = os.environ.get("AGENTD_IDENTITY_PATH", "/run/agentd/identity")
+
+_lock = threading.Lock()
+_state: dict[str, object] = {
+    "vm_id": None,
+    "generation": 0,
+    "last_event": None,
+    "booted_at": None,
+    "run_payload": None,
+}
+
+
+def regenerate(event: str, *, run_payload: bytes | None = None) -> dict[str, object]:
+    """Mint a fresh per-VM identity and reseed randomness.
+
+    Uses :mod:`secrets` (a CSPRNG pulling fresh kernel entropy) so the new id is
+    unique per VM even though the snapshot is shared, then reseeds :mod:`random`
+    so later non-CSPRNG use also diverges. Persisting is best-effort.
+    """
+    with _lock:
+        _state["vm_id"] = secrets.token_hex(16)
+        _state["generation"] = int(_state["generation"]) + 1
+        _state["last_event"] = event
+        _state["booted_at"] = time.time()
+        if run_payload is not None:
+            _state["run_payload"] = run_payload.decode("utf-8", errors="replace")
+        random.seed(secrets.randbits(256))
+        _persist(str(_state["vm_id"]))
+        return dict(_state)
+
+
+def current() -> dict[str, object]:
+    with _lock:
+        return dict(_state)
+
+
+def _persist(vm_id: str) -> None:
+    try:
+        os.makedirs(os.path.dirname(IDENTITY_PATH), exist_ok=True)
+        with open(IDENTITY_PATH, "w", encoding="utf-8") as fh:
+            fh.write(vm_id)
+    except OSError:
+        # Best-effort: identity is still available in-process via current().
+        pass

agent_sandbox/_bundled_guest/agentd/ports.py

@@ -0,0 +1,17 @@
+"""Well-known ports for the guest, as a single source of truth.
+
+``AGENT_PORT`` is agentd's application API (exec/fs), fronted by the platform's
+auth-token proxy. ``HOOK_PORT`` is the lifecycle-hook server, called only by the
+platform. Both are overridable via env; the Dockerfile bakes the defaults.
+
+Invariant: ``HOOK_PORT`` here (what ``serve.py`` binds) must equal the SDK's
+``agent_sandbox.ports.HOOK_PORT`` (what ``--hooks`` tells the platform). Both
+default to 9000; override consistently on both sides.
+"""
+
+from __future__ import annotations
+
+import os
+
+AGENT_PORT: int = int(os.getenv("AGENTD_PORT", "8080"))
+HOOK_PORT: int = int(os.getenv("AGENTD_HOOK_PORT", "9000"))

agent_sandbox/_bundled_guest/agentd/serve.py

@@ -0,0 +1,58 @@
+"""Process entrypoint: serve agentd and the lifecycle hooks side by side.
+
+One container runs one ``CMD``, so this single process starts both listeners:
+
+* agentd's application API on ``AGENT_PORT`` (main thread) — fronted by the
+  platform's auth-token proxy, this is what the SDK/CLI talk to.
+* the lifecycle-hook server on ``HOOK_PORT`` (daemon thread) — called only by the
+  platform. Bound to ``0.0.0.0`` because the platform reaches it over the guest's
+  network interface, not loopback.
+
+Running both in one process lets the hooks reseed state agentd shares (see
+``identity``/``state``); running the hooks on their own thread keeps ``/suspend``
+responsive even while agentd is busy.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import threading
+
+import uvicorn
+
+from agentd.ports import AGENT_PORT, HOOK_PORT
+
+log = logging.getLogger("agentd.serve")
+
+
+def _serve_hooks() -> None:
+    # The hook server is required: the platform's /ready gates image build, and
+    # /suspend etc. gate runtime transitions. If it ever exits — failed bind,
+    # import error — don't let agentd keep the container alive without it (which
+    # would surface only as an opaque /ready timeout). Crash the process instead.
+    try:
+        uvicorn.run(
+            "agentd.hooks:hooks_app",
+            host="0.0.0.0",
+            port=HOOK_PORT,
+            log_level="warning",
+        )
+    except BaseException:
+        log.exception("hook server crashed")
+    finally:
+        os._exit(1)
+
+
+def main() -> None:
+    threading.Thread(target=_serve_hooks, name="agentd-hooks", daemon=True).start()
+    uvicorn.run(
+        "agentd.app:app",
+        host="0.0.0.0",
+        port=AGENT_PORT,
+        log_level="info",
+    )
+
+
+if __name__ == "__main__":
+    main()

agent_sandbox/_bundled_guest/agentd/state.py

@@ -0,0 +1,58 @@
+"""Shared in-process state, visible to both server threads.
+
+agentd (app plane) and the lifecycle hooks (:mod:`agentd.hooks`) run in the same
+process but on separate threads/event loops (see ``serve.py``), so this module
+uses plain :mod:`threading` primitives rather than asyncio ones:
+
+* :data:`ready` — set once agentd's app has started; the ``/ready`` hook reads it
+  so the platform snapshots only a fully-booted agentd.
+* :data:`inflight` — count of currently-running ``exec`` calls, so the
+  ``/suspend`` hook can drain in-flight work before the VM is frozen.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Iterator
+from contextlib import contextmanager
+
+# Set by agentd's startup event; read by the /ready hook.
+ready = threading.Event()
+
+_lock = threading.Lock()
+_inflight = 0
+_drained = threading.Condition(_lock)
+
+
+@contextmanager
+def track() -> Iterator[None]:
+    """Count one in-flight exec for the duration of the ``with`` block."""
+    global _inflight
+    with _lock:
+        _inflight += 1
+    try:
+        yield
+    finally:
+        with _lock:
+            _inflight -= 1
+            if _inflight == 0:
+                _drained.notify_all()
+
+
+def inflight() -> int:
+    with _lock:
+        return _inflight
+
+
+def wait_drain(timeout: float) -> bool:
+    """Block until no exec is in flight, or ``timeout`` elapses.
+
+    Returns ``True`` if fully drained, ``False`` on timeout. Used by ``/suspend``
+    so the snapshot isn't taken mid-command.
+    """
+    with _drained:
+        if _inflight == 0:
+            return True
+        # wait_for returns True when the predicate held (drained), False on timeout.
+        drained = _drained.wait_for(lambda: _inflight == 0, timeout=timeout)
+    return drained

agent_sandbox/_bundled_guest/requirements.txt

@@ -0,0 +1,3 @@
+fastapi>=0.111
+uvicorn[standard]>=0.30
+pydantic>=2.7