wrenn/python-sdk
2
0
Fork 1
Code Issues Pull Requests Projects Wiki Activity

v0.1.1 #7

Merged
pptx704 merged 29 commits from dev into main 2026-05-01 23:06:55 +00:00
Conversation 0 Commits 29 Files Changed 15 +782 -5565
29 changed files with 782 additions and 5565 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit 3d0eda5c60 - Show all commits

80
AGENTS.md
View File

@ -1,80 +0,0 @@
# AGENTS.md
## What this repo is
Python SDK for **Wrenn** (microVM code execution platform). Communicates with the Control Plane via REST + WebSockets only — no gRPC. The `envd` and `HostAgentService` are internal to the Go backend and never reachable from this SDK.
## Build & dev commands
All commands go through `uv` and the `Makefile`. Never use raw `pip`, `venv`, or `python -m venv`.
```bash
make generate # Fetch openapi.yaml → src/wrenn/models/_generated.py
make lint # ruff check + ruff format --check on src/
make test # runs ONLY tests/test_client.py
make test-integration # runs ALL tests (unit + integration, needs live server)
make check # lint + test (test_client.py only)
```
To run all unit tests (not just test_client.py):
```bash
uv run pytest tests/test_client.py tests/test_sandbox_features.py tests/test_filesystem_pty.py -v
```
To run a single test:
```bash
uv run pytest tests/test_client.py::TestAuth::test_signup -v
```
## Code generation (CRITICAL)
Models in `src/wrenn/models/_generated.py` are generated by `datamodel-codegen` from `api/openapi.yaml`.
1. **Never edit `_generated.py`** — overwritten on next `make generate`.
2. All user-facing models must be re-exported in `src/wrenn/models/__init__.py` via `__all__`.
3. To extend a generated model with custom methods, subclass it (e.g. `Sandbox` in `sandbox.py` subclasses the generated `SandboxModel`).
## Dependency management
```bash
uv add <package> # runtime dep
uv add --dev <package> # dev dep
uv run <command> # run in managed .venv
```
## Implemented resource namespaces
Only these are currently implemented in `client.py`:
- **`client.auth`** — `signup`, `login`
- **`client.api_keys`** — `create`, `list`, `delete`
- **`client.sandboxes`** — `create`, `list`, `get`, `destroy`
- **`client.snapshots`** — `create`, `list`, `delete`
- **`client.hosts`** — `create`, `list`, `get`, `delete`, `regenerate_token`, `list_tags`, `add_tag`, `remove_tag`
Both sync and async variants exist for every resource.
## Architecture notes
- **Sync/async parity**: `WrennClient` + `AsyncWrennClient` in `client.py`, using `httpx.Client`/`httpx.AsyncClient`. Async methods on `Sandbox` are prefixed `async_` (e.g. `async_exec`, `async_upload`).
- **WebSocket library**: `httpx-ws` (not `websockets`). Used for `exec_stream`, `pty`, and `run_code`.
- **Sandbox proxy URL**: `get_url(port)` returns `ws://` or `wss://` scheme. The `http_client` property converts to `http://`/`https://` automatically.
- **`Sandbox`** (in `sandbox.py`) is the main developer-facing class — subclasses generated model, adds lifecycle methods (`exec`, `upload`, `download`, `list_dir`, `mkdir`, `remove`, `pty`, `run_code`, `wait_ready`, `pause`, `resume`, `destroy`, `ping`, `metrics`), context manager support, and proxy helpers.
- **Error handling**: `handle_response()` in `exceptions.py` maps server error `code` field to typed exceptions (not just HTTP status). All inherit from `WrennError` with `.code`, `.message`, `.status_code`.
## Testing
- **HTTP mocking**: `respx` library (not `responses` or `pytest-httpx`). Mock routes with `@respx.mock` decorator or `respx.mock` context manager.
- **Async tests**: use `@pytest.mark.asyncio` (backed by `pytest-asyncio`).
- **Integration tests**: in `test_integration.py`, require env vars `WRENN_API_KEY` or `WRENN_TOKEN` (plus optional `WRENN_BASE_URL`, `WRENN_TEST_EMAIL`, `WRENN_TEST_PASSWORD`). They are skipped via `@requires_auth` if credentials are absent.
- **Fixtures**: test fixtures create `WrennClient(api_key="wrn_test1234567890abcdef12345678")` with context manager cleanup.
## Coding conventions
- **Python 3.13+** with modern syntax (`|` unions, `list[str]` generics).
- **Strict typing** throughout. `pyright`/`mypy` available but not in CI.
- **`ruff`** is the sole linter and formatter. Do not use `black`, `isort`, or `flake8`.
- **Google-style docstrings** on all public APIs.
- **No comments** unless explicitly asked.

542
README.md
View File

@ -1,6 +1,8 @@
# Wrenn Python SDK # Wrenn Python SDK
Python client for the [Wrenn](https://wrenn.dev) microVM code execution platform. Create isolated capsules, execute commands, manage files, run interactive terminals, and execute persistent code — all from Python. Python client for the [Wrenn](https://wrenn.dev) microVM platform. Create isolated capsules, execute commands, manage files, run interactive terminals, and execute persistent code -- all from Python.
Designed as a drop-in replacement for [e2b](https://e2b.dev). If you're migrating, just swap your imports.
## Installation ## Installation
@ -10,97 +12,144 @@ pip install wrenn
Requires Python 3.13+. Requires Python 3.13+.
## Quick Start
```python
from wrenn import WrennClient
client = WrennClient(api_key="wrn_your_api_key_here")
# Create a capsule and run a command
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
cap.wait_ready(timeout=60)
result = cap.exec("echo", args=["hello world"])
print(result.stdout) # "hello world"
print(result.exit_code) # 0
```
## Authentication ## Authentication
The SDK supports two authentication methods: Set the `WRENN_API_KEY` environment variable:
```python ```bash
# API key export WRENN_API_KEY="wrn_your_api_key_here"
client = WrennClient(api_key="wrn_...")
# JWT token
client = WrennClient(token="eyJ...")
``` ```
You can obtain an API key via the dashboard or create one programmatically: Optionally override the API base URL:
```python ```bash
with WrennClient(token="jwt_token") as client: export WRENN_BASE_URL="https://app.wrenn.dev/api" # default
key = client.api_keys.create(name="my-key")
print(key.key) # wrn_...
``` ```
## Capsules You can also pass credentials directly:
Capsules are isolated microVM environments. Create, manage, and interact with them:
```python ```python
# Create from wrenn import Capsule
cap = client.capsules.create(
template="base-python",
vcpus=2,
memory_mb=1024,
timeout_sec=300,
)
# List capsule = Capsule(api_key="wrn_...", base_url="https://...")
for c in client.capsules.list(): ```
print(c.id, c.status)
# Get ---
cap = client.capsules.get("cl-abc123")
# Destroy ## Wrenn Capsules
client.capsules.destroy("cl-abc123")
### Quick Start
```python
from wrenn import Capsule
# Create a capsule (reads WRENN_API_KEY from env)
with Capsule(template="minimal") as capsule:
result = capsule.commands.run("echo hello")
print(result.stdout) # "hello\n"
```
### Creating Capsules
```python
from wrenn import Capsule
# Direct construction (creates immediately)
capsule = Capsule()
capsule = Capsule(template="base-python", vcpus=2, memory_mb=1024, timeout=300)
# With auto-wait (blocks until capsule is running)
capsule = Capsule(template="minimal", wait=True)
# Via factory classmethod
capsule = Capsule.create(template="minimal", wait=True)
``` ```
### Context Manager ### Context Manager
Use capsules as context managers for automatic cleanup: Use capsules as context managers for automatic cleanup (destroys capsule on exit):
```python ```python
with client.capsules.create(template="minimal", timeout_sec=120) as cap: with Capsule(template="minimal", wait=True) as capsule:
cap.wait_ready(timeout=60) capsule.commands.run("echo hello")
cap.exec("python -c 'print(42)'") # capsule is automatically destroyed
# cap.destroy() is called automatically
``` ```
## Command Execution ### Connecting to Existing Capsules
### `exec()` — One-off Commands Attach to a running capsule by ID. If it's paused, it will be resumed automatically:
Starts a fresh process for each call. No state persists between calls.
```python ```python
result = cap.exec("python", args=["-c", "import os; print(os.getcwd())"]) capsule = Capsule.connect("cl-abc123")
print(result.stdout) # "/home/user\n" result = capsule.commands.run("echo still running")
print(result.stderr) # ""
print(result.exit_code) # 0
print(result.duration_ms) # 42
``` ```
### `exec_stream()` — Streaming Output For code interpreter capsules:
Stream real-time output from long-running commands:
```python ```python
for event in cap.exec_stream("python", args=["-u", "train.py"]): from wrenn.code_interpreter import Capsule as CodeCapsule
capsule = CodeCapsule.connect("cl-abc123")
result = capsule.run_code("print('reconnected')")
```
### Lifecycle Management
```python
# Instance methods
capsule.pause()
capsule.resume()
capsule.destroy()
capsule.ping() # reset inactivity timer
capsule.wait_ready() # block until running
info = capsule.get_info()
print(info.status) # "running"
print(capsule.is_running()) # True
# Static methods (no instance needed)
Capsule.destroy("cl-abc123", api_key="wrn_...")
Capsule.pause("cl-abc123")
Capsule.resume("cl-abc123")
info = Capsule.get_info("cl-abc123")
# List all capsules
capsules = Capsule.list()
```
### Command Execution
Commands are accessed via `capsule.commands`:
```python
# Foreground (blocks until complete)
result = capsule.commands.run("python -c 'print(42)'")
print(result.stdout) # "42\n"
print(result.stderr) # ""
print(result.exit_code) # 0
print(result.duration_ms) # 35
# With options
result = capsule.commands.run(
"python train.py",
timeout=120,
envs={"CUDA_VISIBLE_DEVICES": "0"},
cwd="/app",
)
# Background process
handle = capsule.commands.run("python server.py", background=True)
print(handle.pid) # 1234
print(handle.tag) # "exec-abc123"
```
#### Streaming Output
```python
import sys
# Stream a new command
for event in capsule.commands.stream("python", args=["-u", "train.py"]):
match event.type: match event.type:
case "stdout": case "stdout":
print(event.data, end="") print(event.data, end="")
@ -108,77 +157,80 @@ for event in cap.exec_stream("python", args=["-u", "train.py"]):
print(event.data, end="", file=sys.stderr) print(event.data, end="", file=sys.stderr)
case "exit": case "exit":
print(f"\nExited with code {event.exit_code}") print(f"\nExited with code {event.exit_code}")
# Connect to a running background process
for event in capsule.commands.connect(handle.pid):
if event.type == "stdout":
print(event.data, end="")
``` ```
### `run_code()` — Stateful Code Execution #### Process Management
Execute Python code in a persistent Jupyter kernel. Variables, imports, and function definitions survive across calls:
```python ```python
with client.capsules.create(template="python-interpreter-v0-beta") as cap: # List running processes
cap.wait_ready(timeout=60) for proc in capsule.commands.list():
print(proc.pid, proc.cmd, proc.tag)
cap.run_code("x = 42") # Kill a process
r = cap.run_code("x * 2") capsule.commands.kill(pid=1234)
print(r.text) # "84"
cap.run_code("def greet(name): return f'hello {name}'")
r = cap.run_code("greet('world')")
print(r.text) # "'hello world'"
r = cap.run_code("1/0")
print(r.error) # "ZeroDivisionError: division by zero\n..."
``` ```
**`CodeResult` fields:** ### Filesystem
| Field | Type | Description | Files are accessed via `capsule.files`:
|-------|------|-------------|
| `text` | `str \| None` | Plain text representation |
| `data` | `dict \| None` | Rich MIME bundle (e.g. `{"image/png": "..."}`) |
| `stdout` | `str` | Accumulated stdout |
| `stderr` | `str` | Accumulated stderr |
| `error` | `str \| None` | Error traceback string |
## Filesystem
Upload, download, and manage files inside capsules:
```python ```python
# Upload / Download # Write and read files
cap.upload("/app/main.py", b"print('hello')") capsule.files.write("/app/main.py", "print('hello')")
content = cap.download("/app/main.py") content = capsule.files.read("/app/main.py") # str
raw = capsule.files.read_bytes("/app/main.py") # bytes
# Streaming (for large files) # Check existence
capsule.files.exists("/app/main.py") # True
# List directory
entries = capsule.files.list("/home/user", depth=1)
for entry in entries:
print(entry.name, entry.type, entry.size)
# Create directory
capsule.files.make_dir("/app/data")
# Remove file or directory
capsule.files.remove("/app/old_data")
```
#### Streaming (Large Files)
```python
# Streaming upload
def chunks(): def chunks():
yield b"chunk1" yield b"chunk1"
yield b"chunk2" yield b"chunk2"
cap.stream_upload("/data/large.bin", chunks()) capsule.files.upload_stream("/data/large.bin", chunks())
for chunk in cap.stream_download("/data/large.bin"):
# Streaming download
for chunk in capsule.files.download_stream("/data/large.bin"):
process(chunk) process(chunk)
# Directory operations
entries = cap.list_dir("/home/user", depth=1)
for entry in entries:
print(entry.name, entry.type, entry.size)
cap.mkdir("/home/user/data")
cap.remove("/home/user/old_data")
``` ```
## Interactive Terminal (PTY) ### Interactive Terminal (PTY)
Open a full interactive terminal session over WebSocket:
```python ```python
with cap.pty(cmd="/bin/bash", cols=120, rows=40, cwd="/home/user") as term: import sys
with capsule.pty(cmd="/bin/bash", cols=120, rows=40, cwd="/home/user") as term:
term.write(b"ls -la\n") term.write(b"ls -la\n")
for event in term: for event in term:
if event.type == "output": if event.type == "output":
sys.stdout.buffer.write(event.data) sys.stdout.buffer.write(event.data)
elif event.type == "exit": elif event.type == "exit":
break break
# Reconnect to an existing session
with capsule.pty_connect(term.tag) as term:
term.write(b"echo reconnected\n")
``` ```
**PtySession methods:** **PtySession methods:**
@ -188,123 +240,169 @@ with cap.pty(cmd="/bin/bash", cols=120, rows=40, cwd="/home/user") as term:
| `write(data: bytes)` | Send raw bytes to stdin | | `write(data: bytes)` | Send raw bytes to stdin |
| `resize(cols, rows)` | Resize the terminal | | `resize(cols, rows)` | Resize the terminal |
| `kill()` | Send SIGKILL to the process | | `kill()` | Send SIGKILL to the process |
| `tag` | Session tag (available after `started` event) | | `tag` | Session tag (after `started` event) |
| `pid` | Process PID (available after `started` event) | | `pid` | Process PID (after `started` event) |
Reconnect to an existing session using the tag: ### Proxy URL
Access services running inside a capsule:
```python ```python
with cap.pty_connect(term.tag) as term: url = capsule.get_url(8080)
term.write(b"echo reconnected\n") # "wss://8080-cl-abc123.app.wrenn.dev"
``` ```
## Lifecycle ### Snapshots
Pause and resume capsules to save resources: Create reusable templates from running capsules:
```python ```python
cap = client.capsules.create(template="minimal") template = capsule.create_snapshot(name="my-template", overwrite=True)
cap.wait_ready(timeout=60)
# Pause (snapshots and releases resources)
cap.pause()
print(cap.status) # "paused"
# Resume (restores from snapshot)
cap.resume()
cap.wait_ready(timeout=60)
``` ```
Keep a capsule alive with `ping()`: ---
## Code Interpreter
The `wrenn.code_interpreter` module provides a specialized capsule for stateful code execution via a persistent Jupyter kernel.
### Quick Start
```python ```python
cap.ping() # Resets the inactivity timer from wrenn.code_interpreter import Capsule
with Capsule(wait=True) as capsule:
result = capsule.run_code("print('hello')")
print(result.text) # "hello"
``` ```
## Proxy URL ### Stateful Execution
Access services running inside a capsule through the proxy: Variables, imports, and function definitions persist across `run_code` calls:
```python ```python
url = cap.get_url(8888) from wrenn.code_interpreter import Capsule
# "wss://8888-cl-abc123.api.wrenn.dev"
# Pre-configured HTTP client targeting port 8888 with Capsule(wait=True) as capsule:
resp = cap.http_client.get("/api/kernels") capsule.run_code("x = 42")
result = capsule.run_code("x * 2")
print(result.text) # "84"
capsule.run_code("import math")
result = capsule.run_code("math.pi")
print(result.text) # "3.141592653589793"
capsule.run_code("def greet(name): return f'hello {name}'")
result = capsule.run_code("greet('world')")
print(result.text) # "hello world"
``` ```
## Snapshots The `text` field returns the expression result when available. For `print()` calls (which produce no expression result), it falls back to the stripped stdout output.
Create templates from running capsules: ### Error Handling in Code
```python ```python
# Create a snapshot result = capsule.run_code("1 / 0")
template = client.snapshots.create( print(result.error) # "ZeroDivisionError: division by zero\n..."
capsule_id="cl-abc123",
name="my-template",
overwrite=True,
)
# List templates
for t in client.snapshots.list():
print(t.name, t.type)
# Delete
client.snapshots.delete("my-template")
``` ```
## Hosts ### Rich Output
Manage host machines:
```python ```python
host = client.hosts.create(type="regular") result = capsule.run_code("""
client.hosts.list() import matplotlib.pyplot as plt
client.hosts.get("h-1") plt.plot([1, 2, 3])
client.hosts.delete("h-1") plt.savefig('/tmp/plot.png')
client.hosts.regenerate_token("h-1") plt.show()
client.hosts.list_tags("h-1") """)
client.hosts.add_tag("h-1", "gpu") print(result.data) # {"image/png": "base64...", "text/plain": "..."}
client.hosts.remove_tag("h-1", "gpu")
``` ```
### Custom Templates
By default, `code-runner-beta` template is used. You can specify a custom template:
```python
capsule = Capsule(template="my-custom-jupyter-template", wait=True)
result = capsule.run_code("print('running on custom template')")
```
### CodeResult Fields
| Field | Type | Description |
|-------|------|-------------|
| `text` | `str \| None` | Expression result, or stripped stdout if no expression result |
| `data` | `dict \| None` | Rich MIME bundle (e.g. `{"image/png": "..."}`) |
| `stdout` | `str` | Raw accumulated stdout output |
| `stderr` | `str` | Raw accumulated stderr output |
| `error` | `str \| None` | Error traceback string |
String expression results have quotes stripped automatically (e.g. `'hello'` becomes `hello`).
### Code Interpreter + Commands/Files
The code interpreter capsule inherits all standard capsule features:
```python
from wrenn.code_interpreter import Capsule
with Capsule(wait=True) as capsule:
# Use run_code for Jupyter execution
capsule.run_code("import pandas as pd; df = pd.DataFrame({'a': [1,2,3]})")
capsule.run_code("df.to_csv('/tmp/data.csv', index=False)")
# Use standard file operations
content = capsule.files.read("/tmp/data.csv")
print(content)
# Use standard command execution
result = capsule.commands.run("wc -l /tmp/data.csv")
print(result.stdout)
```
---
## Async Support ## Async Support
All operations have async variants. Use `AsyncWrennClient` and prefix capsule methods with `async_`: All operations have async variants via `AsyncCapsule`:
### Async Capsule
```python ```python
from wrenn import AsyncWrennClient from wrenn import AsyncCapsule
async with AsyncWrennClient(api_key="wrn_...") as client: async with await AsyncCapsule.create(template="minimal", wait=True) as capsule:
cap = await client.capsules.create(template="minimal") result = await capsule.commands.run("echo hello")
await cap.async_wait_ready(timeout=60) print(result.stdout)
result = await cap.async_exec("echo", args=["hello"]) await capsule.files.write("/app/file.txt", "data")
await cap.async_upload("/app/file.txt", b"data") entries = await capsule.files.list("/app")
entries = await cap.async_list_dir("/home/user")
r = await cap.async_run_code("42 * 2")
await cap.async_destroy() await capsule.pause()
await capsule.resume()
``` ```
**Async method mapping:** ### Async Code Interpreter
| Sync | Async | ```python
|------|-------| from wrenn.code_interpreter import AsyncCapsule
| `exec()` | `async_exec()` |
| `upload()` | `async_upload()` | async with await AsyncCapsule.create(wait=True) as capsule:
| `download()` | `async_download()` | result = await capsule.run_code("2 + 2")
| `stream_upload()` | `async_stream_upload()` | print(result.text) # "4"
| `stream_download()` | `async_stream_download()` | ```
| `list_dir()` | `async_list_dir()` |
| `mkdir()` | `async_mkdir()` | ### Async PTY
| `remove()` | `async_remove()` |
| `wait_ready()` | `async_wait_ready()` | ```python
| `pause()` | `async_pause()` | async with capsule.pty(cmd="/bin/bash") as term:
| `resume()` | `async_resume()` | await term.write(b"ls -la\n")
| `destroy()` | `async_destroy()` | async for event in term:
| `ping()` | `async_ping()` | if event.type == "output":
| `run_code()` | `async_run_code()` | sys.stdout.buffer.write(event.data)
```
---
## Error Handling ## Error Handling
@ -318,14 +416,14 @@ from wrenn import (
WrennForbiddenError, # 403 WrennForbiddenError, # 403
WrennNotFoundError, # 404 WrennNotFoundError, # 404
WrennConflictError, # 409 WrennConflictError, # 409
WrennHostHasCapsulesError, # 409 — host has running capsules WrennHostHasCapsulesError, # 409 (host has running capsules)
WrennAgentError, # 502 WrennAgentError, # 502
WrennInternalError, # 500 WrennInternalError, # 500
WrennHostUnavailableError, # 503 WrennHostUnavailableError, # 503
) )
try: try:
client.capsules.get("nonexistent") Capsule.get_info("nonexistent")
except WrennNotFoundError as e: except WrennNotFoundError as e:
print(e.code) # "not_found" print(e.code) # "not_found"
print(e.message) # "capsule not found" print(e.message) # "capsule not found"
@ -334,6 +432,67 @@ except WrennNotFoundError as e:
All exceptions inherit from `WrennError` and expose `.code`, `.message`, and `.status_code`. All exceptions inherit from `WrennError` and expose `.code`, `.message`, and `.status_code`.
---
## Migrating from e2b
Replace your imports:
```python
# Before
from e2b import Sandbox
sandbox = Sandbox()
# After
from wrenn import Capsule
capsule = Capsule()
```
For code interpreter:
```python
# Before
from e2b_code_interpreter import Sandbox
sandbox = Sandbox()
result = sandbox.run_code("print('hello')")
# After
from wrenn.code_interpreter import Capsule
capsule = Capsule()
result = capsule.run_code("print('hello')")
```
The `Sandbox` name is available as a deprecated alias in both modules:
```python
from wrenn import Sandbox # works, emits FutureWarning
from wrenn.code_interpreter import Sandbox # works, emits FutureWarning
```
---
## Low-Level Client
For direct API access, use `WrennClient` / `AsyncWrennClient`:
```python
from wrenn import WrennClient
with WrennClient(api_key="wrn_...") as client:
capsule = client.capsules.create(template="minimal")
client.capsules.pause(capsule.id)
client.capsules.resume(capsule.id)
client.capsules.ping(capsule.id)
client.capsules.destroy(capsule.id)
# Snapshots
template = client.snapshots.create(capsule_id="cl-abc", name="my-snap")
templates = client.snapshots.list()
client.snapshots.delete("my-snap")
```
---
## Development ## Development
This project uses [uv](https://docs.astral.sh/uv/) for dependency management. This project uses [uv](https://docs.astral.sh/uv/) for dependency management.
@ -350,14 +509,11 @@ make test
# Run all tests (including integration) # Run all tests (including integration)
make test-integration make test-integration
# Regenerate models from OpenAPI spec
make generate
``` ```
### Running Integration Tests ### Running Integration Tests
Integration tests require a live Wrenn server. Set environment variables: Integration tests require a live Wrenn server:
```bash ```bash
export WRENN_API_KEY="wrn_..." export WRENN_API_KEY="wrn_..."

14
src/wrenn/async_capsule.py
View File

@ -63,6 +63,7 @@ class AsyncCapsule:
memory_mb: int | None = None, memory_mb: int | None = None,
timeout: int | None = None, timeout: int | None = None,
*, *,
wait: bool = False,
api_key: str | None = None, api_key: str | None = None,
base_url: str | None = None, base_url: str | None = None,
) -> AsyncCapsule: ) -> AsyncCapsule:
@ -74,11 +75,14 @@ class AsyncCapsule:
memory_mb=memory_mb, memory_mb=memory_mb,
timeout_sec=timeout, timeout_sec=timeout,
) )
return cls( capsule = cls(
_capsule_id=info.id, _capsule_id=info.id,
_client=client, _client=client,
_info=info, _info=info,
) )
if wait:
await capsule.wait_ready()
return capsule
@classmethod @classmethod
async def connect( async def connect(
@ -103,16 +107,16 @@ class AsyncCapsule:
# ── Dual instance/static lifecycle ────────────────────────── # ── Dual instance/static lifecycle ──────────────────────────
kill = _DualMethod("_instance_kill", "_static_kill") destroy = _DualMethod("_instance_destroy", "_static_destroy")
pause = _DualMethod("_instance_pause", "_static_pause") pause = _DualMethod("_instance_pause", "_static_pause")
resume = _DualMethod("_instance_resume", "_static_resume") resume = _DualMethod("_instance_resume", "_static_resume")
get_info = _DualMethod("_instance_get_info", "_static_get_info") get_info = _DualMethod("_instance_get_info", "_static_get_info")
async def _instance_kill(self) -> None: async def _instance_destroy(self) -> None:
await self._client.capsules.destroy(self._id) await self._client.capsules.destroy(self._id)
@classmethod @classmethod
async def _static_kill( async def _static_destroy(
cls, cls,
capsule_id: str, capsule_id: str,
*, *,
@ -260,7 +264,7 @@ class AsyncCapsule:
exc_tb: object, exc_tb: object,
) -> None: ) -> None:
try: try:
await self._instance_kill() await self._instance_destroy()
except Exception: except Exception:
pass pass
try: try:

14
src/wrenn/capsule.py
View File

@ -66,6 +66,7 @@ class Capsule:
memory_mb: int | None = None, memory_mb: int | None = None,
timeout: int | None = None, timeout: int | None = None,
*, *,
wait: bool = False,
api_key: str | None = None, api_key: str | None = None,
base_url: str | None = None, base_url: str | None = None,
# Private: used by classmethods to skip creation # Private: used by classmethods to skip creation
@ -93,6 +94,9 @@ class Capsule:
self.commands = Commands(self._id, self._client.http) self.commands = Commands(self._id, self._client.http)
self.files = Files(self._id, self._client.http) self.files = Files(self._id, self._client.http)
if wait:
self.wait_ready()
# ── Properties ────────────────────────────────────────────── # ── Properties ──────────────────────────────────────────────
@property @property
@ -113,6 +117,7 @@ class Capsule:
memory_mb: int | None = None, memory_mb: int | None = None,
timeout: int | None = None, timeout: int | None = None,
*, *,
wait: bool = False,
api_key: str | None = None, api_key: str | None = None,
base_url: str | None = None, base_url: str | None = None,
) -> Capsule: ) -> Capsule:
@ -122,6 +127,7 @@ class Capsule:
vcpus=vcpus, vcpus=vcpus,
memory_mb=memory_mb, memory_mb=memory_mb,
timeout=timeout, timeout=timeout,
wait=wait,
api_key=api_key, api_key=api_key,
base_url=base_url, base_url=base_url,
) )
@ -149,17 +155,17 @@ class Capsule:
# ── Dual instance/static lifecycle ────────────────────────── # ── Dual instance/static lifecycle ──────────────────────────
kill = _DualMethod("_instance_kill", "_static_kill") destroy = _DualMethod("_instance_destroy", "_static_destroy")
pause = _DualMethod("_instance_pause", "_static_pause") pause = _DualMethod("_instance_pause", "_static_pause")
resume = _DualMethod("_instance_resume", "_static_resume") resume = _DualMethod("_instance_resume", "_static_resume")
get_info = _DualMethod("_instance_get_info", "_static_get_info") get_info = _DualMethod("_instance_get_info", "_static_get_info")
def _instance_kill(self) -> None: def _instance_destroy(self) -> None:
"""Destroy this capsule.""" """Destroy this capsule."""
self._client.capsules.destroy(self._id) self._client.capsules.destroy(self._id)
@classmethod @classmethod
def _static_kill( def _static_destroy(
cls, cls,
capsule_id: str, capsule_id: str,
*, *,
@ -321,7 +327,7 @@ class Capsule:
exc_tb: object, exc_tb: object,
) -> None: ) -> None:
try: try:
self._instance_kill() self._instance_destroy()
except Exception: except Exception:
pass pass
try: try:

20
src/wrenn/code_interpreter/__init__.py
View File

@ -1,8 +1,26 @@
from wrenn.code_interpreter.capsule import Capsule, CodeResult
from wrenn.code_interpreter.async_capsule import AsyncCapsule from wrenn.code_interpreter.async_capsule import AsyncCapsule
from wrenn.code_interpreter.capsule import Capsule, CodeResult
__all__ = [ __all__ = [
"AsyncCapsule", "AsyncCapsule",
"Capsule", "Capsule",
"CodeResult", "CodeResult",
"Sandbox",
] ]
def __getattr__(name: str) -> type:
import sys
import warnings
_module = sys.modules[__name__]
if name == "Sandbox":
warnings.warn(
"'Sandbox' is deprecated, use 'Capsule' instead",
FutureWarning,
stacklevel=2,
)
setattr(_module, name, Capsule)
return Capsule
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

32
src/wrenn/code_interpreter/async_capsule.py
View File

@ -41,6 +41,7 @@ class AsyncCapsule(BaseAsyncCapsule):
memory_mb: int | None = None, memory_mb: int | None = None,
timeout: int | None = None, timeout: int | None = None,
*, *,
wait: bool = False,
api_key: str | None = None, api_key: str | None = None,
base_url: str | None = None, base_url: str | None = None,
) -> AsyncCapsule: ) -> AsyncCapsule:
@ -51,11 +52,14 @@ class AsyncCapsule(BaseAsyncCapsule):
memory_mb=memory_mb, memory_mb=memory_mb,
timeout_sec=timeout, timeout_sec=timeout,
) )
return cls( capsule = cls(
_capsule_id=info.id, _capsule_id=info.id,
_client=client, _client=client,
_info=info, _info=info,
) )
if wait:
await capsule.wait_ready()
return capsule
def _get_proxy_client(self) -> httpx.AsyncClient: def _get_proxy_client(self) -> httpx.AsyncClient:
if self._proxy_client is None: if self._proxy_client is None:
@ -80,11 +84,20 @@ class AsyncCapsule(BaseAsyncCapsule):
while time.monotonic() < deadline: while time.monotonic() < deadline:
try: try:
resp = await client.post("/api/kernels") # Try to reuse an existing kernel
resp = await client.get("/api/kernels")
if resp.status_code < 500: if resp.status_code < 500:
resp.raise_for_status() resp.raise_for_status()
self._kernel_id = resp.json()["id"] kernels = resp.json()
return self._kernel_id if kernels:
self._kernel_id = kernels[0]["id"]
return self._kernel_id
# No existing kernels, create a new one
resp = await client.post("/api/kernels")
if resp.status_code < 500:
resp.raise_for_status()
self._kernel_id = resp.json()["id"]
return self._kernel_id
last_exc = httpx.HTTPStatusError( last_exc = httpx.HTTPStatusError(
f"Jupyter returned {resp.status_code}", f"Jupyter returned {resp.status_code}",
request=resp.request, request=resp.request,
@ -180,7 +193,13 @@ class AsyncCapsule(BaseAsyncCapsule):
result.stdout += content.get("text", "") result.stdout += content.get("text", "")
elif msg_type == "execute_result": elif msg_type == "execute_result":
bundle = content.get("data", {}) bundle = content.get("data", {})
result.text = bundle.get("text/plain") text = bundle.get("text/plain")
if text and (
(text.startswith("'") and text.endswith("'"))
or (text.startswith('"') and text.endswith('"'))
):
text = text[1:-1]
result.text = text
result.data = bundle result.data = bundle
elif msg_type == "error": elif msg_type == "error":
traceback = content.get("traceback", []) traceback = content.get("traceback", [])
@ -188,6 +207,9 @@ class AsyncCapsule(BaseAsyncCapsule):
elif msg_type == "status" and content.get("execution_state") == "idle": elif msg_type == "status" and content.get("execution_state") == "idle":
break break
if result.text is None and result.stdout:
result.text = result.stdout.strip()
return result return result
async def __aexit__(self, *args) -> None: async def __aexit__(self, *args) -> None:

28
src/wrenn/code_interpreter/capsule.py
View File

@ -80,6 +80,7 @@ class Capsule(BaseCapsule):
memory_mb: int | None = None, memory_mb: int | None = None,
timeout: int | None = None, timeout: int | None = None,
*, *,
wait: bool = False,
api_key: str | None = None, api_key: str | None = None,
base_url: str | None = None, base_url: str | None = None,
) -> Capsule: ) -> Capsule:
@ -88,6 +89,7 @@ class Capsule(BaseCapsule):
vcpus=vcpus, vcpus=vcpus,
memory_mb=memory_mb, memory_mb=memory_mb,
timeout=timeout, timeout=timeout,
wait=wait,
api_key=api_key, api_key=api_key,
base_url=base_url, base_url=base_url,
) )
@ -115,11 +117,20 @@ class Capsule(BaseCapsule):
while time.monotonic() < deadline: while time.monotonic() < deadline:
try: try:
resp = client.post("/api/kernels") # Try to reuse an existing kernel
resp = client.get("/api/kernels")
if resp.status_code < 500: if resp.status_code < 500:
resp.raise_for_status() resp.raise_for_status()
self._kernel_id = resp.json()["id"] kernels = resp.json()
return self._kernel_id if kernels:
self._kernel_id = kernels[0]["id"]
return self._kernel_id
# No existing kernels, create a new one
resp = client.post("/api/kernels")
if resp.status_code < 500:
resp.raise_for_status()
self._kernel_id = resp.json()["id"]
return self._kernel_id
last_exc = httpx.HTTPStatusError( last_exc = httpx.HTTPStatusError(
f"Jupyter returned {resp.status_code}", f"Jupyter returned {resp.status_code}",
request=resp.request, request=resp.request,
@ -225,7 +236,13 @@ class Capsule(BaseCapsule):
result.stdout += content.get("text", "") result.stdout += content.get("text", "")
elif msg_type == "execute_result": elif msg_type == "execute_result":
bundle = content.get("data", {}) bundle = content.get("data", {})
result.text = bundle.get("text/plain") text = bundle.get("text/plain")
if text and (
(text.startswith("'") and text.endswith("'"))
or (text.startswith('"') and text.endswith('"'))
):
text = text[1:-1]
result.text = text
result.data = bundle result.data = bundle
elif msg_type == "error": elif msg_type == "error":
traceback = content.get("traceback", []) traceback = content.get("traceback", [])
@ -233,6 +250,9 @@ class Capsule(BaseCapsule):
elif msg_type == "status" and content.get("execution_state") == "idle": elif msg_type == "status" and content.get("execution_state") == "idle":
break break
if result.text is None and result.stdout:
result.text = result.stdout.strip()
return result return result
def __exit__(self, *args) -> None: def __exit__(self, *args) -> None:

4
tests/test_capsule_features.py
View File

@ -68,9 +68,9 @@ class TestCapsuleCreate:
class TestCapsuleStaticMethods: class TestCapsuleStaticMethods:
@respx.mock @respx.mock
def test_static_kill(self): def test_static_destroy(self):
route = respx.delete(f"{BASE}/v1/capsules/cl-1").respond(204) route = respx.delete(f"{BASE}/v1/capsules/cl-1").respond(204)
Capsule._static_kill("cl-1", api_key="wrn_test1234567890abcdef12345678") Capsule._static_destroy("cl-1", api_key="wrn_test1234567890abcdef12345678")
assert route.called assert route.called
@respx.mock @respx.mock