forked from wrenn/python-sdk
Introduces the core Wrenn client and a dedicated sandbox execution environment. This includes automated model generation and a custom exception hierarchy to support robust integration. - Add `WrennClient` in `src/wrenn/client.py` for API interaction. - Implement `Sandbox` in `src/wrenn/sandbox.py` for isolated execution. - Add Pydantic/model support via `_generated.py`. - Define project-specific error types in `exceptions.py`. - Include AGENTS.md documentation for specialized logic. - Add comprehensive unit and integration tests. - Update build system (Makefile, uv.lock, pyproject.toml) and LICENSE.
253 lines
12 KiB
Markdown
253 lines
12 KiB
Markdown
# AGENTS.md
|
|
|
|
This file provides strict guidance to AI coding agents and assistants when modifying code in the `wrenn-python-sdk` repository. Read this entirely before writing or refactoring any code.
|
|
|
|
## Project Overview
|
|
|
|
This is the official Python SDK for **Wrenn**, a microVM-based code execution platform. The SDK provides developers and AI agents with a clean, typed interface to interact with the Wrenn Control Plane over REST and WebSockets.
|
|
|
|
**Important:** The SDK communicates exclusively with the Control Plane over HTTP/HTTPS and WebSockets. It does **not** generate or use gRPC stubs. The `envd` guest agent and `HostAgentService` are internal RPCs between the control plane and host agents — they are never reachable from the SDK. All data-plane operations (exec, file I/O) are proxied through the control plane's REST/WS endpoints.
|
|
|
|
## Repository Architecture & Structure
|
|
|
|
This is a modern Python package managed entirely by `uv`. It uses a flattened `src/` layout.
|
|
|
|
```text
|
|
.
|
|
├── LICENSE
|
|
├── Makefile # Central command runner
|
|
├── pyproject.toml # uv dependency and build config
|
|
├── uv.lock # Exact dependency resolution
|
|
├── internal/
|
|
│ └── api/
|
|
│ └── openapi.yaml # Cached OpenAPI spec from the Go backend
|
|
├── src/
|
|
│ └── wrenn/ # The actual importable Python package
|
|
│ ├── __init__.py # Version + top-level re-exports
|
|
│ ├── client.py # WrennClient & AsyncWrennClient (httpx transport)
|
|
│ ├── sandbox.py # Sandbox class (exec, files, context manager)
|
|
│ ├── exceptions.py # Typed exception hierarchy
|
|
│ ├── py.typed # PEP 561 marker
|
|
│ └── models/
|
|
│ ├── __init__.py # Public re-exports via __all__
|
|
│ └── _generated.py # DO NOT EDIT — generated by datamodel-codegen
|
|
└── tests/ # Pytest suite
|
|
```
|
|
|
|
## Build & Development Commands
|
|
|
|
Never use raw `pip`, `venv`, or `python -m venv`. **All dependency management and script execution goes through `uv` and the `Makefile`.**
|
|
|
|
```bash
|
|
make generate # Fetches openapi.yaml and runs datamodel-codegen → models/_generated.py
|
|
make lint # Runs ruff check and ruff format
|
|
make test # Runs pytest
|
|
make check # Runs lint + test
|
|
```
|
|
|
|
There is no `make proto`. The SDK does not generate gRPC stubs — the `envd` and `HostAgentService` protos are internal to the Go backend.
|
|
|
|
## Dependency Management (`uv`)
|
|
|
|
- **Adding a runtime dependency:** `uv add <package>` (e.g., `uv add httpx pydantic`)
|
|
- **Adding a dev dependency:** `uv add --dev <package>` (e.g., `uv add --dev pytest ruff`)
|
|
- **Running isolated scripts:** Use `uv run <command>`. `uv` implicitly manages the `.venv`; do not try to manually activate it in automation scripts.
|
|
|
|
## Code Generation Invariants (CRITICAL)
|
|
|
|
The data models for this SDK are generated directly from the Go backend's OpenAPI contract (`internal/api/openapi.yaml`).
|
|
|
|
1. **Never manually edit `src/wrenn/models/_generated.py`.** Any custom logic placed here will be destroyed on the next `make generate`.
|
|
2. If the Go API contract changes, run `make generate`.
|
|
3. **Export routing:** The `_generated.py` file is large. Users must never import from it directly. All user-facing models must be explicitly re-exported in `src/wrenn/models/__init__.py` using the `__all__` dunder list.
|
|
4. **Extending models:** If a generated Pydantic model needs custom Python methods, subclass it in a new file (e.g., `src/wrenn/sandbox.py` extends the generated `Sandbox` model) and export the subclass.
|
|
|
|
## Authentication
|
|
|
|
The SDK supports two authentication mechanisms, set via the `WrennClient` constructor:
|
|
|
|
1. **API Key (primary):** Pass `api_key="wrn_..."` to the constructor. Sent as `X-API-Key` header. Format: `wrn_` + 32 hex chars. Used for programmatic/agent access.
|
|
2. **JWT (secondary):** Pass `token="<jwt>"` to the constructor. Sent as `Authorization: Bearer <jwt>` header. Used for user-facing tooling. Tokens expire after 6 hours.
|
|
|
|
Host tokens (`X-Host-Token`) are for the host agent binary only and are **not** exposed in the SDK.
|
|
|
|
```python
|
|
client = WrennClient(api_key="wrn_ab12cd34...") # typical usage
|
|
client = WrennClient(token="eyJhbGci...") # alternative
|
|
```
|
|
|
|
## Core SDK Design Patterns
|
|
|
|
### 1. Sync and Async Parity
|
|
|
|
The SDK must natively support both synchronous and asynchronous workflows.
|
|
- Core logic lives in `WrennClient` and `AsyncWrennClient` inside `client.py`.
|
|
- Under the hood, rely on `httpx.Client` and `httpx.AsyncClient`.
|
|
- Resource namespaces are injected via constructor.
|
|
|
|
### 2. Resource Namespaces
|
|
|
|
The client exposes resources as plural namespaces matching the API path convention:
|
|
|
|
```python
|
|
client = WrennClient(api_key="wrn_...")
|
|
client.sandboxes.create(template="base-python")
|
|
client.sandboxes.list()
|
|
client.snapshots.create(sandbox_id="cl-...")
|
|
client.api_keys.create(name="my-key")
|
|
client.hosts.list()
|
|
client.teams.list()
|
|
client.audit.list(limit=50)
|
|
client.builds.list() # admin-only
|
|
```
|
|
|
|
### 3. The Sandbox Class
|
|
|
|
The `Sandbox` object is the primary developer-facing interface. It wraps the generated `Sandbox` model with lifecycle and data-plane methods:
|
|
|
|
```python
|
|
with client.sandboxes.create("base-python") as sb:
|
|
sb.wait_ready(timeout=30)
|
|
|
|
result = sb.exec("echo hello")
|
|
print(result.stdout) # "hello\n"
|
|
print(result.exit_code) # 0
|
|
|
|
sb.upload("/app/main.py", b"print('hello')")
|
|
data = sb.download("/app/main.py")
|
|
|
|
sb.ping()
|
|
sb.pause()
|
|
sb.resume()
|
|
# Exiting the block automatically calls sb.destroy()
|
|
```
|
|
|
|
**Key methods:**
|
|
|
|
| Method | Endpoint | Description |
|
|
|--------|----------|-------------|
|
|
| `sb.exec(cmd)` | `POST /v1/sandboxes/{id}/exec` | Synchronous exec. Returns `ExecResult` with `stdout`, `stderr`, `exit_code`, `duration_ms`. |
|
|
| `sb.exec_stream(cmd)` | `WS GET /v1/sandboxes/{id}/exec/stream` | Streaming exec via WebSocket. Returns an `Iterator[StreamEvent]` yielding `start`, `stdout`, `stderr`, `exit`, `error` events. |
|
|
| `sb.upload(path, data)` | `POST /v1/sandboxes/{id}/files/write` | Upload a small file (multipart form-data). |
|
|
| `sb.download(path)` | `POST /v1/sandboxes/{id}/files/read` | Download a small file. Returns bytes. |
|
|
| `sb.stream_upload(path, stream)` | `POST /v1/sandboxes/{id}/files/stream/write` | Streaming multipart upload for large files. No in-memory buffering. |
|
|
| `sb.stream_download(path)` | `POST /v1/sandboxes/{id}/files/stream/read` | Streaming chunked download for large files. Returns `Iterator[bytes]`. |
|
|
| `sb.wait_ready(timeout=30)` | Polls `GET /v1/sandboxes/{id}` | Blocks until status is `running`. Raises `TimeoutError` on expiry. |
|
|
| `sb.ping()` | `POST /v1/sandboxes/{id}/ping` | Resets inactivity timer. |
|
|
| `sb.pause()` | `POST /v1/sandboxes/{id}/pause` | Snapshots and releases resources. |
|
|
| `sb.resume()` | `POST /v1/sandboxes/{id}/resume` | Restores from snapshot. |
|
|
| `sb.destroy()` | `DELETE /v1/sandboxes/{id}` | Tears down the sandbox. Called automatically by context manager. |
|
|
| `sb.metrics(range="10m")` | `GET /v1/sandboxes/{id}/metrics` | Returns CPU, memory, disk time-series. |
|
|
| `sb.run_code(code, language="python")` | Jupyter kernel via proxy WS | Stateful code execution in any language with a Jupyter kernel. Variables persist across calls. Returns `CodeResult` with `.text`, `.stdout`, `.stderr`, `.error`, `.data`. See `CODE_EXECUTION.md`. |
|
|
|
|
### 4. Context Managers
|
|
|
|
Sandboxes are ephemeral. The SDK must use context managers (`with` and `async with`) to guarantee cleanup:
|
|
|
|
```python
|
|
with client.sandboxes.create("base-python") as sb:
|
|
sb.wait_ready(timeout=30)
|
|
result = sb.exec("python -c 'print(42)'")
|
|
# __exit__ calls sb.destroy() / DELETE /v1/sandboxes/{id}
|
|
```
|
|
|
|
### 5. Streaming Executions
|
|
|
|
There are two distinct exec endpoints:
|
|
|
|
**Synchronous exec** — `sb.exec(cmd, args=[], timeout_sec=30)`
|
|
- Calls `POST /v1/sandboxes/{id}/exec`. Blocks until the command completes.
|
|
- Returns an `ExecResult` with `stdout`, `stderr`, `exit_code`, `duration_ms`, `encoding`.
|
|
|
|
**Streaming exec** — `sb.exec_stream(cmd, args=[])`
|
|
- Opens a WebSocket to `GET /v1/sandboxes/{id}/exec/stream`.
|
|
- Returns an `Iterator[StreamEvent]` (or `AsyncIterator[StreamEvent]` for async).
|
|
- The client sends `{"type": "start", "cmd": "...", "args": [...]}` as the first message.
|
|
- The server sends events: `StreamStartEvent(pid)`, `StreamStdoutEvent(data)`, `StreamStderrEvent(data)`, `StreamExitEvent(exit_code)`, `StreamErrorEvent(data)`.
|
|
- The connection closes after the process exits. The client can send `{"type": "stop"}` to terminate early.
|
|
|
|
### 6. Error Handling
|
|
|
|
Do not leak raw `httpx.HTTPStatusError` to the user. The server returns errors as:
|
|
|
|
```json
|
|
{"error": {"code": "not_found", "message": "sandbox not found"}}
|
|
```
|
|
|
|
Map the `code` field (not just HTTP status) to typed exceptions:
|
|
|
|
| Error code | HTTP status | Exception |
|
|
|-----------|-------------|-----------|
|
|
| `invalid_request` | 400 | `WrennValidationError` |
|
|
| `unauthorized` | 401 | `WrennAuthenticationError` |
|
|
| `forbidden` | 403 | `WrennForbiddenError` |
|
|
| `not_found` | 404 | `WrennNotFoundError` |
|
|
| `invalid_state` | 409 | `WrennConflictError` |
|
|
| `conflict` | 409 | `WrennConflictError` |
|
|
| `host_has_sandboxes` | 409 | `WrennHostHasSandboxesError` (includes `sandbox_ids`) |
|
|
| `host_unavailable` | 503 | `WrennHostUnavailableError` |
|
|
| `agent_error` | 502 | `WrennAgentError` |
|
|
| `internal_error` | 500 | `WrennInternalError` |
|
|
|
|
All exceptions inherit from `WrennError` and expose `.code`, `.message`, and `.status_code`.
|
|
|
|
### 7. Resource Coverage
|
|
|
|
The full API surface exposed through resource namespaces:
|
|
|
|
**`client.sandboxes`** — `create`, `list`, `get`, `destroy`, `get_stats`
|
|
**`client.snapshots`** — `create`, `list`, `delete`
|
|
**`client.api_keys`** — `create`, `list`, `delete`
|
|
**`client.hosts`** — `create`, `list`, `get`, `delete`, `delete_preview`, `regenerate_token`, `list_tags`, `add_tag`, `remove_tag`
|
|
**`client.teams`** — `list`, `create`, `get`, `rename`, `delete`, `list_members`, `add_member`, `update_member_role`, `remove_member`, `leave`
|
|
**`client.audit`** — `list` (paginated with `before`/`before_id` cursors)
|
|
**`client.builds`** — `create`, `list`, `get`, `cancel` (admin-only)
|
|
**`client.admin`** — `set_team_byoc`, `list_templates`, `delete_template`
|
|
|
|
### 8. Sandbox Proxy / Port Forwarding
|
|
|
|
Services running inside a sandbox are accessible via a reverse proxy. The control plane intercepts requests whose `Host` header matches `{port}-{sandbox_id}.{domain}` and forwards them to the host agent.
|
|
|
|
The SDK exposes two helpers on the `Sandbox` object:
|
|
|
|
**`sb.get_url(port) -> str`**
|
|
- Constructs the proxy URL from the client's `base_url`.
|
|
- Derivation: parse `base_url` host, build `http://{port}-{sandbox_id}.{host}`.
|
|
- Example: `base_url="https://api.wrenn.dev"`, `sb.id="cl-abc123"` → `"http://8888-cl-abc123.api.wrenn.dev"`
|
|
- Example: `base_url="http://localhost:8080"`, `sb.id="cl-abc123"` → `"http://8888-cl-abc123.localhost:8080"`
|
|
|
|
**`sb.http_client -> httpx.Client`**
|
|
- A pre-configured `httpx.Client` with:
|
|
- `base_url` set to the proxy URL (root `/` maps to the proxied service)
|
|
- `X-API-Key` header set from the parent client's API key
|
|
- Allows direct HTTP interaction with services inside the sandbox without manual header management.
|
|
- Closed automatically when the sandbox context manager exits.
|
|
|
|
**Auth:** Proxy requests require the `X-API-Key` header. JWT is not supported for proxy routes. If the client was constructed with a JWT token only, `sb.get_url()` and `sb.http_client` must raise `WrennAuthenticationError`.
|
|
|
|
**Example: Jupyter inside a sandbox**
|
|
|
|
```python
|
|
with client.sandboxes.create("python-jupyter") as sb:
|
|
sb.wait_ready(timeout=60)
|
|
|
|
# High-level: stateful code execution (see CODE_EXECUTION.md)
|
|
result = sb.run_code("print('hello from persistent kernel')")
|
|
print(result.stdout)
|
|
|
|
# Low-level: direct HTTP to Jupyter REST API
|
|
resp = sb.http_client.get("/api/kernels")
|
|
print(resp.json())
|
|
|
|
# Low-level: direct proxy URL for browser access
|
|
jupyter_url = sb.get_url(8888)
|
|
```
|
|
|
|
## Coding Conventions & Typing
|
|
|
|
- **Python Target:** `3.13+`. Use modern syntax (`|` for Unions, standard library generics like `list[str]`).
|
|
- **Typing:** Everything must be strictly typed. Use `pyright` for validation.
|
|
- **Formatting:** `ruff` is the sole linter and formatter. Do not use `black`, `isort`, or `flake8`.
|
|
- **Docstrings:** Use Google-style docstrings. These surface to end-users via IDE hover.
|
|
- **No comments:** Do not add comments unless explicitly asked.
|