feat: implement client architecture and sandbox environment

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.
2026-04-10 22:24:50 +06:00
parent 4e02061bc9
commit f51a962fff
15 changed files with 3099 additions and 14 deletions
--- a/AGENTS.md
+++ b/AGENTS.md
@ -0,0 +1,252 @@
+# 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.