v0.1.5 #13
1
.gitignore
vendored
1
.gitignore
vendored
@ -174,3 +174,4 @@ cython_debug/
|
|||||||
# PyPI configuration file
|
# PyPI configuration file
|
||||||
.pypirc
|
.pypirc
|
||||||
|
|
||||||
|
CODE_EXECUTION.md
|
||||||
|
|||||||
46
.woodpecker/check.yml
Normal file
46
.woodpecker/check.yml
Normal file
@ -0,0 +1,46 @@
|
|||||||
|
when:
|
||||||
|
event: push
|
||||||
|
branch:
|
||||||
|
- main
|
||||||
|
- dev
|
||||||
|
|
||||||
|
variables:
|
||||||
|
- &python_image "ghcr.io/astral-sh/uv:python3.13-bookworm-slim"
|
||||||
|
- &uv_cache_dir "/root/.cache/uv"
|
||||||
|
|
||||||
|
steps:
|
||||||
|
- name: restore-cache
|
||||||
|
image: woodpeckerci/plugin-cache
|
||||||
|
settings:
|
||||||
|
restore: true
|
||||||
|
cache_key: "uv-{{ checksum \"uv.lock\" }}"
|
||||||
|
mount:
|
||||||
|
- /root/.cache/uv
|
||||||
|
|
||||||
|
- name: lint
|
||||||
|
image: *python_image
|
||||||
|
environment:
|
||||||
|
UV_CACHE_DIR: *uv_cache_dir
|
||||||
|
UV_FROZEN: 1
|
||||||
|
commands:
|
||||||
|
- uv sync --no-install-project
|
||||||
|
- make lint
|
||||||
|
|
||||||
|
- name: test
|
||||||
|
image: *python_image
|
||||||
|
environment:
|
||||||
|
UV_CACHE_DIR: *uv_cache_dir
|
||||||
|
UV_FROZEN: 1
|
||||||
|
commands:
|
||||||
|
- uv sync --no-install-project
|
||||||
|
- make test
|
||||||
|
|
||||||
|
- name: rebuild-cache
|
||||||
|
image: woodpeckerci/plugin-cache
|
||||||
|
when:
|
||||||
|
- status: [success]
|
||||||
|
settings:
|
||||||
|
rebuild: true
|
||||||
|
cache_key: "uv-{{ checksum \"uv.lock\" }}"
|
||||||
|
mount:
|
||||||
|
- /root/.cache/uv
|
||||||
20
LICENSE
20
LICENSE
@ -1,18 +1,18 @@
|
|||||||
MIT License
|
MIT License
|
||||||
|
|
||||||
Copyright (c) 2026 wrenn
|
Copyright (c) 2026 M/S Omukk, Bangladesh
|
||||||
|
|
||||||
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
|
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
|
||||||
associated documentation files (the "Software"), to deal in the Software without restriction, including
|
associated documentation files (the "Software"), to deal in the Software without restriction, including
|
||||||
without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||||
copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the
|
copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the
|
||||||
following conditions:
|
following conditions:
|
||||||
|
|
||||||
The above copyright notice and this permission notice shall be included in all copies or substantial
|
The above copyright notice and this permission notice shall be included in all copies or substantial
|
||||||
portions of the Software.
|
portions of the Software.
|
||||||
|
|
||||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
|
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
|
||||||
LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO
|
LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO
|
||||||
EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
|
EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
|
||||||
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
|
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
|
||||||
USE OR OTHER DEALINGS IN THE SOFTWARE.
|
USE OR OTHER DEALINGS IN THE SOFTWARE.
|
||||||
|
|||||||
20
Makefile
20
Makefile
@ -1,8 +1,8 @@
|
|||||||
# Makefile
|
# Makefile
|
||||||
.PHONY: generate
|
.PHONY: generate lint test check test-integration
|
||||||
|
|
||||||
# Variables
|
# Variables
|
||||||
SPEC_URL = "https://git.omukk.dev/wrenn/sandbox/raw/branch/main/internal/api/openapi.yaml"
|
SPEC_URL = "https://git.omukk.dev/wrenn/wrenn/raw/branch/dev/internal/api/openapi.yaml"
|
||||||
SPEC_PATH = "api/openapi.yaml"
|
SPEC_PATH = "api/openapi.yaml"
|
||||||
|
|
||||||
generate:
|
generate:
|
||||||
@ -21,4 +21,18 @@ generate:
|
|||||||
--use-schema-description \
|
--use-schema-description \
|
||||||
--target-python-version 3.13 \
|
--target-python-version 3.13 \
|
||||||
--use-annotated \
|
--use-annotated \
|
||||||
--openapi-scopes schemas
|
--openapi-scopes schemas \
|
||||||
|
--formatters ruff-format ruff-check \
|
||||||
|
--input-file-type openapi
|
||||||
|
|
||||||
|
lint:
|
||||||
|
uv run ruff check src/
|
||||||
|
uv run ruff format --check src/
|
||||||
|
|
||||||
|
test:
|
||||||
|
uv run pytest tests/test_client.py -v
|
||||||
|
|
||||||
|
test-integration:
|
||||||
|
uv run pytest tests/ -v -m "integration or not integration"
|
||||||
|
|
||||||
|
check: lint test
|
||||||
|
|||||||
527
README.md
527
README.md
@ -1,3 +1,526 @@
|
|||||||
# python-sdk
|
# Wrenn Python SDK
|
||||||
|
|
||||||
Python SDK for wrenn
|
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
|
||||||
|
|
||||||
|
```bash
|
||||||
|
pip install wrenn
|
||||||
|
```
|
||||||
|
|
||||||
|
Requires Python 3.13+.
|
||||||
|
|
||||||
|
## Authentication
|
||||||
|
|
||||||
|
Set the `WRENN_API_KEY` environment variable:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export WRENN_API_KEY="wrn_your_api_key_here"
|
||||||
|
```
|
||||||
|
|
||||||
|
Optionally override the API base URL:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export WRENN_BASE_URL="https://app.wrenn.dev/api" # default
|
||||||
|
```
|
||||||
|
|
||||||
|
You can also pass credentials directly:
|
||||||
|
|
||||||
|
```python
|
||||||
|
from wrenn import Capsule
|
||||||
|
|
||||||
|
capsule = Capsule(api_key="wrn_...", base_url="https://...")
|
||||||
|
```
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Wrenn Capsules
|
||||||
|
|
||||||
|
### 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
|
||||||
|
|
||||||
|
Use capsules as context managers for automatic cleanup (destroys capsule on exit):
|
||||||
|
|
||||||
|
```python
|
||||||
|
with Capsule(template="minimal", wait=True) as capsule:
|
||||||
|
capsule.commands.run("echo hello")
|
||||||
|
# capsule is automatically destroyed
|
||||||
|
```
|
||||||
|
|
||||||
|
### Connecting to Existing Capsules
|
||||||
|
|
||||||
|
Attach to a running capsule by ID. If it's paused, it will be resumed automatically:
|
||||||
|
|
||||||
|
```python
|
||||||
|
capsule = Capsule.connect("cl-abc123")
|
||||||
|
result = capsule.commands.run("echo still running")
|
||||||
|
```
|
||||||
|
|
||||||
|
For code interpreter capsules:
|
||||||
|
|
||||||
|
```python
|
||||||
|
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:
|
||||||
|
case "stdout":
|
||||||
|
print(event.data, end="")
|
||||||
|
case "stderr":
|
||||||
|
print(event.data, end="", file=sys.stderr)
|
||||||
|
case "exit":
|
||||||
|
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="")
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Process Management
|
||||||
|
|
||||||
|
```python
|
||||||
|
# List running processes
|
||||||
|
for proc in capsule.commands.list():
|
||||||
|
print(proc.pid, proc.cmd, proc.tag)
|
||||||
|
|
||||||
|
# Kill a process
|
||||||
|
capsule.commands.kill(pid=1234)
|
||||||
|
```
|
||||||
|
|
||||||
|
### Filesystem
|
||||||
|
|
||||||
|
Files are accessed via `capsule.files`:
|
||||||
|
|
||||||
|
```python
|
||||||
|
# Write and read files
|
||||||
|
capsule.files.write("/app/main.py", "print('hello')")
|
||||||
|
content = capsule.files.read("/app/main.py") # str
|
||||||
|
raw = capsule.files.read_bytes("/app/main.py") # bytes
|
||||||
|
|
||||||
|
# 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():
|
||||||
|
yield b"chunk1"
|
||||||
|
yield b"chunk2"
|
||||||
|
|
||||||
|
capsule.files.upload_stream("/data/large.bin", chunks())
|
||||||
|
|
||||||
|
# Streaming download
|
||||||
|
for chunk in capsule.files.download_stream("/data/large.bin"):
|
||||||
|
process(chunk)
|
||||||
|
```
|
||||||
|
|
||||||
|
### Interactive Terminal (PTY)
|
||||||
|
|
||||||
|
```python
|
||||||
|
import sys
|
||||||
|
|
||||||
|
with capsule.pty(cmd="/bin/bash", cols=120, rows=40, cwd="/home/user") as term:
|
||||||
|
term.write(b"ls -la\n")
|
||||||
|
for event in term:
|
||||||
|
if event.type == "output":
|
||||||
|
sys.stdout.buffer.write(event.data)
|
||||||
|
elif event.type == "exit":
|
||||||
|
break
|
||||||
|
|
||||||
|
# Reconnect to an existing session
|
||||||
|
with capsule.pty_connect(term.tag) as term:
|
||||||
|
term.write(b"echo reconnected\n")
|
||||||
|
```
|
||||||
|
|
||||||
|
**PtySession methods:**
|
||||||
|
|
||||||
|
| Method | Description |
|
||||||
|
|--------|-------------|
|
||||||
|
| `write(data: bytes)` | Send raw bytes to stdin |
|
||||||
|
| `resize(cols, rows)` | Resize the terminal |
|
||||||
|
| `kill()` | Send SIGKILL to the process |
|
||||||
|
| `tag` | Session tag (after `started` event) |
|
||||||
|
| `pid` | Process PID (after `started` event) |
|
||||||
|
|
||||||
|
### Proxy URL
|
||||||
|
|
||||||
|
Access services running inside a capsule:
|
||||||
|
|
||||||
|
```python
|
||||||
|
url = capsule.get_url(8080)
|
||||||
|
# "wss://8080-cl-abc123.app.wrenn.dev"
|
||||||
|
```
|
||||||
|
|
||||||
|
### Snapshots
|
||||||
|
|
||||||
|
Create reusable templates from running capsules:
|
||||||
|
|
||||||
|
```python
|
||||||
|
template = capsule.create_snapshot(name="my-template", overwrite=True)
|
||||||
|
```
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Code Interpreter
|
||||||
|
|
||||||
|
The `wrenn.code_interpreter` module provides a specialized capsule for stateful code execution via a persistent Jupyter kernel.
|
||||||
|
|
||||||
|
### Quick Start
|
||||||
|
|
||||||
|
```python
|
||||||
|
from wrenn.code_interpreter import Capsule
|
||||||
|
|
||||||
|
with Capsule(wait=True) as capsule:
|
||||||
|
result = capsule.run_code("print('hello')")
|
||||||
|
print(result.text) # "hello"
|
||||||
|
```
|
||||||
|
|
||||||
|
### Stateful Execution
|
||||||
|
|
||||||
|
Variables, imports, and function definitions persist across `run_code` calls:
|
||||||
|
|
||||||
|
```python
|
||||||
|
from wrenn.code_interpreter import Capsule
|
||||||
|
|
||||||
|
with Capsule(wait=True) as capsule:
|
||||||
|
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"
|
||||||
|
```
|
||||||
|
|
||||||
|
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.
|
||||||
|
|
||||||
|
### Error Handling in Code
|
||||||
|
|
||||||
|
```python
|
||||||
|
result = capsule.run_code("1 / 0")
|
||||||
|
print(result.error) # "ZeroDivisionError: division by zero\n..."
|
||||||
|
```
|
||||||
|
|
||||||
|
### Rich Output
|
||||||
|
|
||||||
|
```python
|
||||||
|
result = capsule.run_code("""
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
plt.plot([1, 2, 3])
|
||||||
|
plt.savefig('/tmp/plot.png')
|
||||||
|
plt.show()
|
||||||
|
""")
|
||||||
|
print(result.data) # {"image/png": "base64...", "text/plain": "..."}
|
||||||
|
```
|
||||||
|
|
||||||
|
### 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
|
||||||
|
|
||||||
|
All operations have async variants via `AsyncCapsule`:
|
||||||
|
|
||||||
|
### Async Capsule
|
||||||
|
|
||||||
|
```python
|
||||||
|
from wrenn import AsyncCapsule
|
||||||
|
|
||||||
|
async with await AsyncCapsule.create(template="minimal", wait=True) as capsule:
|
||||||
|
result = await capsule.commands.run("echo hello")
|
||||||
|
print(result.stdout)
|
||||||
|
|
||||||
|
await capsule.files.write("/app/file.txt", "data")
|
||||||
|
entries = await capsule.files.list("/app")
|
||||||
|
|
||||||
|
await capsule.pause()
|
||||||
|
await capsule.resume()
|
||||||
|
```
|
||||||
|
|
||||||
|
### Async Code Interpreter
|
||||||
|
|
||||||
|
```python
|
||||||
|
from wrenn.code_interpreter import AsyncCapsule
|
||||||
|
|
||||||
|
async with await AsyncCapsule.create(wait=True) as capsule:
|
||||||
|
result = await capsule.run_code("2 + 2")
|
||||||
|
print(result.text) # "4"
|
||||||
|
```
|
||||||
|
|
||||||
|
### Async PTY
|
||||||
|
|
||||||
|
```python
|
||||||
|
async with capsule.pty(cmd="/bin/bash") as term:
|
||||||
|
await term.write(b"ls -la\n")
|
||||||
|
async for event in term:
|
||||||
|
if event.type == "output":
|
||||||
|
sys.stdout.buffer.write(event.data)
|
||||||
|
```
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Error Handling
|
||||||
|
|
||||||
|
The SDK maps server error codes to typed exceptions:
|
||||||
|
|
||||||
|
```python
|
||||||
|
from wrenn import (
|
||||||
|
WrennError,
|
||||||
|
WrennValidationError, # 400
|
||||||
|
WrennAuthenticationError, # 401
|
||||||
|
WrennForbiddenError, # 403
|
||||||
|
WrennNotFoundError, # 404
|
||||||
|
WrennConflictError, # 409
|
||||||
|
WrennHostHasCapsulesError, # 409 (host has running capsules)
|
||||||
|
WrennAgentError, # 502
|
||||||
|
WrennInternalError, # 500
|
||||||
|
WrennHostUnavailableError, # 503
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
Capsule.get_info("nonexistent")
|
||||||
|
except WrennNotFoundError as e:
|
||||||
|
print(e.code) # "not_found"
|
||||||
|
print(e.message) # "capsule not found"
|
||||||
|
print(e.status_code) # 404
|
||||||
|
```
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
This project uses [uv](https://docs.astral.sh/uv/) for dependency management.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Install dependencies
|
||||||
|
uv sync
|
||||||
|
|
||||||
|
# Run linting
|
||||||
|
make lint
|
||||||
|
|
||||||
|
# Run unit tests
|
||||||
|
make test
|
||||||
|
|
||||||
|
# Run all tests (including integration)
|
||||||
|
make test-integration
|
||||||
|
```
|
||||||
|
|
||||||
|
### Running Integration Tests
|
||||||
|
|
||||||
|
Integration tests require a live Wrenn server:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export WRENN_API_KEY="wrn_..."
|
||||||
|
export WRENN_BASE_URL="http://localhost:8080" # optional
|
||||||
|
make test-integration
|
||||||
|
```
|
||||||
|
|
||||||
|
## License
|
||||||
|
|
||||||
|
MIT
|
||||||
|
|||||||
1671
api/openapi.yaml
1671
api/openapi.yaml
File diff suppressed because it is too large
Load Diff
@ -8,8 +8,11 @@ authors = [
|
|||||||
]
|
]
|
||||||
requires-python = ">=3.13"
|
requires-python = ">=3.13"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
|
"email-validator>=2.3.0",
|
||||||
"httpx>=0.28.1",
|
"httpx>=0.28.1",
|
||||||
|
"httpx-ws>=0.9.0",
|
||||||
"pydantic>=2.12.5",
|
"pydantic>=2.12.5",
|
||||||
|
"python-dotenv>=1.2.2",
|
||||||
]
|
]
|
||||||
|
|
||||||
[build-system]
|
[build-system]
|
||||||
@ -18,9 +21,15 @@ build-backend = "hatchling.build"
|
|||||||
|
|
||||||
[dependency-groups]
|
[dependency-groups]
|
||||||
dev = [
|
dev = [
|
||||||
"datamodel-code-generator>=0.56.0",
|
"datamodel-code-generator[ruff]>=0.56.0",
|
||||||
"mypy>=1.20.0",
|
"mypy>=1.20.0",
|
||||||
"pytest>=9.0.3",
|
"pytest>=9.0.3",
|
||||||
"pytest-asyncio>=1.3.0",
|
"pytest-asyncio>=1.3.0",
|
||||||
|
"respx>=0.23.1",
|
||||||
"ruff>=0.15.10",
|
"ruff>=0.15.10",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[tool.pytest.ini_options]
|
||||||
|
markers = [
|
||||||
|
"integration: integration tests (require live server)",
|
||||||
|
]
|
||||||
|
|||||||
@ -1,2 +1,89 @@
|
|||||||
def hello() -> str:
|
from wrenn.async_capsule import AsyncCapsule
|
||||||
return "Hello from wrenn!"
|
from wrenn.capsule import Capsule
|
||||||
|
from wrenn.client import AsyncWrennClient, WrennClient
|
||||||
|
from wrenn.commands import (
|
||||||
|
CommandHandle,
|
||||||
|
CommandResult,
|
||||||
|
ProcessInfo,
|
||||||
|
StreamErrorEvent,
|
||||||
|
StreamEvent,
|
||||||
|
StreamExitEvent,
|
||||||
|
StreamStartEvent,
|
||||||
|
StreamStderrEvent,
|
||||||
|
StreamStdoutEvent,
|
||||||
|
)
|
||||||
|
from wrenn.exceptions import (
|
||||||
|
WrennAgentError,
|
||||||
|
WrennAuthenticationError,
|
||||||
|
WrennConflictError,
|
||||||
|
WrennError,
|
||||||
|
WrennForbiddenError,
|
||||||
|
WrennHostHasCapsulesError,
|
||||||
|
WrennHostUnavailableError,
|
||||||
|
WrennInternalError,
|
||||||
|
WrennNotFoundError,
|
||||||
|
WrennValidationError,
|
||||||
|
)
|
||||||
|
from wrenn.models import FileEntry
|
||||||
|
from wrenn.pty import AsyncPtySession, PtyEvent, PtyEventType, PtySession
|
||||||
|
|
||||||
|
__version__ = "0.1.0"
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"__version__",
|
||||||
|
"AsyncCapsule",
|
||||||
|
"AsyncPtySession",
|
||||||
|
"AsyncWrennClient",
|
||||||
|
"Capsule",
|
||||||
|
"CommandHandle",
|
||||||
|
"CommandResult",
|
||||||
|
"FileEntry",
|
||||||
|
"ProcessInfo",
|
||||||
|
"PtyEvent",
|
||||||
|
"PtyEventType",
|
||||||
|
"PtySession",
|
||||||
|
"Sandbox",
|
||||||
|
"StreamErrorEvent",
|
||||||
|
"StreamEvent",
|
||||||
|
"StreamExitEvent",
|
||||||
|
"StreamStartEvent",
|
||||||
|
"StreamStderrEvent",
|
||||||
|
"StreamStdoutEvent",
|
||||||
|
"WrennAgentError",
|
||||||
|
"WrennAuthenticationError",
|
||||||
|
"WrennClient",
|
||||||
|
"WrennConflictError",
|
||||||
|
"WrennError",
|
||||||
|
"WrennForbiddenError",
|
||||||
|
"WrennHostHasCapsulesError",
|
||||||
|
"WrennHostHasSandboxesError",
|
||||||
|
"WrennHostUnavailableError",
|
||||||
|
"WrennInternalError",
|
||||||
|
"WrennNotFoundError",
|
||||||
|
"WrennValidationError",
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
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
|
||||||
|
if name == "WrennHostHasSandboxesError":
|
||||||
|
warnings.warn(
|
||||||
|
"'WrennHostHasSandboxesError' is deprecated, use 'WrennHostHasCapsulesError' instead",
|
||||||
|
FutureWarning,
|
||||||
|
stacklevel=2,
|
||||||
|
)
|
||||||
|
setattr(_module, name, WrennHostHasCapsulesError)
|
||||||
|
return WrennHostHasCapsulesError
|
||||||
|
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||||
|
|||||||
33
src/wrenn/_config.py
Normal file
33
src/wrenn/_config.py
Normal file
@ -0,0 +1,33 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
DEFAULT_BASE_URL = "https://app.wrenn.dev/api"
|
||||||
|
ENV_API_KEY = "WRENN_API_KEY"
|
||||||
|
ENV_BASE_URL = "WRENN_BASE_URL"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class ConnectionConfig:
|
||||||
|
"""Resolved credentials and base URL for Wrenn API calls."""
|
||||||
|
|
||||||
|
api_key: str
|
||||||
|
base_url: str
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_env(
|
||||||
|
cls,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> ConnectionConfig:
|
||||||
|
resolved_key = api_key or os.environ.get(ENV_API_KEY)
|
||||||
|
if not resolved_key:
|
||||||
|
raise ValueError(
|
||||||
|
f"No API key provided. Pass api_key= or set the {ENV_API_KEY} environment variable."
|
||||||
|
)
|
||||||
|
resolved_url = base_url or os.environ.get(ENV_BASE_URL, DEFAULT_BASE_URL)
|
||||||
|
return cls(api_key=resolved_key, base_url=resolved_url)
|
||||||
|
|
||||||
|
def auth_headers(self) -> dict[str, str]:
|
||||||
|
return {"X-API-Key": self.api_key}
|
||||||
273
src/wrenn/async_capsule.py
Normal file
273
src/wrenn/async_capsule.py
Normal file
@ -0,0 +1,273 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import time
|
||||||
|
from collections.abc import AsyncIterator
|
||||||
|
from contextlib import asynccontextmanager
|
||||||
|
|
||||||
|
import httpx_ws
|
||||||
|
|
||||||
|
from wrenn.capsule import _DualMethod, _build_proxy_url
|
||||||
|
from wrenn.client import AsyncWrennClient
|
||||||
|
from wrenn.commands import AsyncCommands
|
||||||
|
from wrenn.files import AsyncFiles
|
||||||
|
from wrenn.models import Capsule as CapsuleModel
|
||||||
|
from wrenn.models import Status, Template
|
||||||
|
from wrenn.pty import AsyncPtySession
|
||||||
|
|
||||||
|
|
||||||
|
class AsyncCapsule:
|
||||||
|
"""Async Wrenn capsule with e2b-compatible interface.
|
||||||
|
|
||||||
|
Create via classmethod::
|
||||||
|
|
||||||
|
capsule = await AsyncCapsule.create(template="minimal")
|
||||||
|
|
||||||
|
Use as async context manager::
|
||||||
|
|
||||||
|
async with await AsyncCapsule.create() as capsule:
|
||||||
|
await capsule.commands.run("echo hello")
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
_capsule_id: str,
|
||||||
|
_client: AsyncWrennClient,
|
||||||
|
_info: CapsuleModel | None = None,
|
||||||
|
) -> None:
|
||||||
|
self._id = _capsule_id
|
||||||
|
self._client = _client
|
||||||
|
self._info = _info
|
||||||
|
|
||||||
|
self.commands = AsyncCommands(_capsule_id, _client.http)
|
||||||
|
self.files = AsyncFiles(_capsule_id, _client.http)
|
||||||
|
|
||||||
|
# ── Properties ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
@property
|
||||||
|
def capsule_id(self) -> str:
|
||||||
|
return self._id
|
||||||
|
|
||||||
|
@property
|
||||||
|
def info(self) -> CapsuleModel | None:
|
||||||
|
return self._info
|
||||||
|
|
||||||
|
# ── Factory classmethods ────────────────────────────────────
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
async def create(
|
||||||
|
cls,
|
||||||
|
template: str | None = None,
|
||||||
|
vcpus: int | None = None,
|
||||||
|
memory_mb: int | None = None,
|
||||||
|
timeout: int | None = None,
|
||||||
|
*,
|
||||||
|
wait: bool = False,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> AsyncCapsule:
|
||||||
|
"""Create a new capsule."""
|
||||||
|
client = AsyncWrennClient(api_key=api_key, base_url=base_url)
|
||||||
|
info = await client.capsules.create(
|
||||||
|
template=template,
|
||||||
|
vcpus=vcpus,
|
||||||
|
memory_mb=memory_mb,
|
||||||
|
timeout_sec=timeout,
|
||||||
|
)
|
||||||
|
capsule = cls(
|
||||||
|
_capsule_id=info.id,
|
||||||
|
_client=client,
|
||||||
|
_info=info,
|
||||||
|
)
|
||||||
|
if wait:
|
||||||
|
await capsule.wait_ready()
|
||||||
|
return capsule
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
async def connect(
|
||||||
|
cls,
|
||||||
|
capsule_id: str,
|
||||||
|
*,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> AsyncCapsule:
|
||||||
|
"""Connect to an existing capsule. Resumes it if paused."""
|
||||||
|
client = AsyncWrennClient(api_key=api_key, base_url=base_url)
|
||||||
|
info = await client.capsules.get(capsule_id)
|
||||||
|
|
||||||
|
if info.status == Status.paused:
|
||||||
|
info = await client.capsules.resume(capsule_id)
|
||||||
|
|
||||||
|
return cls(
|
||||||
|
_capsule_id=capsule_id,
|
||||||
|
_client=client,
|
||||||
|
_info=info,
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── Dual instance/static lifecycle ──────────────────────────
|
||||||
|
|
||||||
|
destroy = _DualMethod("_instance_destroy", "_static_destroy")
|
||||||
|
pause = _DualMethod("_instance_pause", "_static_pause")
|
||||||
|
resume = _DualMethod("_instance_resume", "_static_resume")
|
||||||
|
get_info = _DualMethod("_instance_get_info", "_static_get_info")
|
||||||
|
|
||||||
|
async def _instance_destroy(self) -> None:
|
||||||
|
await self._client.capsules.destroy(self._id)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
async def _static_destroy(
|
||||||
|
cls,
|
||||||
|
capsule_id: str,
|
||||||
|
*,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
async with AsyncWrennClient(api_key=api_key, base_url=base_url) as client:
|
||||||
|
await client.capsules.destroy(capsule_id)
|
||||||
|
|
||||||
|
async def _instance_pause(self) -> CapsuleModel:
|
||||||
|
self._info = await self._client.capsules.pause(self._id)
|
||||||
|
return self._info
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
async def _static_pause(
|
||||||
|
cls,
|
||||||
|
capsule_id: str,
|
||||||
|
*,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> CapsuleModel:
|
||||||
|
async with AsyncWrennClient(api_key=api_key, base_url=base_url) as client:
|
||||||
|
return await client.capsules.pause(capsule_id)
|
||||||
|
|
||||||
|
async def _instance_resume(self) -> CapsuleModel:
|
||||||
|
self._info = await self._client.capsules.resume(self._id)
|
||||||
|
return self._info
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
async def _static_resume(
|
||||||
|
cls,
|
||||||
|
capsule_id: str,
|
||||||
|
*,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> CapsuleModel:
|
||||||
|
async with AsyncWrennClient(api_key=api_key, base_url=base_url) as client:
|
||||||
|
return await client.capsules.resume(capsule_id)
|
||||||
|
|
||||||
|
async def _instance_get_info(self) -> CapsuleModel:
|
||||||
|
self._info = await self._client.capsules.get(self._id)
|
||||||
|
return self._info
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
async def _static_get_info(
|
||||||
|
cls,
|
||||||
|
capsule_id: str,
|
||||||
|
*,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> CapsuleModel:
|
||||||
|
async with AsyncWrennClient(api_key=api_key, base_url=base_url) as client:
|
||||||
|
return await client.capsules.get(capsule_id)
|
||||||
|
|
||||||
|
# ── Instance-only methods ───────────────────────────────────
|
||||||
|
|
||||||
|
async def ping(self) -> None:
|
||||||
|
await self._client.capsules.ping(self._id)
|
||||||
|
|
||||||
|
async def wait_ready(self, timeout: float = 30, interval: float = 0.5) -> None:
|
||||||
|
deadline = time.monotonic() + timeout
|
||||||
|
while time.monotonic() < deadline:
|
||||||
|
info = await self._client.capsules.get(self._id)
|
||||||
|
if info.status == Status.running:
|
||||||
|
self._info = info
|
||||||
|
return
|
||||||
|
if info.status in (Status.error, Status.stopped, Status.paused):
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Capsule entered {info.status} state while waiting"
|
||||||
|
)
|
||||||
|
await asyncio.sleep(interval)
|
||||||
|
raise TimeoutError(
|
||||||
|
f"Capsule {self._id} did not become ready within {timeout}s"
|
||||||
|
)
|
||||||
|
|
||||||
|
async def is_running(self) -> bool:
|
||||||
|
info = await self._instance_get_info()
|
||||||
|
return info.status == Status.running
|
||||||
|
|
||||||
|
# ── Static list ─────────────────────────────────────────────
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
async def list(
|
||||||
|
cls,
|
||||||
|
*,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> list[CapsuleModel]:
|
||||||
|
async with AsyncWrennClient(api_key=api_key, base_url=base_url) as client:
|
||||||
|
return await client.capsules.list()
|
||||||
|
|
||||||
|
# ── PTY ─────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
@asynccontextmanager
|
||||||
|
async def pty(
|
||||||
|
self,
|
||||||
|
cmd: str = "/bin/bash",
|
||||||
|
args: list[str] | None = None,
|
||||||
|
cols: int = 80,
|
||||||
|
rows: int = 24,
|
||||||
|
envs: dict[str, str] | None = None,
|
||||||
|
cwd: str | None = None,
|
||||||
|
) -> AsyncIterator[AsyncPtySession]:
|
||||||
|
async with httpx_ws.aconnect_ws(
|
||||||
|
f"/v1/capsules/{self._id}/pty", client=self._client.http
|
||||||
|
) as ws:
|
||||||
|
session = AsyncPtySession(ws, self._id)
|
||||||
|
await session._send_start(
|
||||||
|
cmd=cmd, args=args, cols=cols, rows=rows, envs=envs, cwd=cwd
|
||||||
|
)
|
||||||
|
yield session
|
||||||
|
|
||||||
|
@asynccontextmanager
|
||||||
|
async def pty_connect(self, tag: str) -> AsyncIterator[AsyncPtySession]:
|
||||||
|
async with httpx_ws.aconnect_ws(
|
||||||
|
f"/v1/capsules/{self._id}/pty", client=self._client.http
|
||||||
|
) as ws:
|
||||||
|
session = AsyncPtySession(ws, self._id)
|
||||||
|
await session._send_connect(tag)
|
||||||
|
yield session
|
||||||
|
|
||||||
|
# ── Proxy helpers ───────────────────────────────────────────
|
||||||
|
|
||||||
|
def get_url(self, port: int) -> str:
|
||||||
|
return _build_proxy_url(self._client._base_url, self._id, port)
|
||||||
|
|
||||||
|
# ── Snapshots ───────────────────────────────────────────────
|
||||||
|
|
||||||
|
async def create_snapshot(
|
||||||
|
self, name: str | None = None, overwrite: bool = False
|
||||||
|
) -> Template:
|
||||||
|
return await self._client.snapshots.create(
|
||||||
|
capsule_id=self._id, name=name, overwrite=overwrite
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── Context manager ─────────────────────────────────────────
|
||||||
|
|
||||||
|
async def __aenter__(self) -> AsyncCapsule:
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(
|
||||||
|
self,
|
||||||
|
exc_type: type[BaseException] | None,
|
||||||
|
exc_val: BaseException | None,
|
||||||
|
exc_tb: object,
|
||||||
|
) -> None:
|
||||||
|
try:
|
||||||
|
await self._instance_destroy()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
await self._client.aclose()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
338
src/wrenn/capsule.py
Normal file
338
src/wrenn/capsule.py
Normal file
@ -0,0 +1,338 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import time
|
||||||
|
from collections.abc import Iterator
|
||||||
|
from contextlib import contextmanager
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
import httpx_ws
|
||||||
|
|
||||||
|
from wrenn.client import WrennClient
|
||||||
|
from wrenn.commands import Commands
|
||||||
|
from wrenn.files import Files
|
||||||
|
from wrenn.models import Capsule as CapsuleModel
|
||||||
|
from wrenn.models import Status, Template
|
||||||
|
from wrenn.pty import PtySession
|
||||||
|
|
||||||
|
|
||||||
|
def _build_proxy_url(base_url: str, capsule_id: str | None, port: int) -> str:
|
||||||
|
parsed = httpx.URL(base_url)
|
||||||
|
host = parsed.host
|
||||||
|
if parsed.port:
|
||||||
|
host = f"{host}:{parsed.port}"
|
||||||
|
scheme = "ws" if parsed.scheme == "http" else "wss"
|
||||||
|
return f"{scheme}://{port}-{capsule_id}.{host}"
|
||||||
|
|
||||||
|
|
||||||
|
class _DualMethod:
|
||||||
|
"""Descriptor that dispatches to instance method or classmethod depending on call site."""
|
||||||
|
|
||||||
|
def __init__(self, instance_fn_name: str, static_fn_name: str) -> None:
|
||||||
|
self._ifn = instance_fn_name
|
||||||
|
self._sfn = static_fn_name
|
||||||
|
|
||||||
|
def __set_name__(self, owner: type, name: str) -> None:
|
||||||
|
self._name = name
|
||||||
|
|
||||||
|
def __get__(self, obj: Any, cls: type) -> Any:
|
||||||
|
if obj is None:
|
||||||
|
return getattr(cls, self._sfn)
|
||||||
|
return getattr(obj, self._ifn)
|
||||||
|
|
||||||
|
|
||||||
|
class Capsule:
|
||||||
|
"""A Wrenn capsule (sandbox) with e2b-compatible interface.
|
||||||
|
|
||||||
|
Create directly::
|
||||||
|
|
||||||
|
capsule = Capsule(api_key="wrn_...")
|
||||||
|
capsule = Capsule(template="minimal") # reads WRENN_API_KEY env
|
||||||
|
|
||||||
|
Or via classmethod::
|
||||||
|
|
||||||
|
capsule = Capsule.create(template="minimal")
|
||||||
|
|
||||||
|
Use as context manager for automatic cleanup::
|
||||||
|
|
||||||
|
with Capsule() as capsule:
|
||||||
|
capsule.commands.run("echo hello")
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
template: str | None = None,
|
||||||
|
vcpus: int | None = None,
|
||||||
|
memory_mb: int | None = None,
|
||||||
|
timeout: int | None = None,
|
||||||
|
*,
|
||||||
|
wait: bool = False,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
# Private: used by classmethods to skip creation
|
||||||
|
_capsule_id: str | None = None,
|
||||||
|
_client: WrennClient | None = None,
|
||||||
|
_info: CapsuleModel | None = None,
|
||||||
|
) -> None:
|
||||||
|
if _capsule_id is not None:
|
||||||
|
# Internal construction path (from create/connect classmethods)
|
||||||
|
assert _client is not None
|
||||||
|
self._id = _capsule_id
|
||||||
|
self._client = _client
|
||||||
|
self._info = _info
|
||||||
|
else:
|
||||||
|
# Public construction: create a capsule immediately
|
||||||
|
self._client = WrennClient(api_key=api_key, base_url=base_url)
|
||||||
|
self._info = self._client.capsules.create(
|
||||||
|
template=template,
|
||||||
|
vcpus=vcpus,
|
||||||
|
memory_mb=memory_mb,
|
||||||
|
timeout_sec=timeout,
|
||||||
|
)
|
||||||
|
self._id = self._info.id
|
||||||
|
|
||||||
|
self.commands = Commands(self._id, self._client.http)
|
||||||
|
self.files = Files(self._id, self._client.http)
|
||||||
|
|
||||||
|
if wait:
|
||||||
|
self.wait_ready()
|
||||||
|
|
||||||
|
# ── Properties ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
@property
|
||||||
|
def capsule_id(self) -> str:
|
||||||
|
return self._id
|
||||||
|
|
||||||
|
@property
|
||||||
|
def info(self) -> CapsuleModel | None:
|
||||||
|
return self._info
|
||||||
|
|
||||||
|
# ── Factory classmethods ────────────────────────────────────
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def create(
|
||||||
|
cls,
|
||||||
|
template: str | None = None,
|
||||||
|
vcpus: int | None = None,
|
||||||
|
memory_mb: int | None = None,
|
||||||
|
timeout: int | None = None,
|
||||||
|
*,
|
||||||
|
wait: bool = False,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> Capsule:
|
||||||
|
"""Create a new capsule. Alias for ``Capsule(...)``."""
|
||||||
|
return cls(
|
||||||
|
template=template,
|
||||||
|
vcpus=vcpus,
|
||||||
|
memory_mb=memory_mb,
|
||||||
|
timeout=timeout,
|
||||||
|
wait=wait,
|
||||||
|
api_key=api_key,
|
||||||
|
base_url=base_url,
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def connect(
|
||||||
|
cls,
|
||||||
|
capsule_id: str,
|
||||||
|
*,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> Capsule:
|
||||||
|
"""Connect to an existing capsule. Resumes it if paused."""
|
||||||
|
client = WrennClient(api_key=api_key, base_url=base_url)
|
||||||
|
info = client.capsules.get(capsule_id)
|
||||||
|
|
||||||
|
if info.status == Status.paused:
|
||||||
|
info = client.capsules.resume(capsule_id)
|
||||||
|
|
||||||
|
return cls(
|
||||||
|
_capsule_id=capsule_id,
|
||||||
|
_client=client,
|
||||||
|
_info=info,
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── Dual instance/static lifecycle ──────────────────────────
|
||||||
|
|
||||||
|
destroy = _DualMethod("_instance_destroy", "_static_destroy")
|
||||||
|
pause = _DualMethod("_instance_pause", "_static_pause")
|
||||||
|
resume = _DualMethod("_instance_resume", "_static_resume")
|
||||||
|
get_info = _DualMethod("_instance_get_info", "_static_get_info")
|
||||||
|
|
||||||
|
def _instance_destroy(self) -> None:
|
||||||
|
"""Destroy this capsule."""
|
||||||
|
self._client.capsules.destroy(self._id)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _static_destroy(
|
||||||
|
cls,
|
||||||
|
capsule_id: str,
|
||||||
|
*,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Destroy a capsule by ID."""
|
||||||
|
with WrennClient(api_key=api_key, base_url=base_url) as client:
|
||||||
|
client.capsules.destroy(capsule_id)
|
||||||
|
|
||||||
|
def _instance_pause(self) -> CapsuleModel:
|
||||||
|
"""Pause this capsule."""
|
||||||
|
self._info = self._client.capsules.pause(self._id)
|
||||||
|
return self._info
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _static_pause(
|
||||||
|
cls,
|
||||||
|
capsule_id: str,
|
||||||
|
*,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> CapsuleModel:
|
||||||
|
"""Pause a capsule by ID."""
|
||||||
|
with WrennClient(api_key=api_key, base_url=base_url) as client:
|
||||||
|
return client.capsules.pause(capsule_id)
|
||||||
|
|
||||||
|
def _instance_resume(self) -> CapsuleModel:
|
||||||
|
"""Resume this capsule."""
|
||||||
|
self._info = self._client.capsules.resume(self._id)
|
||||||
|
return self._info
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _static_resume(
|
||||||
|
cls,
|
||||||
|
capsule_id: str,
|
||||||
|
*,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> CapsuleModel:
|
||||||
|
"""Resume a capsule by ID."""
|
||||||
|
with WrennClient(api_key=api_key, base_url=base_url) as client:
|
||||||
|
return client.capsules.resume(capsule_id)
|
||||||
|
|
||||||
|
def _instance_get_info(self) -> CapsuleModel:
|
||||||
|
"""Get current info for this capsule."""
|
||||||
|
self._info = self._client.capsules.get(self._id)
|
||||||
|
return self._info
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _static_get_info(
|
||||||
|
cls,
|
||||||
|
capsule_id: str,
|
||||||
|
*,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> CapsuleModel:
|
||||||
|
"""Get capsule info by ID."""
|
||||||
|
with WrennClient(api_key=api_key, base_url=base_url) as client:
|
||||||
|
return client.capsules.get(capsule_id)
|
||||||
|
|
||||||
|
# ── Instance-only methods ───────────────────────────────────
|
||||||
|
|
||||||
|
def ping(self) -> None:
|
||||||
|
"""Reset the capsule inactivity timer."""
|
||||||
|
self._client.capsules.ping(self._id)
|
||||||
|
|
||||||
|
def wait_ready(self, timeout: float = 30, interval: float = 0.5) -> None:
|
||||||
|
"""Block until the capsule status is ``running``."""
|
||||||
|
deadline = time.monotonic() + timeout
|
||||||
|
while time.monotonic() < deadline:
|
||||||
|
info = self._client.capsules.get(self._id)
|
||||||
|
if info.status == Status.running:
|
||||||
|
self._info = info
|
||||||
|
return
|
||||||
|
if info.status in (Status.error, Status.stopped, Status.paused):
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Capsule entered {info.status} state while waiting"
|
||||||
|
)
|
||||||
|
time.sleep(interval)
|
||||||
|
raise TimeoutError(
|
||||||
|
f"Capsule {self._id} did not become ready within {timeout}s"
|
||||||
|
)
|
||||||
|
|
||||||
|
def is_running(self) -> bool:
|
||||||
|
info = self._instance_get_info()
|
||||||
|
return info.status == Status.running
|
||||||
|
|
||||||
|
# ── Static list ─────────────────────────────────────────────
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def list(
|
||||||
|
cls,
|
||||||
|
*,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> list[CapsuleModel]:
|
||||||
|
"""List all capsules for the team."""
|
||||||
|
with WrennClient(api_key=api_key, base_url=base_url) as client:
|
||||||
|
return client.capsules.list()
|
||||||
|
|
||||||
|
# ── PTY ─────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def pty(
|
||||||
|
self,
|
||||||
|
cmd: str = "/bin/bash",
|
||||||
|
args: list[str] | None = None,
|
||||||
|
cols: int = 80,
|
||||||
|
rows: int = 24,
|
||||||
|
envs: dict[str, str] | None = None,
|
||||||
|
cwd: str | None = None,
|
||||||
|
) -> Iterator[PtySession]:
|
||||||
|
"""Open an interactive PTY session."""
|
||||||
|
with httpx_ws.connect_ws(
|
||||||
|
f"/v1/capsules/{self._id}/pty", client=self._client.http
|
||||||
|
) as ws:
|
||||||
|
session = PtySession(ws, self._id)
|
||||||
|
session._send_start(
|
||||||
|
cmd=cmd, args=args, cols=cols, rows=rows, envs=envs, cwd=cwd
|
||||||
|
)
|
||||||
|
yield session
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def pty_connect(self, tag: str) -> Iterator[PtySession]:
|
||||||
|
"""Reconnect to an existing PTY session by tag."""
|
||||||
|
with httpx_ws.connect_ws(
|
||||||
|
f"/v1/capsules/{self._id}/pty", client=self._client.http
|
||||||
|
) as ws:
|
||||||
|
session = PtySession(ws, self._id)
|
||||||
|
session._send_connect(tag)
|
||||||
|
yield session
|
||||||
|
|
||||||
|
# ── Proxy helpers ───────────────────────────────────────────
|
||||||
|
|
||||||
|
def get_url(self, port: int) -> str:
|
||||||
|
"""Get the proxy URL for a port inside this capsule."""
|
||||||
|
return _build_proxy_url(self._client._base_url, self._id, port)
|
||||||
|
|
||||||
|
# ── Snapshots ───────────────────────────────────────────────
|
||||||
|
|
||||||
|
def create_snapshot(
|
||||||
|
self, name: str | None = None, overwrite: bool = False
|
||||||
|
) -> Template:
|
||||||
|
"""Create a snapshot template from this capsule."""
|
||||||
|
return self._client.snapshots.create(
|
||||||
|
capsule_id=self._id, name=name, overwrite=overwrite
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── Context manager ─────────────────────────────────────────
|
||||||
|
|
||||||
|
def __enter__(self) -> Capsule:
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __exit__(
|
||||||
|
self,
|
||||||
|
exc_type: type[BaseException] | None,
|
||||||
|
exc_val: BaseException | None,
|
||||||
|
exc_tb: object,
|
||||||
|
) -> None:
|
||||||
|
try:
|
||||||
|
self._instance_destroy()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
self._client.close()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
281
src/wrenn/client.py
Normal file
281
src/wrenn/client.py
Normal file
@ -0,0 +1,281 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
|
||||||
|
from wrenn._config import DEFAULT_BASE_URL, ENV_API_KEY, ENV_BASE_URL
|
||||||
|
from wrenn.exceptions import handle_response
|
||||||
|
from wrenn.models import (
|
||||||
|
Template,
|
||||||
|
)
|
||||||
|
from wrenn.models import (
|
||||||
|
Capsule as CapsuleModel,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _resolve_api_key(api_key: str | None) -> str:
|
||||||
|
resolved = api_key or os.environ.get(ENV_API_KEY)
|
||||||
|
if not resolved:
|
||||||
|
raise ValueError(
|
||||||
|
f"No API key provided. Pass api_key= or set the {ENV_API_KEY} environment variable."
|
||||||
|
)
|
||||||
|
return resolved
|
||||||
|
|
||||||
|
|
||||||
|
class CapsulesResource:
|
||||||
|
"""Sync capsule control-plane operations."""
|
||||||
|
|
||||||
|
def __init__(self, http: httpx.Client) -> None:
|
||||||
|
self._http = http
|
||||||
|
|
||||||
|
def create(
|
||||||
|
self,
|
||||||
|
template: str | None = None,
|
||||||
|
vcpus: int | None = None,
|
||||||
|
memory_mb: int | None = None,
|
||||||
|
timeout_sec: int | None = None,
|
||||||
|
) -> CapsuleModel:
|
||||||
|
payload: dict = {}
|
||||||
|
if template is not None:
|
||||||
|
payload["template"] = template
|
||||||
|
if vcpus is not None:
|
||||||
|
payload["vcpus"] = vcpus
|
||||||
|
if memory_mb is not None:
|
||||||
|
payload["memory_mb"] = memory_mb
|
||||||
|
if timeout_sec is not None:
|
||||||
|
payload["timeout_sec"] = timeout_sec
|
||||||
|
resp = self._http.post("/v1/capsules", json=payload)
|
||||||
|
return CapsuleModel.model_validate(handle_response(resp))
|
||||||
|
|
||||||
|
def list(self) -> list[CapsuleModel]:
|
||||||
|
resp = self._http.get("/v1/capsules")
|
||||||
|
return [CapsuleModel.model_validate(item) for item in handle_response(resp)]
|
||||||
|
|
||||||
|
def get(self, id: str) -> CapsuleModel:
|
||||||
|
resp = self._http.get(f"/v1/capsules/{id}")
|
||||||
|
return CapsuleModel.model_validate(handle_response(resp))
|
||||||
|
|
||||||
|
def destroy(self, id: str) -> None:
|
||||||
|
resp = self._http.delete(f"/v1/capsules/{id}")
|
||||||
|
handle_response(resp)
|
||||||
|
|
||||||
|
def pause(self, id: str) -> CapsuleModel:
|
||||||
|
resp = self._http.post(f"/v1/capsules/{id}/pause")
|
||||||
|
return CapsuleModel.model_validate(handle_response(resp))
|
||||||
|
|
||||||
|
def resume(self, id: str) -> CapsuleModel:
|
||||||
|
resp = self._http.post(f"/v1/capsules/{id}/resume")
|
||||||
|
return CapsuleModel.model_validate(handle_response(resp))
|
||||||
|
|
||||||
|
def ping(self, id: str) -> None:
|
||||||
|
resp = self._http.post(f"/v1/capsules/{id}/ping")
|
||||||
|
handle_response(resp)
|
||||||
|
|
||||||
|
|
||||||
|
class AsyncCapsulesResource:
|
||||||
|
"""Async capsule control-plane operations."""
|
||||||
|
|
||||||
|
def __init__(self, http: httpx.AsyncClient) -> None:
|
||||||
|
self._http = http
|
||||||
|
|
||||||
|
async def create(
|
||||||
|
self,
|
||||||
|
template: str | None = None,
|
||||||
|
vcpus: int | None = None,
|
||||||
|
memory_mb: int | None = None,
|
||||||
|
timeout_sec: int | None = None,
|
||||||
|
) -> CapsuleModel:
|
||||||
|
payload: dict = {}
|
||||||
|
if template is not None:
|
||||||
|
payload["template"] = template
|
||||||
|
if vcpus is not None:
|
||||||
|
payload["vcpus"] = vcpus
|
||||||
|
if memory_mb is not None:
|
||||||
|
payload["memory_mb"] = memory_mb
|
||||||
|
if timeout_sec is not None:
|
||||||
|
payload["timeout_sec"] = timeout_sec
|
||||||
|
resp = await self._http.post("/v1/capsules", json=payload)
|
||||||
|
return CapsuleModel.model_validate(handle_response(resp))
|
||||||
|
|
||||||
|
async def list(self) -> list[CapsuleModel]:
|
||||||
|
resp = await self._http.get("/v1/capsules")
|
||||||
|
return [CapsuleModel.model_validate(item) for item in handle_response(resp)]
|
||||||
|
|
||||||
|
async def get(self, id: str) -> CapsuleModel:
|
||||||
|
resp = await self._http.get(f"/v1/capsules/{id}")
|
||||||
|
return CapsuleModel.model_validate(handle_response(resp))
|
||||||
|
|
||||||
|
async def destroy(self, id: str) -> None:
|
||||||
|
resp = await self._http.delete(f"/v1/capsules/{id}")
|
||||||
|
handle_response(resp)
|
||||||
|
|
||||||
|
async def pause(self, id: str) -> CapsuleModel:
|
||||||
|
resp = await self._http.post(f"/v1/capsules/{id}/pause")
|
||||||
|
return CapsuleModel.model_validate(handle_response(resp))
|
||||||
|
|
||||||
|
async def resume(self, id: str) -> CapsuleModel:
|
||||||
|
resp = await self._http.post(f"/v1/capsules/{id}/resume")
|
||||||
|
return CapsuleModel.model_validate(handle_response(resp))
|
||||||
|
|
||||||
|
async def ping(self, id: str) -> None:
|
||||||
|
resp = await self._http.post(f"/v1/capsules/{id}/ping")
|
||||||
|
handle_response(resp)
|
||||||
|
|
||||||
|
|
||||||
|
class SnapshotsResource:
|
||||||
|
"""Sync snapshot operations."""
|
||||||
|
|
||||||
|
def __init__(self, http: httpx.Client) -> None:
|
||||||
|
self._http = http
|
||||||
|
|
||||||
|
def create(
|
||||||
|
self,
|
||||||
|
capsule_id: str,
|
||||||
|
name: str | None = None,
|
||||||
|
overwrite: bool = False,
|
||||||
|
) -> Template:
|
||||||
|
payload: dict = {"sandbox_id": capsule_id}
|
||||||
|
if name is not None:
|
||||||
|
payload["name"] = name
|
||||||
|
params: dict = {}
|
||||||
|
if overwrite:
|
||||||
|
params["overwrite"] = "true"
|
||||||
|
resp = self._http.post("/v1/snapshots", json=payload, params=params)
|
||||||
|
return Template.model_validate(handle_response(resp))
|
||||||
|
|
||||||
|
def list(self, type: str | None = None) -> list[Template]:
|
||||||
|
params: dict = {}
|
||||||
|
if type is not None:
|
||||||
|
params["type"] = type
|
||||||
|
resp = self._http.get("/v1/snapshots", params=params)
|
||||||
|
return [Template.model_validate(item) for item in handle_response(resp)]
|
||||||
|
|
||||||
|
def delete(self, name: str) -> None:
|
||||||
|
resp = self._http.delete(f"/v1/snapshots/{name}")
|
||||||
|
handle_response(resp)
|
||||||
|
|
||||||
|
|
||||||
|
class AsyncSnapshotsResource:
|
||||||
|
"""Async snapshot operations."""
|
||||||
|
|
||||||
|
def __init__(self, http: httpx.AsyncClient) -> None:
|
||||||
|
self._http = http
|
||||||
|
|
||||||
|
async def create(
|
||||||
|
self,
|
||||||
|
capsule_id: str,
|
||||||
|
name: str | None = None,
|
||||||
|
overwrite: bool = False,
|
||||||
|
) -> Template:
|
||||||
|
payload: dict = {"sandbox_id": capsule_id}
|
||||||
|
if name is not None:
|
||||||
|
payload["name"] = name
|
||||||
|
params: dict = {}
|
||||||
|
if overwrite:
|
||||||
|
params["overwrite"] = "true"
|
||||||
|
resp = await self._http.post("/v1/snapshots", json=payload, params=params)
|
||||||
|
return Template.model_validate(handle_response(resp))
|
||||||
|
|
||||||
|
async def list(self, type: str | None = None) -> list[Template]:
|
||||||
|
params: dict = {}
|
||||||
|
if type is not None:
|
||||||
|
params["type"] = type
|
||||||
|
resp = await self._http.get("/v1/snapshots", params=params)
|
||||||
|
return [Template.model_validate(item) for item in handle_response(resp)]
|
||||||
|
|
||||||
|
async def delete(self, name: str) -> None:
|
||||||
|
resp = await self._http.delete(f"/v1/snapshots/{name}")
|
||||||
|
handle_response(resp)
|
||||||
|
|
||||||
|
|
||||||
|
class WrennClient:
|
||||||
|
"""Synchronous client for the Wrenn API.
|
||||||
|
|
||||||
|
Authenticates with an API key.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
api_key: API key (``wrn_...``). Falls back to ``WRENN_API_KEY`` env var.
|
||||||
|
base_url: Wrenn API base URL.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
self._api_key = _resolve_api_key(api_key)
|
||||||
|
self._base_url = base_url or os.environ.get(ENV_BASE_URL, DEFAULT_BASE_URL)
|
||||||
|
self._http = httpx.Client(
|
||||||
|
base_url=self._base_url,
|
||||||
|
headers={"X-API-Key": self._api_key},
|
||||||
|
)
|
||||||
|
|
||||||
|
self.capsules = CapsulesResource(self._http)
|
||||||
|
self.snapshots = SnapshotsResource(self._http)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def http(self) -> httpx.Client:
|
||||||
|
"""The underlying httpx.Client (for sub-objects that need direct access)."""
|
||||||
|
return self._http
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
"""Close the underlying HTTP connection pool."""
|
||||||
|
self._http.close()
|
||||||
|
|
||||||
|
def __enter__(self) -> WrennClient:
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __exit__(
|
||||||
|
self,
|
||||||
|
exc_type: type[BaseException] | None,
|
||||||
|
exc_val: BaseException | None,
|
||||||
|
exc_tb: object,
|
||||||
|
) -> None:
|
||||||
|
self.close()
|
||||||
|
|
||||||
|
|
||||||
|
class AsyncWrennClient:
|
||||||
|
"""Asynchronous client for the Wrenn API.
|
||||||
|
|
||||||
|
Authenticates with an API key.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
api_key: API key (``wrn_...``). Falls back to ``WRENN_API_KEY`` env var.
|
||||||
|
base_url: Wrenn API base URL. Falls back to ``WRENN_BASE_URL`` env var.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
self._api_key = _resolve_api_key(api_key)
|
||||||
|
self._base_url = base_url or os.environ.get(ENV_BASE_URL, DEFAULT_BASE_URL)
|
||||||
|
self._http = httpx.AsyncClient(
|
||||||
|
base_url=self._base_url,
|
||||||
|
headers={"X-API-Key": self._api_key},
|
||||||
|
)
|
||||||
|
|
||||||
|
self.capsules = AsyncCapsulesResource(self._http)
|
||||||
|
self.snapshots = AsyncSnapshotsResource(self._http)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def http(self) -> httpx.AsyncClient:
|
||||||
|
"""The underlying httpx.AsyncClient."""
|
||||||
|
return self._http
|
||||||
|
|
||||||
|
async def aclose(self) -> None:
|
||||||
|
"""Close the underlying async HTTP connection pool."""
|
||||||
|
await self._http.aclose()
|
||||||
|
|
||||||
|
async def __aenter__(self) -> AsyncWrennClient:
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(
|
||||||
|
self,
|
||||||
|
exc_type: type[BaseException] | None,
|
||||||
|
exc_val: BaseException | None,
|
||||||
|
exc_tb: object,
|
||||||
|
) -> None:
|
||||||
|
await self.aclose()
|
||||||
26
src/wrenn/code_interpreter/__init__.py
Normal file
26
src/wrenn/code_interpreter/__init__.py
Normal file
@ -0,0 +1,26 @@
|
|||||||
|
from wrenn.code_interpreter.async_capsule import AsyncCapsule
|
||||||
|
from wrenn.code_interpreter.capsule import Capsule, CodeResult
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"AsyncCapsule",
|
||||||
|
"Capsule",
|
||||||
|
"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}")
|
||||||
221
src/wrenn/code_interpreter/async_capsule.py
Normal file
221
src/wrenn/code_interpreter/async_capsule.py
Normal file
@ -0,0 +1,221 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import time
|
||||||
|
import uuid
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
import httpx_ws
|
||||||
|
|
||||||
|
from wrenn.async_capsule import AsyncCapsule as BaseAsyncCapsule
|
||||||
|
from wrenn.capsule import _build_proxy_url
|
||||||
|
from wrenn.client import AsyncWrennClient
|
||||||
|
from wrenn.code_interpreter.capsule import CodeResult, DEFAULT_TEMPLATE
|
||||||
|
|
||||||
|
|
||||||
|
class AsyncCapsule(BaseAsyncCapsule):
|
||||||
|
"""Async code interpreter capsule with ``run_code`` support.
|
||||||
|
|
||||||
|
Uses ``code-runner-beta`` template by default::
|
||||||
|
|
||||||
|
from wrenn.code_interpreter import AsyncCapsule
|
||||||
|
|
||||||
|
capsule = await AsyncCapsule.create()
|
||||||
|
result = await capsule.run_code("print('hello')")
|
||||||
|
"""
|
||||||
|
|
||||||
|
_kernel_id: str | None
|
||||||
|
_proxy_client: httpx.AsyncClient | None
|
||||||
|
|
||||||
|
def __init__(self, **kwargs) -> None:
|
||||||
|
super().__init__(**kwargs)
|
||||||
|
self._kernel_id = None
|
||||||
|
self._proxy_client = None
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
async def create(
|
||||||
|
cls,
|
||||||
|
template: str | None = None,
|
||||||
|
vcpus: int | None = None,
|
||||||
|
memory_mb: int | None = None,
|
||||||
|
timeout: int | None = None,
|
||||||
|
*,
|
||||||
|
wait: bool = False,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> AsyncCapsule:
|
||||||
|
client = AsyncWrennClient(api_key=api_key, base_url=base_url)
|
||||||
|
info = await client.capsules.create(
|
||||||
|
template=template or DEFAULT_TEMPLATE,
|
||||||
|
vcpus=vcpus,
|
||||||
|
memory_mb=memory_mb,
|
||||||
|
timeout_sec=timeout,
|
||||||
|
)
|
||||||
|
capsule = cls(
|
||||||
|
_capsule_id=info.id,
|
||||||
|
_client=client,
|
||||||
|
_info=info,
|
||||||
|
)
|
||||||
|
if wait:
|
||||||
|
await capsule.wait_ready()
|
||||||
|
return capsule
|
||||||
|
|
||||||
|
def _get_proxy_client(self) -> httpx.AsyncClient:
|
||||||
|
if self._proxy_client is None:
|
||||||
|
url = (
|
||||||
|
_build_proxy_url(self._client._base_url, self._id, 8888)
|
||||||
|
.replace("ws://", "http://")
|
||||||
|
.replace("wss://", "https://")
|
||||||
|
)
|
||||||
|
self._proxy_client = httpx.AsyncClient(
|
||||||
|
base_url=url,
|
||||||
|
headers={"X-API-Key": self._client._api_key},
|
||||||
|
)
|
||||||
|
return self._proxy_client
|
||||||
|
|
||||||
|
async def _ensure_kernel(self, jupyter_timeout: float = 30) -> str:
|
||||||
|
if self._kernel_id is not None:
|
||||||
|
return self._kernel_id
|
||||||
|
|
||||||
|
client = self._get_proxy_client()
|
||||||
|
deadline = time.monotonic() + jupyter_timeout
|
||||||
|
last_exc: Exception | None = None
|
||||||
|
|
||||||
|
while time.monotonic() < deadline:
|
||||||
|
try:
|
||||||
|
# Try to reuse an existing kernel
|
||||||
|
resp = await client.get("/api/kernels")
|
||||||
|
if resp.status_code < 500:
|
||||||
|
resp.raise_for_status()
|
||||||
|
kernels = resp.json()
|
||||||
|
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(
|
||||||
|
f"Jupyter returned {resp.status_code}",
|
||||||
|
request=resp.request,
|
||||||
|
response=resp,
|
||||||
|
)
|
||||||
|
except httpx.HTTPStatusError:
|
||||||
|
raise
|
||||||
|
except Exception as exc:
|
||||||
|
last_exc = exc
|
||||||
|
await asyncio.sleep(0.5)
|
||||||
|
|
||||||
|
raise TimeoutError(
|
||||||
|
f"Jupyter not available within {jupyter_timeout}s: {last_exc}"
|
||||||
|
)
|
||||||
|
|
||||||
|
def _jupyter_ws_url(self, kernel_id: str) -> str:
|
||||||
|
proxy = _build_proxy_url(self._client._base_url, self._id, 8888)
|
||||||
|
return f"{proxy}/api/kernels/{kernel_id}/channels"
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _jupyter_execute_request(code: str) -> dict:
|
||||||
|
msg_id = str(uuid.uuid4())
|
||||||
|
return {
|
||||||
|
"header": {
|
||||||
|
"msg_id": msg_id,
|
||||||
|
"msg_type": "execute_request",
|
||||||
|
"username": "wrenn-sdk",
|
||||||
|
"session": str(uuid.uuid4()),
|
||||||
|
"date": time.strftime("%Y-%m-%dT%H:%M:%S.000Z", time.gmtime()),
|
||||||
|
"version": "5.3",
|
||||||
|
},
|
||||||
|
"parent_header": {},
|
||||||
|
"metadata": {},
|
||||||
|
"content": {
|
||||||
|
"code": code,
|
||||||
|
"silent": False,
|
||||||
|
"store_history": True,
|
||||||
|
"user_expressions": {},
|
||||||
|
"allow_stdin": False,
|
||||||
|
"stop_on_error": True,
|
||||||
|
},
|
||||||
|
"buffers": [],
|
||||||
|
"channel": "shell",
|
||||||
|
"msg_id": msg_id,
|
||||||
|
"msg_type": "execute_request",
|
||||||
|
}
|
||||||
|
|
||||||
|
async def run_code(
|
||||||
|
self,
|
||||||
|
code: str,
|
||||||
|
language: str = "python",
|
||||||
|
timeout: float = 30,
|
||||||
|
jupyter_timeout: float = 30,
|
||||||
|
) -> CodeResult:
|
||||||
|
"""Execute code in a persistent Jupyter kernel (async)."""
|
||||||
|
kernel_id = await self._ensure_kernel(jupyter_timeout=jupyter_timeout)
|
||||||
|
ws_url = self._jupyter_ws_url(kernel_id)
|
||||||
|
|
||||||
|
msg = self._jupyter_execute_request(code)
|
||||||
|
msg_id = msg["msg_id"]
|
||||||
|
|
||||||
|
result = CodeResult()
|
||||||
|
deadline = time.monotonic() + timeout
|
||||||
|
headers = {"X-API-Key": self._client._api_key}
|
||||||
|
|
||||||
|
async with httpx_ws.aconnect_ws(ws_url, headers=headers) as ws:
|
||||||
|
await ws.send_text(json.dumps(msg))
|
||||||
|
while time.monotonic() < deadline:
|
||||||
|
time_left = deadline - time.monotonic()
|
||||||
|
if time_left <= 0:
|
||||||
|
break
|
||||||
|
try:
|
||||||
|
data = await asyncio.wait_for(
|
||||||
|
ws.receive_json(), timeout=time_left
|
||||||
|
)
|
||||||
|
except (asyncio.TimeoutError, Exception):
|
||||||
|
break
|
||||||
|
if not data:
|
||||||
|
break
|
||||||
|
parent = data.get("parent_header", {}).get("msg_id")
|
||||||
|
if parent != msg_id:
|
||||||
|
continue
|
||||||
|
msg_type = data.get("msg_type") or data.get("header", {}).get(
|
||||||
|
"msg_type"
|
||||||
|
)
|
||||||
|
content = data.get("content", {})
|
||||||
|
|
||||||
|
if msg_type == "stream":
|
||||||
|
name = content.get("name", "stdout")
|
||||||
|
if name == "stderr":
|
||||||
|
result.stderr += content.get("text", "")
|
||||||
|
else:
|
||||||
|
result.stdout += content.get("text", "")
|
||||||
|
elif msg_type == "execute_result":
|
||||||
|
bundle = content.get("data", {})
|
||||||
|
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
|
||||||
|
elif msg_type == "error":
|
||||||
|
traceback = content.get("traceback", [])
|
||||||
|
result.error = "\n".join(traceback)
|
||||||
|
elif msg_type == "status" and content.get("execution_state") == "idle":
|
||||||
|
break
|
||||||
|
|
||||||
|
if result.text is None and result.stdout:
|
||||||
|
result.text = result.stdout.strip()
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
async def __aexit__(self, *args) -> None:
|
||||||
|
if self._proxy_client is not None:
|
||||||
|
try:
|
||||||
|
await self._proxy_client.aclose()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
await super().__aexit__(*args)
|
||||||
264
src/wrenn/code_interpreter/capsule.py
Normal file
264
src/wrenn/code_interpreter/capsule.py
Normal file
@ -0,0 +1,264 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
import time
|
||||||
|
import uuid
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
import httpx_ws
|
||||||
|
|
||||||
|
from wrenn.capsule import Capsule as BaseCapsule
|
||||||
|
from wrenn.capsule import _build_proxy_url
|
||||||
|
|
||||||
|
|
||||||
|
DEFAULT_TEMPLATE = "code-runner-beta"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class CodeResult:
|
||||||
|
"""Result from stateful code execution.
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
text: text/plain representation of the result.
|
||||||
|
data: rich MIME bundle (e.g. ``{"image/png": "..."}``).
|
||||||
|
stdout: accumulated stdout output.
|
||||||
|
stderr: accumulated stderr output.
|
||||||
|
error: language-specific error/traceback string.
|
||||||
|
"""
|
||||||
|
|
||||||
|
text: str | None = None
|
||||||
|
data: dict[str, str] | None = None
|
||||||
|
stdout: str = ""
|
||||||
|
stderr: str = ""
|
||||||
|
error: str | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class Capsule(BaseCapsule):
|
||||||
|
"""Code interpreter capsule with ``run_code`` support.
|
||||||
|
|
||||||
|
Uses ``code-runner-beta`` template by default::
|
||||||
|
|
||||||
|
from wrenn.code_interpreter import Capsule
|
||||||
|
|
||||||
|
capsule = Capsule()
|
||||||
|
result = capsule.run_code("print('hello')")
|
||||||
|
print(result.stdout) # "hello\\n"
|
||||||
|
"""
|
||||||
|
|
||||||
|
_kernel_id: str | None
|
||||||
|
_proxy_client: httpx.Client | None
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
template: str | None = None,
|
||||||
|
vcpus: int | None = None,
|
||||||
|
memory_mb: int | None = None,
|
||||||
|
timeout: int | None = None,
|
||||||
|
*,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
**kwargs,
|
||||||
|
) -> None:
|
||||||
|
super().__init__(
|
||||||
|
template=template or DEFAULT_TEMPLATE,
|
||||||
|
vcpus=vcpus,
|
||||||
|
memory_mb=memory_mb,
|
||||||
|
timeout=timeout,
|
||||||
|
api_key=api_key,
|
||||||
|
base_url=base_url,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
self._kernel_id = None
|
||||||
|
self._proxy_client = None
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def create(
|
||||||
|
cls,
|
||||||
|
template: str | None = None,
|
||||||
|
vcpus: int | None = None,
|
||||||
|
memory_mb: int | None = None,
|
||||||
|
timeout: int | None = None,
|
||||||
|
*,
|
||||||
|
wait: bool = False,
|
||||||
|
api_key: str | None = None,
|
||||||
|
base_url: str | None = None,
|
||||||
|
) -> Capsule:
|
||||||
|
return cls(
|
||||||
|
template=template or DEFAULT_TEMPLATE,
|
||||||
|
vcpus=vcpus,
|
||||||
|
memory_mb=memory_mb,
|
||||||
|
timeout=timeout,
|
||||||
|
wait=wait,
|
||||||
|
api_key=api_key,
|
||||||
|
base_url=base_url,
|
||||||
|
)
|
||||||
|
|
||||||
|
def _get_proxy_client(self) -> httpx.Client:
|
||||||
|
if self._proxy_client is None:
|
||||||
|
url = (
|
||||||
|
_build_proxy_url(self._client._base_url, self._id, 8888)
|
||||||
|
.replace("ws://", "http://")
|
||||||
|
.replace("wss://", "https://")
|
||||||
|
)
|
||||||
|
self._proxy_client = httpx.Client(
|
||||||
|
base_url=url,
|
||||||
|
headers={"X-API-Key": self._client._api_key},
|
||||||
|
)
|
||||||
|
return self._proxy_client
|
||||||
|
|
||||||
|
def _ensure_kernel(self, jupyter_timeout: float = 30) -> str:
|
||||||
|
if self._kernel_id is not None:
|
||||||
|
return self._kernel_id
|
||||||
|
|
||||||
|
client = self._get_proxy_client()
|
||||||
|
deadline = time.monotonic() + jupyter_timeout
|
||||||
|
last_exc: Exception | None = None
|
||||||
|
|
||||||
|
while time.monotonic() < deadline:
|
||||||
|
try:
|
||||||
|
# Try to reuse an existing kernel
|
||||||
|
resp = client.get("/api/kernels")
|
||||||
|
if resp.status_code < 500:
|
||||||
|
resp.raise_for_status()
|
||||||
|
kernels = resp.json()
|
||||||
|
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(
|
||||||
|
f"Jupyter returned {resp.status_code}",
|
||||||
|
request=resp.request,
|
||||||
|
response=resp,
|
||||||
|
)
|
||||||
|
except httpx.HTTPStatusError:
|
||||||
|
raise
|
||||||
|
except Exception as exc:
|
||||||
|
last_exc = exc
|
||||||
|
time.sleep(0.5)
|
||||||
|
|
||||||
|
raise TimeoutError(
|
||||||
|
f"Jupyter not available within {jupyter_timeout}s: {last_exc}"
|
||||||
|
)
|
||||||
|
|
||||||
|
def _jupyter_ws_url(self, kernel_id: str) -> str:
|
||||||
|
proxy = _build_proxy_url(self._client._base_url, self._id, 8888)
|
||||||
|
return f"{proxy}/api/kernels/{kernel_id}/channels"
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _jupyter_execute_request(code: str) -> dict:
|
||||||
|
msg_id = str(uuid.uuid4())
|
||||||
|
return {
|
||||||
|
"header": {
|
||||||
|
"msg_id": msg_id,
|
||||||
|
"msg_type": "execute_request",
|
||||||
|
"username": "wrenn-sdk",
|
||||||
|
"session": str(uuid.uuid4()),
|
||||||
|
"date": time.strftime("%Y-%m-%dT%H:%M:%S.000Z", time.gmtime()),
|
||||||
|
"version": "5.3",
|
||||||
|
},
|
||||||
|
"parent_header": {},
|
||||||
|
"metadata": {},
|
||||||
|
"content": {
|
||||||
|
"code": code,
|
||||||
|
"silent": False,
|
||||||
|
"store_history": True,
|
||||||
|
"user_expressions": {},
|
||||||
|
"allow_stdin": False,
|
||||||
|
"stop_on_error": True,
|
||||||
|
},
|
||||||
|
"buffers": [],
|
||||||
|
"channel": "shell",
|
||||||
|
"msg_id": msg_id,
|
||||||
|
"msg_type": "execute_request",
|
||||||
|
}
|
||||||
|
|
||||||
|
def run_code(
|
||||||
|
self,
|
||||||
|
code: str,
|
||||||
|
language: str = "python",
|
||||||
|
timeout: float = 30,
|
||||||
|
jupyter_timeout: float = 30,
|
||||||
|
) -> CodeResult:
|
||||||
|
"""Execute code in a persistent Jupyter kernel.
|
||||||
|
|
||||||
|
Variables, imports, and function definitions survive across calls.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
code: Code string to execute.
|
||||||
|
language: Execution backend language. Currently only ``"python"``.
|
||||||
|
timeout: Maximum seconds to wait for execution to complete.
|
||||||
|
jupyter_timeout: Maximum seconds to wait for Jupyter to become available.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
A ``CodeResult`` with ``.text``, ``.data``, ``.stdout``, ``.stderr``, ``.error``.
|
||||||
|
"""
|
||||||
|
kernel_id = self._ensure_kernel(jupyter_timeout=jupyter_timeout)
|
||||||
|
ws_url = self._jupyter_ws_url(kernel_id)
|
||||||
|
|
||||||
|
msg = self._jupyter_execute_request(code)
|
||||||
|
msg_id = msg["msg_id"]
|
||||||
|
|
||||||
|
result = CodeResult()
|
||||||
|
deadline = time.monotonic() + timeout
|
||||||
|
headers = {"X-API-Key": self._client._api_key}
|
||||||
|
|
||||||
|
with httpx_ws.connect_ws(ws_url, headers=headers) as ws:
|
||||||
|
ws.send_text(json.dumps(msg))
|
||||||
|
while time.monotonic() < deadline:
|
||||||
|
time_left = deadline - time.monotonic()
|
||||||
|
if time_left <= 0:
|
||||||
|
break
|
||||||
|
try:
|
||||||
|
data = ws.receive_json(timeout=time_left)
|
||||||
|
except (TimeoutError, Exception):
|
||||||
|
break
|
||||||
|
if not data:
|
||||||
|
break
|
||||||
|
parent = data.get("parent_header", {}).get("msg_id")
|
||||||
|
if parent != msg_id:
|
||||||
|
continue
|
||||||
|
msg_type = data.get("msg_type") or data.get("header", {}).get(
|
||||||
|
"msg_type"
|
||||||
|
)
|
||||||
|
content = data.get("content", {})
|
||||||
|
|
||||||
|
if msg_type == "stream":
|
||||||
|
name = content.get("name", "stdout")
|
||||||
|
if name == "stderr":
|
||||||
|
result.stderr += content.get("text", "")
|
||||||
|
else:
|
||||||
|
result.stdout += content.get("text", "")
|
||||||
|
elif msg_type == "execute_result":
|
||||||
|
bundle = content.get("data", {})
|
||||||
|
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
|
||||||
|
elif msg_type == "error":
|
||||||
|
traceback = content.get("traceback", [])
|
||||||
|
result.error = "\n".join(traceback)
|
||||||
|
elif msg_type == "status" and content.get("execution_state") == "idle":
|
||||||
|
break
|
||||||
|
|
||||||
|
if result.text is None and result.stdout:
|
||||||
|
result.text = result.stdout.strip()
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
def __exit__(self, *args) -> None:
|
||||||
|
if self._proxy_client is not None:
|
||||||
|
try:
|
||||||
|
self._proxy_client.close()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
super().__exit__(*args)
|
||||||
366
src/wrenn/commands.py
Normal file
366
src/wrenn/commands.py
Normal file
@ -0,0 +1,366 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import json
|
||||||
|
from collections.abc import AsyncIterator, Iterator
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import overload, Literal
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
import httpx_ws
|
||||||
|
|
||||||
|
from wrenn.exceptions import handle_response
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class CommandResult:
|
||||||
|
"""Result from a foreground command execution."""
|
||||||
|
|
||||||
|
stdout: str
|
||||||
|
stderr: str
|
||||||
|
exit_code: int
|
||||||
|
duration_ms: int | None = None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class CommandHandle:
|
||||||
|
"""Handle for a background process."""
|
||||||
|
|
||||||
|
pid: int
|
||||||
|
tag: str
|
||||||
|
capsule_id: str
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class ProcessInfo:
|
||||||
|
"""Information about a running process."""
|
||||||
|
|
||||||
|
pid: int
|
||||||
|
tag: str | None = None
|
||||||
|
cmd: str | None = None
|
||||||
|
args: list[str] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class StreamEvent:
|
||||||
|
"""Base class for streaming exec events."""
|
||||||
|
|
||||||
|
__slots__ = ("type",)
|
||||||
|
|
||||||
|
def __init__(self, type: str) -> None:
|
||||||
|
self.type = type
|
||||||
|
|
||||||
|
|
||||||
|
class StreamStartEvent(StreamEvent):
|
||||||
|
__slots__ = ("pid",)
|
||||||
|
|
||||||
|
def __init__(self, pid: int) -> None:
|
||||||
|
super().__init__("start")
|
||||||
|
self.pid = pid
|
||||||
|
|
||||||
|
|
||||||
|
class StreamStdoutEvent(StreamEvent):
|
||||||
|
__slots__ = ("data",)
|
||||||
|
|
||||||
|
def __init__(self, data: str) -> None:
|
||||||
|
super().__init__("stdout")
|
||||||
|
self.data = data
|
||||||
|
|
||||||
|
|
||||||
|
class StreamStderrEvent(StreamEvent):
|
||||||
|
__slots__ = ("data",)
|
||||||
|
|
||||||
|
def __init__(self, data: str) -> None:
|
||||||
|
super().__init__("stderr")
|
||||||
|
self.data = data
|
||||||
|
|
||||||
|
|
||||||
|
class StreamExitEvent(StreamEvent):
|
||||||
|
__slots__ = ("exit_code",)
|
||||||
|
|
||||||
|
def __init__(self, exit_code: int) -> None:
|
||||||
|
super().__init__("exit")
|
||||||
|
self.exit_code = exit_code
|
||||||
|
|
||||||
|
|
||||||
|
class StreamErrorEvent(StreamEvent):
|
||||||
|
__slots__ = ("data",)
|
||||||
|
|
||||||
|
def __init__(self, data: str) -> None:
|
||||||
|
super().__init__("error")
|
||||||
|
self.data = data
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_stream_event(raw: dict) -> StreamEvent:
|
||||||
|
t = raw.get("type")
|
||||||
|
if t == "start":
|
||||||
|
return StreamStartEvent(pid=raw.get("pid", 0))
|
||||||
|
if t == "stdout":
|
||||||
|
return StreamStdoutEvent(data=raw.get("data", ""))
|
||||||
|
if t == "stderr":
|
||||||
|
return StreamStderrEvent(data=raw.get("data", ""))
|
||||||
|
if t == "exit":
|
||||||
|
return StreamExitEvent(exit_code=raw.get("exit_code", -1))
|
||||||
|
if t == "error":
|
||||||
|
return StreamErrorEvent(data=raw.get("data", ""))
|
||||||
|
return StreamEvent(type=t or "unknown")
|
||||||
|
|
||||||
|
|
||||||
|
def _decode_exec_response(data: dict) -> CommandResult:
|
||||||
|
stdout = data.get("stdout") or ""
|
||||||
|
stderr = data.get("stderr") or ""
|
||||||
|
if data.get("encoding") == "base64":
|
||||||
|
stdout = base64.b64decode(stdout).decode("utf-8", errors="replace")
|
||||||
|
if stderr:
|
||||||
|
stderr = base64.b64decode(stderr).decode("utf-8", errors="replace")
|
||||||
|
return CommandResult(
|
||||||
|
stdout=stdout,
|
||||||
|
stderr=stderr,
|
||||||
|
exit_code=data.get("exit_code", -1),
|
||||||
|
duration_ms=data.get("duration_ms"),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class Commands:
|
||||||
|
"""Sync command execution interface. Accessed via ``capsule.commands``."""
|
||||||
|
|
||||||
|
def __init__(self, capsule_id: str, http: httpx.Client) -> None:
|
||||||
|
self._capsule_id = capsule_id
|
||||||
|
self._http = http
|
||||||
|
|
||||||
|
@overload
|
||||||
|
def run(
|
||||||
|
self,
|
||||||
|
cmd: str,
|
||||||
|
*,
|
||||||
|
background: Literal[False] = ...,
|
||||||
|
timeout: int | None = 30,
|
||||||
|
envs: dict[str, str] | None = None,
|
||||||
|
cwd: str | None = None,
|
||||||
|
tag: str | None = None,
|
||||||
|
) -> CommandResult: ...
|
||||||
|
|
||||||
|
@overload
|
||||||
|
def run(
|
||||||
|
self,
|
||||||
|
cmd: str,
|
||||||
|
*,
|
||||||
|
background: Literal[True],
|
||||||
|
timeout: int | None = 30,
|
||||||
|
envs: dict[str, str] | None = None,
|
||||||
|
cwd: str | None = None,
|
||||||
|
tag: str | None = None,
|
||||||
|
) -> CommandHandle: ...
|
||||||
|
|
||||||
|
def run(
|
||||||
|
self,
|
||||||
|
cmd: str,
|
||||||
|
*,
|
||||||
|
background: bool = False,
|
||||||
|
timeout: int | None = 30,
|
||||||
|
envs: dict[str, str] | None = None,
|
||||||
|
cwd: str | None = None,
|
||||||
|
tag: str | None = None,
|
||||||
|
) -> CommandResult | CommandHandle:
|
||||||
|
payload: dict = {"cmd": cmd, "background": background}
|
||||||
|
if timeout is not None and not background:
|
||||||
|
payload["timeout_sec"] = timeout
|
||||||
|
if envs is not None:
|
||||||
|
payload["envs"] = envs
|
||||||
|
if cwd is not None:
|
||||||
|
payload["cwd"] = cwd
|
||||||
|
if tag is not None:
|
||||||
|
payload["tag"] = tag
|
||||||
|
|
||||||
|
resp = self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/exec", json=payload
|
||||||
|
)
|
||||||
|
data = handle_response(resp)
|
||||||
|
|
||||||
|
if background:
|
||||||
|
return CommandHandle(
|
||||||
|
pid=data.get("pid", 0),
|
||||||
|
tag=data.get("tag", ""),
|
||||||
|
capsule_id=self._capsule_id,
|
||||||
|
)
|
||||||
|
return _decode_exec_response(data)
|
||||||
|
|
||||||
|
def list(self) -> list[ProcessInfo]:
|
||||||
|
resp = self._http.get(f"/v1/capsules/{self._capsule_id}/processes")
|
||||||
|
data = handle_response(resp)
|
||||||
|
return [
|
||||||
|
ProcessInfo(
|
||||||
|
pid=p.get("pid", 0),
|
||||||
|
tag=p.get("tag"),
|
||||||
|
cmd=p.get("cmd"),
|
||||||
|
args=p.get("args"),
|
||||||
|
)
|
||||||
|
for p in data.get("processes", [])
|
||||||
|
]
|
||||||
|
|
||||||
|
def kill(self, pid: int) -> None:
|
||||||
|
resp = self._http.delete(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/processes/{pid}"
|
||||||
|
)
|
||||||
|
handle_response(resp)
|
||||||
|
|
||||||
|
def connect(self, pid: int) -> Iterator[StreamEvent]:
|
||||||
|
"""Connect to a running background process and stream its output."""
|
||||||
|
with httpx_ws.connect_ws(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/processes/{pid}/stream",
|
||||||
|
self._http,
|
||||||
|
) as ws:
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
raw = ws.receive_json()
|
||||||
|
event = _parse_stream_event(raw)
|
||||||
|
yield event
|
||||||
|
if event.type in ("exit", "error"):
|
||||||
|
break
|
||||||
|
except httpx_ws.WebSocketDisconnect:
|
||||||
|
break
|
||||||
|
|
||||||
|
def stream(
|
||||||
|
self, cmd: str, args: list[str] | None = None
|
||||||
|
) -> Iterator[StreamEvent]:
|
||||||
|
"""Execute a command via WebSocket, yielding ``StreamEvent`` objects."""
|
||||||
|
with httpx_ws.connect_ws(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/exec/stream",
|
||||||
|
self._http,
|
||||||
|
) as ws:
|
||||||
|
start_msg: dict = {"type": "start", "cmd": cmd}
|
||||||
|
if args:
|
||||||
|
start_msg["args"] = args
|
||||||
|
ws.send_text(json.dumps(start_msg))
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
raw = ws.receive_json()
|
||||||
|
event = _parse_stream_event(raw)
|
||||||
|
yield event
|
||||||
|
if event.type in ("exit", "error"):
|
||||||
|
break
|
||||||
|
except httpx_ws.WebSocketDisconnect:
|
||||||
|
break
|
||||||
|
|
||||||
|
|
||||||
|
class AsyncCommands:
|
||||||
|
"""Async command execution interface. Accessed via ``capsule.commands``."""
|
||||||
|
|
||||||
|
def __init__(self, capsule_id: str, http: httpx.AsyncClient) -> None:
|
||||||
|
self._capsule_id = capsule_id
|
||||||
|
self._http = http
|
||||||
|
|
||||||
|
@overload
|
||||||
|
async def run(
|
||||||
|
self,
|
||||||
|
cmd: str,
|
||||||
|
*,
|
||||||
|
background: Literal[False] = ...,
|
||||||
|
timeout: int | None = 30,
|
||||||
|
envs: dict[str, str] | None = None,
|
||||||
|
cwd: str | None = None,
|
||||||
|
tag: str | None = None,
|
||||||
|
) -> CommandResult: ...
|
||||||
|
|
||||||
|
@overload
|
||||||
|
async def run(
|
||||||
|
self,
|
||||||
|
cmd: str,
|
||||||
|
*,
|
||||||
|
background: Literal[True],
|
||||||
|
timeout: int | None = 30,
|
||||||
|
envs: dict[str, str] | None = None,
|
||||||
|
cwd: str | None = None,
|
||||||
|
tag: str | None = None,
|
||||||
|
) -> CommandHandle: ...
|
||||||
|
|
||||||
|
async def run(
|
||||||
|
self,
|
||||||
|
cmd: str,
|
||||||
|
*,
|
||||||
|
background: bool = False,
|
||||||
|
timeout: int | None = 30,
|
||||||
|
envs: dict[str, str] | None = None,
|
||||||
|
cwd: str | None = None,
|
||||||
|
tag: str | None = None,
|
||||||
|
) -> CommandResult | CommandHandle:
|
||||||
|
payload: dict = {"cmd": cmd, "background": background}
|
||||||
|
if timeout is not None and not background:
|
||||||
|
payload["timeout_sec"] = timeout
|
||||||
|
if envs is not None:
|
||||||
|
payload["envs"] = envs
|
||||||
|
if cwd is not None:
|
||||||
|
payload["cwd"] = cwd
|
||||||
|
if tag is not None:
|
||||||
|
payload["tag"] = tag
|
||||||
|
|
||||||
|
resp = await self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/exec", json=payload
|
||||||
|
)
|
||||||
|
data = handle_response(resp)
|
||||||
|
|
||||||
|
if background:
|
||||||
|
return CommandHandle(
|
||||||
|
pid=data.get("pid", 0),
|
||||||
|
tag=data.get("tag", ""),
|
||||||
|
capsule_id=self._capsule_id,
|
||||||
|
)
|
||||||
|
return _decode_exec_response(data)
|
||||||
|
|
||||||
|
async def list(self) -> list[ProcessInfo]:
|
||||||
|
resp = await self._http.get(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/processes"
|
||||||
|
)
|
||||||
|
data = handle_response(resp)
|
||||||
|
return [
|
||||||
|
ProcessInfo(
|
||||||
|
pid=p.get("pid", 0),
|
||||||
|
tag=p.get("tag"),
|
||||||
|
cmd=p.get("cmd"),
|
||||||
|
args=p.get("args"),
|
||||||
|
)
|
||||||
|
for p in data.get("processes", [])
|
||||||
|
]
|
||||||
|
|
||||||
|
async def kill(self, pid: int) -> None:
|
||||||
|
resp = await self._http.delete(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/processes/{pid}"
|
||||||
|
)
|
||||||
|
handle_response(resp)
|
||||||
|
|
||||||
|
async def connect(self, pid: int) -> AsyncIterator[StreamEvent]:
|
||||||
|
"""Connect to a running background process and stream its output."""
|
||||||
|
async with httpx_ws.aconnect_ws(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/processes/{pid}/stream",
|
||||||
|
self._http,
|
||||||
|
) as ws:
|
||||||
|
try:
|
||||||
|
while True:
|
||||||
|
raw = await ws.receive_json()
|
||||||
|
event = _parse_stream_event(raw)
|
||||||
|
yield event
|
||||||
|
if event.type in ("exit", "error"):
|
||||||
|
break
|
||||||
|
except httpx_ws.WebSocketDisconnect:
|
||||||
|
pass
|
||||||
|
|
||||||
|
async def stream(
|
||||||
|
self, cmd: str, args: list[str] | None = None
|
||||||
|
) -> AsyncIterator[StreamEvent]:
|
||||||
|
"""Execute a command via WebSocket, yielding ``StreamEvent`` objects."""
|
||||||
|
async with httpx_ws.aconnect_ws(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/exec/stream",
|
||||||
|
self._http,
|
||||||
|
) as ws:
|
||||||
|
start_msg: dict = {"type": "start", "cmd": cmd}
|
||||||
|
if args:
|
||||||
|
start_msg["args"] = args
|
||||||
|
await ws.send_text(json.dumps(start_msg))
|
||||||
|
try:
|
||||||
|
while True:
|
||||||
|
raw = await ws.receive_json()
|
||||||
|
event = _parse_stream_event(raw)
|
||||||
|
yield event
|
||||||
|
if event.type in ("exit", "error"):
|
||||||
|
break
|
||||||
|
except httpx_ws.WebSocketDisconnect:
|
||||||
|
pass
|
||||||
126
src/wrenn/exceptions.py
Normal file
126
src/wrenn/exceptions.py
Normal file
@ -0,0 +1,126 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import warnings
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
|
||||||
|
|
||||||
|
class WrennError(Exception):
|
||||||
|
"""Base exception for all Wrenn SDK errors."""
|
||||||
|
|
||||||
|
def __init__(self, code: str, message: str, status_code: int) -> None:
|
||||||
|
self.code = code
|
||||||
|
self.message = message
|
||||||
|
self.status_code = status_code
|
||||||
|
super().__init__(message)
|
||||||
|
|
||||||
|
|
||||||
|
class WrennValidationError(WrennError):
|
||||||
|
"""400 — Invalid request parameters."""
|
||||||
|
|
||||||
|
|
||||||
|
class WrennAuthenticationError(WrennError):
|
||||||
|
"""401 — Invalid or missing authentication."""
|
||||||
|
|
||||||
|
|
||||||
|
class WrennForbiddenError(WrennError):
|
||||||
|
"""403 — Authenticated but not authorized."""
|
||||||
|
|
||||||
|
|
||||||
|
class WrennNotFoundError(WrennError):
|
||||||
|
"""404 — Resource not found."""
|
||||||
|
|
||||||
|
|
||||||
|
class WrennConflictError(WrennError):
|
||||||
|
"""409 — State conflict (e.g. invalid_state)."""
|
||||||
|
|
||||||
|
|
||||||
|
class WrennHostHasCapsulesError(WrennConflictError):
|
||||||
|
"""409 — Host still has running capsules."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self, code: str, message: str, status_code: int, capsule_ids: list[str]
|
||||||
|
) -> None:
|
||||||
|
self.capsule_ids = capsule_ids
|
||||||
|
super().__init__(code, message, status_code)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def sandbox_ids(self) -> list[str]:
|
||||||
|
warnings.warn(
|
||||||
|
"'sandbox_ids' is deprecated, use 'capsule_ids' instead",
|
||||||
|
DeprecationWarning,
|
||||||
|
stacklevel=2,
|
||||||
|
)
|
||||||
|
return self.capsule_ids
|
||||||
|
|
||||||
|
|
||||||
|
class WrennHostUnavailableError(WrennError):
|
||||||
|
"""503 — No suitable host available."""
|
||||||
|
|
||||||
|
|
||||||
|
class WrennAgentError(WrennError):
|
||||||
|
"""502 — Host agent returned an error."""
|
||||||
|
|
||||||
|
|
||||||
|
class WrennInternalError(WrennError):
|
||||||
|
"""500 — Unexpected server error."""
|
||||||
|
|
||||||
|
|
||||||
|
_ERROR_MAP: dict[str, type[WrennError]] = {
|
||||||
|
"invalid_request": WrennValidationError,
|
||||||
|
"unauthorized": WrennAuthenticationError,
|
||||||
|
"forbidden": WrennForbiddenError,
|
||||||
|
"not_found": WrennNotFoundError,
|
||||||
|
"invalid_state": WrennConflictError,
|
||||||
|
"conflict": WrennConflictError,
|
||||||
|
"host_has_sandboxes": WrennHostHasCapsulesError,
|
||||||
|
"host_has_capsules": WrennHostHasCapsulesError,
|
||||||
|
"host_unavailable": WrennHostUnavailableError,
|
||||||
|
"agent_error": WrennAgentError,
|
||||||
|
"internal_error": WrennInternalError,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def handle_response(resp: httpx.Response) -> dict | list:
|
||||||
|
if resp.status_code >= 400:
|
||||||
|
try:
|
||||||
|
body = resp.json()
|
||||||
|
except Exception:
|
||||||
|
resp.raise_for_status()
|
||||||
|
raise
|
||||||
|
|
||||||
|
err = body.get("error", {})
|
||||||
|
code = err.get("code", "internal_error")
|
||||||
|
message = err.get("message", resp.text)
|
||||||
|
|
||||||
|
exc_cls = _ERROR_MAP.get(code, WrennError)
|
||||||
|
|
||||||
|
if exc_cls is WrennHostHasCapsulesError:
|
||||||
|
raise WrennHostHasCapsulesError(
|
||||||
|
code=code,
|
||||||
|
message=message,
|
||||||
|
status_code=resp.status_code,
|
||||||
|
capsule_ids=body.get("sandbox_ids", []),
|
||||||
|
)
|
||||||
|
|
||||||
|
raise exc_cls(
|
||||||
|
code=code,
|
||||||
|
message=message,
|
||||||
|
status_code=resp.status_code,
|
||||||
|
)
|
||||||
|
|
||||||
|
if resp.status_code == 204:
|
||||||
|
return {}
|
||||||
|
|
||||||
|
return resp.json()
|
||||||
|
|
||||||
|
|
||||||
|
def __getattr__(name: str) -> type:
|
||||||
|
if name == "WrennHostHasSandboxesError":
|
||||||
|
warnings.warn(
|
||||||
|
"'WrennHostHasSandboxesError' is deprecated, use 'WrennHostHasCapsulesError' instead",
|
||||||
|
DeprecationWarning,
|
||||||
|
stacklevel=2,
|
||||||
|
)
|
||||||
|
return WrennHostHasCapsulesError
|
||||||
|
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||||
241
src/wrenn/files.py
Normal file
241
src/wrenn/files.py
Normal file
@ -0,0 +1,241 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
from collections.abc import AsyncIterator, Iterator
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
|
||||||
|
from wrenn.exceptions import WrennNotFoundError, handle_response
|
||||||
|
from wrenn.models import FileEntry, ListDirResponse, MakeDirResponse
|
||||||
|
|
||||||
|
|
||||||
|
class Files:
|
||||||
|
"""Sync filesystem interface. Accessed via ``capsule.files``."""
|
||||||
|
|
||||||
|
def __init__(self, capsule_id: str, http: httpx.Client) -> None:
|
||||||
|
self._capsule_id = capsule_id
|
||||||
|
self._http = http
|
||||||
|
|
||||||
|
def read(self, path: str) -> str:
|
||||||
|
"""Read a file as a UTF-8 string."""
|
||||||
|
return self.read_bytes(path).decode("utf-8", errors="replace")
|
||||||
|
|
||||||
|
def read_bytes(self, path: str) -> bytes:
|
||||||
|
"""Read a file as raw bytes."""
|
||||||
|
resp = self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/read",
|
||||||
|
json={"path": path},
|
||||||
|
)
|
||||||
|
resp.raise_for_status()
|
||||||
|
return resp.content
|
||||||
|
|
||||||
|
def write(self, path: str, data: str | bytes) -> None:
|
||||||
|
"""Write data to a file inside the capsule."""
|
||||||
|
if isinstance(data, str):
|
||||||
|
data = data.encode("utf-8")
|
||||||
|
resp = self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/write",
|
||||||
|
files={"file": ("upload", data)},
|
||||||
|
data={"path": path},
|
||||||
|
)
|
||||||
|
resp.raise_for_status()
|
||||||
|
|
||||||
|
def list(self, path: str, depth: int = 1) -> list[FileEntry]:
|
||||||
|
"""List directory contents."""
|
||||||
|
resp = self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/list",
|
||||||
|
json={"path": path, "depth": depth},
|
||||||
|
)
|
||||||
|
parsed = ListDirResponse.model_validate(handle_response(resp))
|
||||||
|
return parsed.entries or []
|
||||||
|
|
||||||
|
def exists(self, path: str) -> bool:
|
||||||
|
"""Check whether a path exists inside the capsule."""
|
||||||
|
parent = os.path.dirname(path)
|
||||||
|
name = os.path.basename(path)
|
||||||
|
try:
|
||||||
|
entries = self.list(parent, depth=1)
|
||||||
|
except WrennNotFoundError:
|
||||||
|
return False
|
||||||
|
return any(e.name == name for e in entries)
|
||||||
|
|
||||||
|
def make_dir(self, path: str) -> FileEntry:
|
||||||
|
"""Create a directory (with parents). Idempotent."""
|
||||||
|
resp = self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/mkdir",
|
||||||
|
json={"path": path},
|
||||||
|
)
|
||||||
|
if resp.status_code == 409:
|
||||||
|
try:
|
||||||
|
body = resp.json()
|
||||||
|
if body.get("error", {}).get("code") == "conflict":
|
||||||
|
parent = os.path.dirname(path)
|
||||||
|
name = os.path.basename(path)
|
||||||
|
for entry in self.list(parent, depth=1):
|
||||||
|
if entry.name == name:
|
||||||
|
return entry
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
parsed = MakeDirResponse.model_validate(handle_response(resp))
|
||||||
|
if parsed.entry is None:
|
||||||
|
raise RuntimeError("mkdir response missing entry")
|
||||||
|
return parsed.entry
|
||||||
|
|
||||||
|
def remove(self, path: str) -> None:
|
||||||
|
"""Remove a file or directory recursively."""
|
||||||
|
resp = self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/remove",
|
||||||
|
json={"path": path},
|
||||||
|
)
|
||||||
|
handle_response(resp)
|
||||||
|
|
||||||
|
def upload_stream(self, path: str, stream: Iterator[bytes]) -> None:
|
||||||
|
"""Streaming upload for large files."""
|
||||||
|
boundary = os.urandom(16).hex().encode("utf-8")
|
||||||
|
|
||||||
|
def _multipart() -> Iterator[bytes]:
|
||||||
|
yield b"--" + boundary + b"\r\n"
|
||||||
|
yield b'Content-Disposition: form-data; name="path"\r\n\r\n'
|
||||||
|
yield path.encode("utf-8") + b"\r\n"
|
||||||
|
yield b"--" + boundary + b"\r\n"
|
||||||
|
yield b'Content-Disposition: form-data; name="file"; filename="upload.bin"\r\n'
|
||||||
|
yield b"Content-Type: application/octet-stream\r\n\r\n"
|
||||||
|
for chunk in stream:
|
||||||
|
yield chunk if isinstance(chunk, bytes) else chunk.encode("utf-8")
|
||||||
|
yield b"\r\n--" + boundary + b"--\r\n"
|
||||||
|
|
||||||
|
resp = self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/stream/write",
|
||||||
|
content=_multipart(),
|
||||||
|
headers={
|
||||||
|
"Content-Type": f"multipart/form-data; boundary={boundary.decode('utf-8')}"
|
||||||
|
},
|
||||||
|
)
|
||||||
|
resp.raise_for_status()
|
||||||
|
|
||||||
|
def download_stream(self, path: str) -> Iterator[bytes]:
|
||||||
|
"""Streaming download for large files."""
|
||||||
|
with self._http.stream(
|
||||||
|
"POST",
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/stream/read",
|
||||||
|
json={"path": path},
|
||||||
|
) as resp:
|
||||||
|
resp.raise_for_status()
|
||||||
|
yield from resp.iter_bytes()
|
||||||
|
|
||||||
|
|
||||||
|
class AsyncFiles:
|
||||||
|
"""Async filesystem interface. Accessed via ``capsule.files``."""
|
||||||
|
|
||||||
|
def __init__(self, capsule_id: str, http: httpx.AsyncClient) -> None:
|
||||||
|
self._capsule_id = capsule_id
|
||||||
|
self._http = http
|
||||||
|
|
||||||
|
async def read(self, path: str) -> str:
|
||||||
|
"""Read a file as a UTF-8 string."""
|
||||||
|
data = await self.read_bytes(path)
|
||||||
|
return data.decode("utf-8", errors="replace")
|
||||||
|
|
||||||
|
async def read_bytes(self, path: str) -> bytes:
|
||||||
|
"""Read a file as raw bytes."""
|
||||||
|
resp = await self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/read",
|
||||||
|
json={"path": path},
|
||||||
|
)
|
||||||
|
resp.raise_for_status()
|
||||||
|
return resp.content
|
||||||
|
|
||||||
|
async def write(self, path: str, data: str | bytes) -> None:
|
||||||
|
"""Write data to a file inside the capsule."""
|
||||||
|
if isinstance(data, str):
|
||||||
|
data = data.encode("utf-8")
|
||||||
|
resp = await self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/write",
|
||||||
|
files={"file": ("upload", data)},
|
||||||
|
data={"path": path},
|
||||||
|
)
|
||||||
|
resp.raise_for_status()
|
||||||
|
|
||||||
|
async def list(self, path: str, depth: int = 1) -> list[FileEntry]:
|
||||||
|
"""List directory contents."""
|
||||||
|
resp = await self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/list",
|
||||||
|
json={"path": path, "depth": depth},
|
||||||
|
)
|
||||||
|
parsed = ListDirResponse.model_validate(handle_response(resp))
|
||||||
|
return parsed.entries or []
|
||||||
|
|
||||||
|
async def exists(self, path: str) -> bool:
|
||||||
|
"""Check whether a path exists inside the capsule."""
|
||||||
|
parent = os.path.dirname(path)
|
||||||
|
name = os.path.basename(path)
|
||||||
|
try:
|
||||||
|
entries = await self.list(parent, depth=1)
|
||||||
|
except WrennNotFoundError:
|
||||||
|
return False
|
||||||
|
return any(e.name == name for e in entries)
|
||||||
|
|
||||||
|
async def make_dir(self, path: str) -> FileEntry:
|
||||||
|
"""Create a directory (with parents). Idempotent."""
|
||||||
|
resp = await self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/mkdir",
|
||||||
|
json={"path": path},
|
||||||
|
)
|
||||||
|
if resp.status_code == 409:
|
||||||
|
try:
|
||||||
|
body = resp.json()
|
||||||
|
if body.get("error", {}).get("code") == "conflict":
|
||||||
|
parent = os.path.dirname(path)
|
||||||
|
name = os.path.basename(path)
|
||||||
|
for entry in await self.list(parent, depth=1):
|
||||||
|
if entry.name == name:
|
||||||
|
return entry
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
parsed = MakeDirResponse.model_validate(handle_response(resp))
|
||||||
|
if parsed.entry is None:
|
||||||
|
raise RuntimeError("mkdir response missing entry")
|
||||||
|
return parsed.entry
|
||||||
|
|
||||||
|
async def remove(self, path: str) -> None:
|
||||||
|
"""Remove a file or directory recursively."""
|
||||||
|
resp = await self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/remove",
|
||||||
|
json={"path": path},
|
||||||
|
)
|
||||||
|
handle_response(resp)
|
||||||
|
|
||||||
|
async def upload_stream(self, path: str, stream: AsyncIterator[bytes]) -> None:
|
||||||
|
"""Streaming upload for large files."""
|
||||||
|
boundary = os.urandom(16).hex().encode("utf-8")
|
||||||
|
|
||||||
|
async def _multipart() -> AsyncIterator[bytes]:
|
||||||
|
yield b"--" + boundary + b"\r\n"
|
||||||
|
yield b'Content-Disposition: form-data; name="path"\r\n\r\n'
|
||||||
|
yield path.encode("utf-8") + b"\r\n"
|
||||||
|
yield b"--" + boundary + b"\r\n"
|
||||||
|
yield b'Content-Disposition: form-data; name="file"; filename="upload.bin"\r\n'
|
||||||
|
yield b"Content-Type: application/octet-stream\r\n\r\n"
|
||||||
|
async for chunk in stream:
|
||||||
|
yield chunk if isinstance(chunk, bytes) else chunk.encode("utf-8")
|
||||||
|
yield b"\r\n--" + boundary + b"--\r\n"
|
||||||
|
|
||||||
|
resp = await self._http.post(
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/stream/write",
|
||||||
|
content=_multipart(),
|
||||||
|
headers={
|
||||||
|
"Content-Type": f"multipart/form-data; boundary={boundary.decode('utf-8')}"
|
||||||
|
},
|
||||||
|
)
|
||||||
|
resp.raise_for_status()
|
||||||
|
|
||||||
|
async def download_stream(self, path: str) -> AsyncIterator[bytes]:
|
||||||
|
"""Streaming download for large files."""
|
||||||
|
async with self._http.stream(
|
||||||
|
"POST",
|
||||||
|
f"/v1/capsules/{self._capsule_id}/files/stream/read",
|
||||||
|
json={"path": path},
|
||||||
|
) as resp:
|
||||||
|
resp.raise_for_status()
|
||||||
|
async for chunk in resp.aiter_bytes():
|
||||||
|
yield chunk
|
||||||
67
src/wrenn/models/__init__.py
Normal file
67
src/wrenn/models/__init__.py
Normal file
@ -0,0 +1,67 @@
|
|||||||
|
from wrenn.models._generated import (
|
||||||
|
APIKeyResponse,
|
||||||
|
AuthResponse,
|
||||||
|
Capsule,
|
||||||
|
CreateAPIKeyRequest,
|
||||||
|
CreateCapsuleRequest,
|
||||||
|
CreateHostRequest,
|
||||||
|
CreateHostResponse,
|
||||||
|
CreateSnapshotRequest,
|
||||||
|
Encoding,
|
||||||
|
Error,
|
||||||
|
Error1,
|
||||||
|
ExecRequest,
|
||||||
|
ExecResponse,
|
||||||
|
FileEntry,
|
||||||
|
Host,
|
||||||
|
ListDirRequest,
|
||||||
|
ListDirResponse,
|
||||||
|
LoginRequest,
|
||||||
|
MakeDirRequest,
|
||||||
|
MakeDirResponse,
|
||||||
|
ReadFileRequest,
|
||||||
|
RegisterHostRequest,
|
||||||
|
RegisterHostResponse,
|
||||||
|
RemoveRequest,
|
||||||
|
SignupRequest,
|
||||||
|
Status,
|
||||||
|
Status1,
|
||||||
|
Template,
|
||||||
|
Type,
|
||||||
|
Type1,
|
||||||
|
Type2,
|
||||||
|
)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"APIKeyResponse",
|
||||||
|
"AuthResponse",
|
||||||
|
"CreateAPIKeyRequest",
|
||||||
|
"CreateHostRequest",
|
||||||
|
"CreateHostResponse",
|
||||||
|
"CreateCapsuleRequest",
|
||||||
|
"CreateSnapshotRequest",
|
||||||
|
"Encoding",
|
||||||
|
"Error",
|
||||||
|
"Error1",
|
||||||
|
"ExecRequest",
|
||||||
|
"ExecResponse",
|
||||||
|
"FileEntry",
|
||||||
|
"Host",
|
||||||
|
"ListDirRequest",
|
||||||
|
"ListDirResponse",
|
||||||
|
"LoginRequest",
|
||||||
|
"MakeDirRequest",
|
||||||
|
"MakeDirResponse",
|
||||||
|
"ReadFileRequest",
|
||||||
|
"RegisterHostRequest",
|
||||||
|
"RegisterHostResponse",
|
||||||
|
"RemoveRequest",
|
||||||
|
"Capsule",
|
||||||
|
"SignupRequest",
|
||||||
|
"Status",
|
||||||
|
"Status1",
|
||||||
|
"Template",
|
||||||
|
"Type",
|
||||||
|
"Type1",
|
||||||
|
"Type2",
|
||||||
|
]
|
||||||
577
src/wrenn/models/_generated.py
Normal file
577
src/wrenn/models/_generated.py
Normal file
@ -0,0 +1,577 @@
|
|||||||
|
# generated by datamodel-codegen:
|
||||||
|
# filename: openapi.yaml
|
||||||
|
# timestamp: 2026-04-15T08:37:41+00:00
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
from pydantic import AwareDatetime, BaseModel, EmailStr, Field
|
||||||
|
from typing import Annotated
|
||||||
|
from enum import StrEnum
|
||||||
|
|
||||||
|
|
||||||
|
class SignupRequest(BaseModel):
|
||||||
|
email: EmailStr
|
||||||
|
password: Annotated[str, Field(min_length=8)]
|
||||||
|
name: Annotated[str, Field(max_length=100)]
|
||||||
|
|
||||||
|
|
||||||
|
class LoginRequest(BaseModel):
|
||||||
|
email: EmailStr
|
||||||
|
password: str
|
||||||
|
|
||||||
|
|
||||||
|
class AuthResponse(BaseModel):
|
||||||
|
token: Annotated[str | None, Field(description="JWT token (valid for 6 hours)")] = (
|
||||||
|
None
|
||||||
|
)
|
||||||
|
user_id: str | None = None
|
||||||
|
team_id: str | None = None
|
||||||
|
email: str | None = None
|
||||||
|
name: str | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class CreateAPIKeyRequest(BaseModel):
|
||||||
|
name: str | None = "Unnamed API Key"
|
||||||
|
|
||||||
|
|
||||||
|
class APIKeyResponse(BaseModel):
|
||||||
|
id: str | None = None
|
||||||
|
team_id: str | None = None
|
||||||
|
name: str | None = None
|
||||||
|
key_prefix: Annotated[
|
||||||
|
str | None, Field(description='Display prefix (e.g. "wrn_ab12cd34...")')
|
||||||
|
] = None
|
||||||
|
created_at: AwareDatetime | None = None
|
||||||
|
last_used: AwareDatetime | None = None
|
||||||
|
key: Annotated[
|
||||||
|
str | None,
|
||||||
|
Field(
|
||||||
|
description="Full plaintext key. Only returned on creation, never again."
|
||||||
|
),
|
||||||
|
] = None
|
||||||
|
|
||||||
|
|
||||||
|
class CreateCapsuleRequest(BaseModel):
|
||||||
|
template: str | None = "minimal"
|
||||||
|
vcpus: int | None = 1
|
||||||
|
memory_mb: int | None = 512
|
||||||
|
timeout_sec: Annotated[
|
||||||
|
int | None,
|
||||||
|
Field(
|
||||||
|
description="Auto-pause TTL in seconds. The capsule is automatically paused after this duration of inactivity (no exec or ping). 0 means no auto-pause.\n"
|
||||||
|
),
|
||||||
|
] = 0
|
||||||
|
|
||||||
|
|
||||||
|
class Range(StrEnum):
|
||||||
|
field_5m = "5m"
|
||||||
|
field_1h = "1h"
|
||||||
|
field_6h = "6h"
|
||||||
|
field_24h = "24h"
|
||||||
|
field_30d = "30d"
|
||||||
|
|
||||||
|
|
||||||
|
class Current(BaseModel):
|
||||||
|
running_count: int | None = None
|
||||||
|
vcpus_reserved: int | None = None
|
||||||
|
memory_mb_reserved: int | None = None
|
||||||
|
sampled_at: AwareDatetime | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class Peaks(BaseModel):
|
||||||
|
"""
|
||||||
|
Maximum values over the last 30 days.
|
||||||
|
"""
|
||||||
|
|
||||||
|
running_count: int | None = None
|
||||||
|
vcpus: int | None = None
|
||||||
|
memory_mb: int | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class Series(BaseModel):
|
||||||
|
"""
|
||||||
|
Parallel arrays for chart rendering.
|
||||||
|
"""
|
||||||
|
|
||||||
|
labels: list[AwareDatetime] | None = None
|
||||||
|
running: list[int] | None = None
|
||||||
|
vcpus: list[int] | None = None
|
||||||
|
memory_mb: list[int] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class CapsuleStats(BaseModel):
|
||||||
|
range: Range | None = None
|
||||||
|
current: Current | None = None
|
||||||
|
peaks: Annotated[
|
||||||
|
Peaks | None, Field(description="Maximum values over the last 30 days.")
|
||||||
|
] = None
|
||||||
|
series: Annotated[
|
||||||
|
Series | None, Field(description="Parallel arrays for chart rendering.")
|
||||||
|
] = None
|
||||||
|
|
||||||
|
|
||||||
|
class Status(StrEnum):
|
||||||
|
pending = "pending"
|
||||||
|
starting = "starting"
|
||||||
|
running = "running"
|
||||||
|
paused = "paused"
|
||||||
|
hibernated = "hibernated"
|
||||||
|
stopped = "stopped"
|
||||||
|
missing = "missing"
|
||||||
|
error = "error"
|
||||||
|
|
||||||
|
|
||||||
|
class Capsule(BaseModel):
|
||||||
|
id: str | None = None
|
||||||
|
status: Status | None = None
|
||||||
|
template: str | None = None
|
||||||
|
vcpus: int | None = None
|
||||||
|
memory_mb: int | None = None
|
||||||
|
timeout_sec: int | None = None
|
||||||
|
guest_ip: str | None = None
|
||||||
|
host_ip: str | None = None
|
||||||
|
created_at: AwareDatetime | None = None
|
||||||
|
started_at: AwareDatetime | None = None
|
||||||
|
last_active_at: AwareDatetime | None = None
|
||||||
|
last_updated: AwareDatetime | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class CreateSnapshotRequest(BaseModel):
|
||||||
|
sandbox_id: Annotated[
|
||||||
|
str, Field(description="ID of the running capsule to snapshot.")
|
||||||
|
]
|
||||||
|
name: Annotated[
|
||||||
|
str | None,
|
||||||
|
Field(description="Name for the snapshot template. Auto-generated if omitted."),
|
||||||
|
] = None
|
||||||
|
|
||||||
|
|
||||||
|
class Type(StrEnum):
|
||||||
|
base = "base"
|
||||||
|
snapshot = "snapshot"
|
||||||
|
|
||||||
|
|
||||||
|
class Template(BaseModel):
|
||||||
|
name: str | None = None
|
||||||
|
type: Type | None = None
|
||||||
|
vcpus: int | None = None
|
||||||
|
memory_mb: int | None = None
|
||||||
|
size_bytes: int | None = None
|
||||||
|
created_at: AwareDatetime | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class ExecRequest(BaseModel):
|
||||||
|
cmd: str
|
||||||
|
args: list[str] | None = None
|
||||||
|
timeout_sec: Annotated[
|
||||||
|
int | None,
|
||||||
|
Field(description="Timeout in seconds (foreground exec only, default 30)"),
|
||||||
|
] = 30
|
||||||
|
background: Annotated[
|
||||||
|
bool | None,
|
||||||
|
Field(
|
||||||
|
description="If true, starts the process in the background and returns immediately with a PID and tag (HTTP 202)"
|
||||||
|
),
|
||||||
|
] = False
|
||||||
|
tag: Annotated[
|
||||||
|
str | None,
|
||||||
|
Field(
|
||||||
|
description="Optional user-chosen tag for the background process. Auto-generated if omitted. Only used when background is true."
|
||||||
|
),
|
||||||
|
] = None
|
||||||
|
envs: Annotated[
|
||||||
|
dict[str, str] | None,
|
||||||
|
Field(
|
||||||
|
description="Environment variables for the process (background exec only)"
|
||||||
|
),
|
||||||
|
] = None
|
||||||
|
cwd: Annotated[
|
||||||
|
str | None,
|
||||||
|
Field(description="Working directory for the process (background exec only)"),
|
||||||
|
] = None
|
||||||
|
|
||||||
|
|
||||||
|
class BackgroundExecResponse(BaseModel):
|
||||||
|
sandbox_id: str | None = None
|
||||||
|
cmd: str | None = None
|
||||||
|
pid: int | None = None
|
||||||
|
tag: str | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class ProcessEntry(BaseModel):
|
||||||
|
pid: int | None = None
|
||||||
|
tag: str | None = None
|
||||||
|
cmd: str | None = None
|
||||||
|
args: list[str] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class ProcessListResponse(BaseModel):
|
||||||
|
processes: list[ProcessEntry] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class Encoding(StrEnum):
|
||||||
|
"""
|
||||||
|
Output encoding. "base64" when stdout/stderr contain binary data.
|
||||||
|
"""
|
||||||
|
|
||||||
|
utf_8 = "utf-8"
|
||||||
|
base64 = "base64"
|
||||||
|
|
||||||
|
|
||||||
|
class ExecResponse(BaseModel):
|
||||||
|
sandbox_id: str | None = None
|
||||||
|
cmd: str | None = None
|
||||||
|
stdout: str | None = None
|
||||||
|
stderr: str | None = None
|
||||||
|
exit_code: int | None = None
|
||||||
|
duration_ms: int | None = None
|
||||||
|
encoding: Annotated[
|
||||||
|
Encoding | None,
|
||||||
|
Field(
|
||||||
|
description='Output encoding. "base64" when stdout/stderr contain binary data.'
|
||||||
|
),
|
||||||
|
] = None
|
||||||
|
|
||||||
|
|
||||||
|
class ReadFileRequest(BaseModel):
|
||||||
|
path: Annotated[str, Field(description="Absolute file path inside the capsule")]
|
||||||
|
|
||||||
|
|
||||||
|
class ListDirRequest(BaseModel):
|
||||||
|
path: Annotated[str, Field(description="Directory path inside the capsule")]
|
||||||
|
depth: Annotated[
|
||||||
|
int | None,
|
||||||
|
Field(
|
||||||
|
description="Recursion depth (0 = non-recursive, 1 = immediate children)"
|
||||||
|
),
|
||||||
|
] = 1
|
||||||
|
|
||||||
|
|
||||||
|
class Type1(StrEnum):
|
||||||
|
file = "file"
|
||||||
|
directory = "directory"
|
||||||
|
symlink = "symlink"
|
||||||
|
|
||||||
|
|
||||||
|
class FileEntry(BaseModel):
|
||||||
|
name: str | None = None
|
||||||
|
path: str | None = None
|
||||||
|
type: Type1 | None = None
|
||||||
|
size: int | None = None
|
||||||
|
mode: int | None = None
|
||||||
|
permissions: Annotated[
|
||||||
|
str | None, Field(description='Human-readable permissions (e.g. "-rwxr-xr-x")')
|
||||||
|
] = None
|
||||||
|
owner: str | None = None
|
||||||
|
group: str | None = None
|
||||||
|
modified_at: Annotated[
|
||||||
|
int | None, Field(description="Unix timestamp (seconds)")
|
||||||
|
] = None
|
||||||
|
symlink_target: str | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class MakeDirRequest(BaseModel):
|
||||||
|
path: Annotated[
|
||||||
|
str, Field(description="Directory path to create inside the capsule")
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class MakeDirResponse(BaseModel):
|
||||||
|
entry: FileEntry | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class RemoveRequest(BaseModel):
|
||||||
|
path: Annotated[str, Field(description="Path to remove inside the capsule")]
|
||||||
|
|
||||||
|
|
||||||
|
class Type2(StrEnum):
|
||||||
|
"""
|
||||||
|
Host type. Regular hosts are shared; BYOC hosts belong to a team.
|
||||||
|
"""
|
||||||
|
|
||||||
|
regular = "regular"
|
||||||
|
byoc = "byoc"
|
||||||
|
|
||||||
|
|
||||||
|
class CreateHostRequest(BaseModel):
|
||||||
|
type: Annotated[
|
||||||
|
Type2,
|
||||||
|
Field(
|
||||||
|
description="Host type. Regular hosts are shared; BYOC hosts belong to a team."
|
||||||
|
),
|
||||||
|
]
|
||||||
|
team_id: Annotated[str | None, Field(description="Required for BYOC hosts.")] = None
|
||||||
|
provider: Annotated[
|
||||||
|
str | None,
|
||||||
|
Field(description="Cloud provider (e.g. aws, gcp, hetzner, bare-metal)."),
|
||||||
|
] = None
|
||||||
|
availability_zone: Annotated[
|
||||||
|
str | None, Field(description="Availability zone (e.g. us-east, eu-west).")
|
||||||
|
] = None
|
||||||
|
|
||||||
|
|
||||||
|
class RegisterHostRequest(BaseModel):
|
||||||
|
token: Annotated[
|
||||||
|
str, Field(description="One-time registration token from POST /v1/hosts.")
|
||||||
|
]
|
||||||
|
arch: Annotated[
|
||||||
|
str | None, Field(description="CPU architecture (e.g. x86_64, aarch64).")
|
||||||
|
] = None
|
||||||
|
cpu_cores: int | None = None
|
||||||
|
memory_mb: int | None = None
|
||||||
|
disk_gb: int | None = None
|
||||||
|
address: Annotated[str, Field(description="Host agent address (ip:port).")]
|
||||||
|
|
||||||
|
|
||||||
|
class Type3(StrEnum):
|
||||||
|
regular = "regular"
|
||||||
|
byoc = "byoc"
|
||||||
|
|
||||||
|
|
||||||
|
class Status1(StrEnum):
|
||||||
|
pending = "pending"
|
||||||
|
online = "online"
|
||||||
|
offline = "offline"
|
||||||
|
draining = "draining"
|
||||||
|
unreachable = "unreachable"
|
||||||
|
|
||||||
|
|
||||||
|
class Host(BaseModel):
|
||||||
|
id: str | None = None
|
||||||
|
type: Type3 | None = None
|
||||||
|
team_id: str | None = None
|
||||||
|
provider: str | None = None
|
||||||
|
availability_zone: str | None = None
|
||||||
|
arch: str | None = None
|
||||||
|
cpu_cores: int | None = None
|
||||||
|
memory_mb: int | None = None
|
||||||
|
disk_gb: int | None = None
|
||||||
|
address: str | None = None
|
||||||
|
status: Status1 | None = None
|
||||||
|
last_heartbeat_at: AwareDatetime | None = None
|
||||||
|
created_by: str | None = None
|
||||||
|
created_at: AwareDatetime | None = None
|
||||||
|
updated_at: AwareDatetime | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class RefreshHostTokenRequest(BaseModel):
|
||||||
|
refresh_token: Annotated[
|
||||||
|
str,
|
||||||
|
Field(
|
||||||
|
description="Refresh token obtained from registration or a previous refresh."
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class RefreshHostTokenResponse(BaseModel):
|
||||||
|
host: Host | None = None
|
||||||
|
token: Annotated[
|
||||||
|
str | None, Field(description="New host JWT. Valid for 7 days.")
|
||||||
|
] = None
|
||||||
|
refresh_token: Annotated[
|
||||||
|
str | None,
|
||||||
|
Field(
|
||||||
|
description="New refresh token. Valid for 60 days; old token is revoked."
|
||||||
|
),
|
||||||
|
] = None
|
||||||
|
|
||||||
|
|
||||||
|
class HostDeletePreview(BaseModel):
|
||||||
|
host: Host | None = None
|
||||||
|
sandbox_ids: Annotated[
|
||||||
|
list[str] | None,
|
||||||
|
Field(description="IDs of capsulees that would be destroyed on force-delete."),
|
||||||
|
] = None
|
||||||
|
|
||||||
|
|
||||||
|
class Error(BaseModel):
|
||||||
|
code: Annotated[str | None, Field(examples=["host_has_sandboxes"])] = None
|
||||||
|
message: str | None = None
|
||||||
|
sandbox_ids: Annotated[
|
||||||
|
list[str] | None,
|
||||||
|
Field(description="IDs of active capsulees blocking deletion."),
|
||||||
|
] = None
|
||||||
|
|
||||||
|
|
||||||
|
class HostHasCapsulesError(BaseModel):
|
||||||
|
error: Error | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class AddTagRequest(BaseModel):
|
||||||
|
tag: str
|
||||||
|
|
||||||
|
|
||||||
|
class UserSearchResult(BaseModel):
|
||||||
|
user_id: str | None = None
|
||||||
|
email: str | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class Team(BaseModel):
|
||||||
|
id: str | None = None
|
||||||
|
name: str | None = None
|
||||||
|
slug: Annotated[
|
||||||
|
str | None, Field(description="Immutable 12-char hex slug (e.g. a1b2c3-d1e2f3)")
|
||||||
|
] = None
|
||||||
|
created_at: AwareDatetime | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class Role(StrEnum):
|
||||||
|
owner = "owner"
|
||||||
|
admin = "admin"
|
||||||
|
member = "member"
|
||||||
|
|
||||||
|
|
||||||
|
class TeamWithRole(Team):
|
||||||
|
role: Role | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class TeamMember(BaseModel):
|
||||||
|
user_id: str | None = None
|
||||||
|
email: str | None = None
|
||||||
|
role: Role | None = None
|
||||||
|
joined_at: AwareDatetime | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class TeamDetail(BaseModel):
|
||||||
|
team: Team | None = None
|
||||||
|
members: list[TeamMember] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class Range1(StrEnum):
|
||||||
|
field_5m = "5m"
|
||||||
|
field_10m = "10m"
|
||||||
|
field_1h = "1h"
|
||||||
|
field_2h = "2h"
|
||||||
|
field_6h = "6h"
|
||||||
|
field_12h = "12h"
|
||||||
|
field_24h = "24h"
|
||||||
|
|
||||||
|
|
||||||
|
class MetricPoint(BaseModel):
|
||||||
|
timestamp_unix: int | None = None
|
||||||
|
cpu_pct: Annotated[
|
||||||
|
float | None,
|
||||||
|
Field(
|
||||||
|
description="CPU utilization percentage (0-100), normalized to vCPU count"
|
||||||
|
),
|
||||||
|
] = None
|
||||||
|
mem_bytes: Annotated[
|
||||||
|
int | None,
|
||||||
|
Field(description="Resident memory in bytes (VmRSS of Firecracker process)"),
|
||||||
|
] = None
|
||||||
|
disk_bytes: Annotated[
|
||||||
|
int | None, Field(description="Allocated disk bytes for the CoW sparse file")
|
||||||
|
] = None
|
||||||
|
|
||||||
|
|
||||||
|
class Provider(StrEnum):
|
||||||
|
discord = "discord"
|
||||||
|
slack = "slack"
|
||||||
|
teams = "teams"
|
||||||
|
googlechat = "googlechat"
|
||||||
|
telegram = "telegram"
|
||||||
|
matrix = "matrix"
|
||||||
|
webhook = "webhook"
|
||||||
|
|
||||||
|
|
||||||
|
class Event(StrEnum):
|
||||||
|
capsule_created = "capsule.created"
|
||||||
|
capsule_running = "capsule.running"
|
||||||
|
capsule_paused = "capsule.paused"
|
||||||
|
capsule_destroyed = "capsule.destroyed"
|
||||||
|
template_snapshot_created = "template.snapshot.created"
|
||||||
|
template_snapshot_deleted = "template.snapshot.deleted"
|
||||||
|
host_up = "host.up"
|
||||||
|
host_down = "host.down"
|
||||||
|
|
||||||
|
|
||||||
|
class CreateChannelRequest(BaseModel):
|
||||||
|
name: Annotated[str, Field(description="Unique channel name within the team.")]
|
||||||
|
provider: Provider
|
||||||
|
config: Annotated[
|
||||||
|
dict[str, str],
|
||||||
|
Field(
|
||||||
|
description='Provider-specific configuration fields. Discord/Slack/Teams/Google Chat: {"webhook_url": "..."}. Telegram: {"bot_token": "...", "chat_id": "..."}. Matrix: {"homeserver_url": "...", "access_token": "...", "room_id": "..."}. Webhook: {"url": "...", "secret": "..."} (secret is auto-generated if omitted).\n'
|
||||||
|
),
|
||||||
|
]
|
||||||
|
events: list[Event]
|
||||||
|
|
||||||
|
|
||||||
|
class TestChannelRequest(BaseModel):
|
||||||
|
provider: Provider
|
||||||
|
config: Annotated[
|
||||||
|
dict[str, str],
|
||||||
|
Field(
|
||||||
|
description="Provider-specific configuration fields (same as CreateChannelRequest.config)."
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class RotateConfigRequest(BaseModel):
|
||||||
|
config: Annotated[
|
||||||
|
dict[str, str],
|
||||||
|
Field(
|
||||||
|
description="New provider configuration fields. Must include all required fields for the channel's provider. Replaces the existing config entirely.\n"
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class UpdateChannelRequest(BaseModel):
|
||||||
|
name: str
|
||||||
|
events: list[Event]
|
||||||
|
|
||||||
|
|
||||||
|
class ChannelResponse(BaseModel):
|
||||||
|
id: str | None = None
|
||||||
|
team_id: str | None = None
|
||||||
|
name: str | None = None
|
||||||
|
provider: Provider | None = None
|
||||||
|
events: list[str] | None = None
|
||||||
|
created_at: AwareDatetime | None = None
|
||||||
|
updated_at: AwareDatetime | None = None
|
||||||
|
secret: Annotated[
|
||||||
|
str | None,
|
||||||
|
Field(description="Webhook secret. Only returned on creation, never again."),
|
||||||
|
] = None
|
||||||
|
|
||||||
|
|
||||||
|
class Error2(BaseModel):
|
||||||
|
code: str | None = None
|
||||||
|
message: str | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class Error1(BaseModel):
|
||||||
|
error: Error2 | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class ListDirResponse(BaseModel):
|
||||||
|
entries: list[FileEntry] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class CreateHostResponse(BaseModel):
|
||||||
|
host: Host | None = None
|
||||||
|
registration_token: Annotated[
|
||||||
|
str | None,
|
||||||
|
Field(
|
||||||
|
description="One-time registration token for the host agent. Expires in 1 hour."
|
||||||
|
),
|
||||||
|
] = None
|
||||||
|
|
||||||
|
|
||||||
|
class RegisterHostResponse(BaseModel):
|
||||||
|
host: Host | None = None
|
||||||
|
token: Annotated[
|
||||||
|
str | None,
|
||||||
|
Field(description="Host JWT for X-Host-Token header. Valid for 7 days."),
|
||||||
|
] = None
|
||||||
|
refresh_token: Annotated[
|
||||||
|
str | None,
|
||||||
|
Field(
|
||||||
|
description="Refresh token for obtaining new JWTs. Valid for 60 days; rotated on each use."
|
||||||
|
),
|
||||||
|
] = None
|
||||||
|
|
||||||
|
|
||||||
|
class CapsuleMetrics(BaseModel):
|
||||||
|
sandbox_id: str | None = None
|
||||||
|
range: Range1 | None = None
|
||||||
|
points: list[MetricPoint] | None = None
|
||||||
306
src/wrenn/pty.py
Normal file
306
src/wrenn/pty.py
Normal file
@ -0,0 +1,306 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import json
|
||||||
|
from collections.abc import AsyncIterator, Iterator
|
||||||
|
from enum import StrEnum
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import httpx_ws
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
|
||||||
|
class PtyEventType(StrEnum):
|
||||||
|
started = "started"
|
||||||
|
output = "output"
|
||||||
|
exit = "exit"
|
||||||
|
error = "error"
|
||||||
|
ping = "ping"
|
||||||
|
|
||||||
|
|
||||||
|
class PtyEvent(BaseModel):
|
||||||
|
type: PtyEventType
|
||||||
|
pid: int | None = None
|
||||||
|
tag: str | None = None
|
||||||
|
data: bytes | str | None = None
|
||||||
|
exit_code: int | None = None
|
||||||
|
fatal: bool | None = None
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_pty_event(raw: dict[str, Any]) -> PtyEvent:
|
||||||
|
msg_type = raw.get("type", "")
|
||||||
|
if msg_type == "started":
|
||||||
|
return PtyEvent(
|
||||||
|
type=PtyEventType.started,
|
||||||
|
pid=raw.get("pid"),
|
||||||
|
tag=raw.get("tag"),
|
||||||
|
)
|
||||||
|
if msg_type == "output":
|
||||||
|
raw_data = raw.get("data", "")
|
||||||
|
decoded = base64.b64decode(raw_data) if raw_data else b""
|
||||||
|
return PtyEvent(type=PtyEventType.output, data=decoded)
|
||||||
|
if msg_type == "exit":
|
||||||
|
return PtyEvent(type=PtyEventType.exit, exit_code=raw.get("exit_code", -1))
|
||||||
|
if msg_type == "error":
|
||||||
|
return PtyEvent(
|
||||||
|
type=PtyEventType.error,
|
||||||
|
data=raw.get("data", ""),
|
||||||
|
fatal=raw.get("fatal", False),
|
||||||
|
)
|
||||||
|
if msg_type == "ping":
|
||||||
|
return PtyEvent(type=PtyEventType.ping)
|
||||||
|
return PtyEvent(type=PtyEventType(msg_type) if msg_type else PtyEventType.ping)
|
||||||
|
|
||||||
|
|
||||||
|
class PtySession:
|
||||||
|
"""Interactive PTY session backed by a WebSocket.
|
||||||
|
|
||||||
|
Use as a context manager and iterate over events::
|
||||||
|
|
||||||
|
with sb.pty(cmd="/bin/bash") as term:
|
||||||
|
term.write(b"ls -la\\n")
|
||||||
|
for event in term:
|
||||||
|
if event.type == "output":
|
||||||
|
sys.stdout.buffer.write(event.data)
|
||||||
|
elif event.type == "exit":
|
||||||
|
break
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, ws: httpx_ws.WebSocketSession, capsule_id: str) -> None:
|
||||||
|
self._ws = ws
|
||||||
|
self._capsule_id = capsule_id
|
||||||
|
self._tag: str | None = None
|
||||||
|
self._pid: int | None = None
|
||||||
|
self._done = False
|
||||||
|
|
||||||
|
@property
|
||||||
|
def tag(self) -> str | None:
|
||||||
|
"""Session tag. Available after the ``started`` event."""
|
||||||
|
return self._tag
|
||||||
|
|
||||||
|
@property
|
||||||
|
def pid(self) -> int | None:
|
||||||
|
"""Process PID. Available after the ``started`` event."""
|
||||||
|
return self._pid
|
||||||
|
|
||||||
|
def _send_start(
|
||||||
|
self,
|
||||||
|
cmd: str = "/bin/bash",
|
||||||
|
args: list[str] | None = None,
|
||||||
|
cols: int = 80,
|
||||||
|
rows: int = 24,
|
||||||
|
envs: dict[str, str] | None = None,
|
||||||
|
cwd: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
msg: dict[str, Any] = {
|
||||||
|
"type": "start",
|
||||||
|
"cmd": cmd,
|
||||||
|
"cols": cols or 80,
|
||||||
|
"rows": rows or 24,
|
||||||
|
}
|
||||||
|
if args:
|
||||||
|
msg["args"] = args
|
||||||
|
if envs:
|
||||||
|
msg["envs"] = envs
|
||||||
|
if cwd:
|
||||||
|
msg["cwd"] = cwd
|
||||||
|
self._ws.send_text(json.dumps(msg))
|
||||||
|
|
||||||
|
def _send_connect(self, tag: str) -> None:
|
||||||
|
self._ws.send_text(json.dumps({"type": "connect", "tag": tag}))
|
||||||
|
|
||||||
|
def write(self, data: bytes) -> None:
|
||||||
|
"""Send raw bytes to the PTY stdin.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
data: Raw bytes to send. Base64-encoded internally.
|
||||||
|
"""
|
||||||
|
encoded = base64.b64encode(data).decode("ascii")
|
||||||
|
self._ws.send_text(json.dumps({"type": "input", "data": encoded}))
|
||||||
|
|
||||||
|
def resize(self, cols: int, rows: int) -> None:
|
||||||
|
"""Resize the PTY terminal.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
cols: New column count. Must be > 0.
|
||||||
|
rows: New row count. Must be > 0.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If cols or rows is 0.
|
||||||
|
"""
|
||||||
|
if cols <= 0 or rows <= 0:
|
||||||
|
raise ValueError("cols and rows must be greater than 0")
|
||||||
|
self._ws.send_text(json.dumps({"type": "resize", "cols": cols, "rows": rows}))
|
||||||
|
|
||||||
|
def kill(self) -> None:
|
||||||
|
"""Send SIGKILL to the PTY process."""
|
||||||
|
self._ws.send_text(json.dumps({"type": "kill"}))
|
||||||
|
|
||||||
|
def __iter__(self) -> Iterator[PtyEvent]:
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __next__(self) -> PtyEvent:
|
||||||
|
if self._done:
|
||||||
|
raise StopIteration
|
||||||
|
try:
|
||||||
|
raw = self._ws.receive_text()
|
||||||
|
except httpx_ws.WebSocketDisconnect:
|
||||||
|
raise StopIteration
|
||||||
|
event = _parse_pty_event(json.loads(raw))
|
||||||
|
if event.type == PtyEventType.started:
|
||||||
|
if event.tag is not None:
|
||||||
|
self._tag = event.tag
|
||||||
|
if event.pid is not None:
|
||||||
|
self._pid = event.pid
|
||||||
|
if event.type == PtyEventType.exit:
|
||||||
|
raise StopIteration
|
||||||
|
if event.type == PtyEventType.error and event.fatal:
|
||||||
|
self._done = True
|
||||||
|
return event
|
||||||
|
return event
|
||||||
|
|
||||||
|
def __enter__(self) -> PtySession:
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __exit__(
|
||||||
|
self,
|
||||||
|
exc_type: type[BaseException] | None,
|
||||||
|
exc_val: BaseException | None,
|
||||||
|
exc_tb: object,
|
||||||
|
) -> None:
|
||||||
|
try:
|
||||||
|
self.kill()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
self._ws.close()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class AsyncPtySession:
|
||||||
|
"""Async interactive PTY session backed by a WebSocket.
|
||||||
|
|
||||||
|
Use as an async context manager and async iterate over events::
|
||||||
|
|
||||||
|
async with sb.pty(cmd="/bin/bash") as term:
|
||||||
|
await term.write(b"ls -la\\n")
|
||||||
|
async for event in term:
|
||||||
|
if event.type == "output":
|
||||||
|
sys.stdout.buffer.write(event.data)
|
||||||
|
elif event.type == "exit":
|
||||||
|
break
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, ws: httpx_ws.AsyncWebSocketSession, capsule_id: str) -> None:
|
||||||
|
self._ws = ws
|
||||||
|
self._capsule_id = capsule_id
|
||||||
|
self._tag: str | None = None
|
||||||
|
self._pid: int | None = None
|
||||||
|
self._done = False
|
||||||
|
|
||||||
|
@property
|
||||||
|
def tag(self) -> str | None:
|
||||||
|
"""Session tag. Available after the ``started`` event."""
|
||||||
|
return self._tag
|
||||||
|
|
||||||
|
@property
|
||||||
|
def pid(self) -> int | None:
|
||||||
|
"""Process PID. Available after the ``started`` event."""
|
||||||
|
return self._pid
|
||||||
|
|
||||||
|
async def _send_start(
|
||||||
|
self,
|
||||||
|
cmd: str = "/bin/bash",
|
||||||
|
args: list[str] | None = None,
|
||||||
|
cols: int = 80,
|
||||||
|
rows: int = 24,
|
||||||
|
envs: dict[str, str] | None = None,
|
||||||
|
cwd: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
msg: dict[str, Any] = {
|
||||||
|
"type": "start",
|
||||||
|
"cmd": cmd,
|
||||||
|
"cols": cols or 80,
|
||||||
|
"rows": rows or 24,
|
||||||
|
}
|
||||||
|
if args:
|
||||||
|
msg["args"] = args
|
||||||
|
if envs:
|
||||||
|
msg["envs"] = envs
|
||||||
|
if cwd:
|
||||||
|
msg["cwd"] = cwd
|
||||||
|
await self._ws.send_text(json.dumps(msg))
|
||||||
|
|
||||||
|
async def _send_connect(self, tag: str) -> None:
|
||||||
|
await self._ws.send_text(json.dumps({"type": "connect", "tag": tag}))
|
||||||
|
|
||||||
|
async def write(self, data: bytes) -> None:
|
||||||
|
"""Send raw bytes to the PTY stdin.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
data: Raw bytes to send. Base64-encoded internally.
|
||||||
|
"""
|
||||||
|
encoded = base64.b64encode(data).decode("ascii")
|
||||||
|
await self._ws.send_text(json.dumps({"type": "input", "data": encoded}))
|
||||||
|
|
||||||
|
async def resize(self, cols: int, rows: int) -> None:
|
||||||
|
"""Resize the PTY terminal.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
cols: New column count. Must be > 0.
|
||||||
|
rows: New row count. Must be > 0.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If cols or rows is 0.
|
||||||
|
"""
|
||||||
|
if cols <= 0 or rows <= 0:
|
||||||
|
raise ValueError("cols and rows must be greater than 0")
|
||||||
|
await self._ws.send_text(
|
||||||
|
json.dumps({"type": "resize", "cols": cols, "rows": rows})
|
||||||
|
)
|
||||||
|
|
||||||
|
async def kill(self) -> None:
|
||||||
|
"""Send SIGKILL to the PTY process."""
|
||||||
|
await self._ws.send_text(json.dumps({"type": "kill"}))
|
||||||
|
|
||||||
|
def __aiter__(self) -> AsyncIterator[PtyEvent]:
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __anext__(self) -> PtyEvent:
|
||||||
|
if self._done:
|
||||||
|
raise StopAsyncIteration
|
||||||
|
try:
|
||||||
|
raw = await self._ws.receive_text()
|
||||||
|
except httpx_ws.WebSocketDisconnect:
|
||||||
|
raise StopAsyncIteration
|
||||||
|
event = _parse_pty_event(json.loads(raw))
|
||||||
|
if event.type == PtyEventType.started:
|
||||||
|
if event.tag is not None:
|
||||||
|
self._tag = event.tag
|
||||||
|
if event.pid is not None:
|
||||||
|
self._pid = event.pid
|
||||||
|
if event.type == PtyEventType.exit:
|
||||||
|
raise StopAsyncIteration
|
||||||
|
if event.type == PtyEventType.error and event.fatal:
|
||||||
|
self._done = True
|
||||||
|
return event
|
||||||
|
return event
|
||||||
|
|
||||||
|
async def __aenter__(self) -> AsyncPtySession:
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(
|
||||||
|
self,
|
||||||
|
exc_type: type[BaseException] | None,
|
||||||
|
exc_val: BaseException | None,
|
||||||
|
exc_tb: object,
|
||||||
|
) -> None:
|
||||||
|
try:
|
||||||
|
await self.kill()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
await self._ws.close()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
22
src/wrenn/sandbox.py
Normal file
22
src/wrenn/sandbox.py
Normal file
@ -0,0 +1,22 @@
|
|||||||
|
import warnings as _warnings
|
||||||
|
|
||||||
|
from wrenn.capsule import Capsule # noqa: F401
|
||||||
|
from wrenn.commands import ( # noqa: F401
|
||||||
|
StreamErrorEvent,
|
||||||
|
StreamEvent,
|
||||||
|
StreamExitEvent,
|
||||||
|
StreamStartEvent,
|
||||||
|
StreamStderrEvent,
|
||||||
|
StreamStdoutEvent,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def __getattr__(name: str) -> type:
|
||||||
|
if name == "Sandbox":
|
||||||
|
_warnings.warn(
|
||||||
|
"'Sandbox' is deprecated, use 'Capsule' instead",
|
||||||
|
FutureWarning,
|
||||||
|
stacklevel=2,
|
||||||
|
)
|
||||||
|
return Capsule
|
||||||
|
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||||
166
tests/test_capsule_features.py
Normal file
166
tests/test_capsule_features.py
Normal file
@ -0,0 +1,166 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import respx
|
||||||
|
|
||||||
|
from wrenn.capsule import Capsule, _build_proxy_url
|
||||||
|
from wrenn.code_interpreter.capsule import CodeResult
|
||||||
|
|
||||||
|
BASE = "https://app.wrenn.dev/api"
|
||||||
|
|
||||||
|
|
||||||
|
class TestBuildProxyUrl:
|
||||||
|
def test_https_production(self):
|
||||||
|
url = _build_proxy_url("https://app.wrenn.dev/api", "cl-abc123", 8888)
|
||||||
|
assert url == "wss://8888-cl-abc123.app.wrenn.dev"
|
||||||
|
|
||||||
|
def test_http_localhost(self):
|
||||||
|
url = _build_proxy_url("http://localhost:8080", "cl-abc123", 3000)
|
||||||
|
assert url == "ws://3000-cl-abc123.localhost:8080"
|
||||||
|
|
||||||
|
def test_https_custom_port(self):
|
||||||
|
url = _build_proxy_url("https://api.example.com:9443", "sb-1", 8080)
|
||||||
|
assert url == "wss://8080-sb-1.api.example.com:9443"
|
||||||
|
|
||||||
|
def test_http_no_port(self):
|
||||||
|
url = _build_proxy_url("http://192.168.1.1", "sb-2", 5000)
|
||||||
|
assert url == "ws://5000-sb-2.192.168.1.1"
|
||||||
|
|
||||||
|
|
||||||
|
class TestCapsuleCreate:
|
||||||
|
@respx.mock
|
||||||
|
def test_capsule_constructor_creates(self):
|
||||||
|
respx.post(f"{BASE}/v1/capsules").respond(
|
||||||
|
201, json={"id": "cl-1", "status": "pending", "template": "minimal"}
|
||||||
|
)
|
||||||
|
cap = Capsule(template="minimal", api_key="wrn_test1234567890abcdef12345678")
|
||||||
|
assert cap.capsule_id == "cl-1"
|
||||||
|
assert hasattr(cap, "commands")
|
||||||
|
assert hasattr(cap, "files")
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_capsule_create_classmethod(self):
|
||||||
|
respx.post(f"{BASE}/v1/capsules").respond(
|
||||||
|
201, json={"id": "cl-2", "status": "pending"}
|
||||||
|
)
|
||||||
|
cap = Capsule.create(api_key="wrn_test1234567890abcdef12345678")
|
||||||
|
assert cap.capsule_id == "cl-2"
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_capsule_context_manager_kills(self):
|
||||||
|
respx.post(f"{BASE}/v1/capsules").respond(
|
||||||
|
201, json={"id": "cl-1", "status": "pending"}
|
||||||
|
)
|
||||||
|
kill_route = respx.delete(f"{BASE}/v1/capsules/cl-1").respond(204)
|
||||||
|
with Capsule(api_key="wrn_test1234567890abcdef12345678") as cap:
|
||||||
|
assert cap.capsule_id == "cl-1"
|
||||||
|
assert kill_route.called
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_capsule_env_var(self, monkeypatch):
|
||||||
|
monkeypatch.setenv("WRENN_API_KEY", "wrn_from_env_key")
|
||||||
|
respx.post(f"{BASE}/v1/capsules").respond(
|
||||||
|
201, json={"id": "cl-3", "status": "pending"}
|
||||||
|
)
|
||||||
|
cap = Capsule()
|
||||||
|
assert cap.capsule_id == "cl-3"
|
||||||
|
|
||||||
|
|
||||||
|
class TestCapsuleStaticMethods:
|
||||||
|
@respx.mock
|
||||||
|
def test_static_destroy(self):
|
||||||
|
route = respx.delete(f"{BASE}/v1/capsules/cl-1").respond(204)
|
||||||
|
Capsule._static_destroy("cl-1", api_key="wrn_test1234567890abcdef12345678")
|
||||||
|
assert route.called
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_static_pause(self):
|
||||||
|
respx.post(f"{BASE}/v1/capsules/cl-1/pause").respond(
|
||||||
|
200, json={"id": "cl-1", "status": "paused"}
|
||||||
|
)
|
||||||
|
info = Capsule._static_pause("cl-1", api_key="wrn_test1234567890abcdef12345678")
|
||||||
|
assert info.status.value == "paused"
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_static_list(self):
|
||||||
|
respx.get(f"{BASE}/v1/capsules").respond(
|
||||||
|
200, json=[{"id": "cl-1", "status": "running"}]
|
||||||
|
)
|
||||||
|
items = Capsule.list(api_key="wrn_test1234567890abcdef12345678")
|
||||||
|
assert len(items) == 1
|
||||||
|
assert items[0].id == "cl-1"
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_static_get_info(self):
|
||||||
|
respx.get(f"{BASE}/v1/capsules/cl-1").respond(
|
||||||
|
200, json={"id": "cl-1", "status": "running"}
|
||||||
|
)
|
||||||
|
info = Capsule._static_get_info("cl-1", api_key="wrn_test1234567890abcdef12345678")
|
||||||
|
assert info.id == "cl-1"
|
||||||
|
|
||||||
|
|
||||||
|
class TestCapsuleConnect:
|
||||||
|
@respx.mock
|
||||||
|
def test_connect_running(self):
|
||||||
|
respx.get(f"{BASE}/v1/capsules/cl-1").respond(
|
||||||
|
200, json={"id": "cl-1", "status": "running"}
|
||||||
|
)
|
||||||
|
cap = Capsule.connect("cl-1", api_key="wrn_test1234567890abcdef12345678")
|
||||||
|
assert cap.capsule_id == "cl-1"
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_connect_paused_resumes(self):
|
||||||
|
respx.get(f"{BASE}/v1/capsules/cl-1").respond(
|
||||||
|
200, json={"id": "cl-1", "status": "paused"}
|
||||||
|
)
|
||||||
|
respx.post(f"{BASE}/v1/capsules/cl-1/resume").respond(
|
||||||
|
200, json={"id": "cl-1", "status": "running"}
|
||||||
|
)
|
||||||
|
cap = Capsule.connect("cl-1", api_key="wrn_test1234567890abcdef12345678")
|
||||||
|
assert cap.capsule_id == "cl-1"
|
||||||
|
|
||||||
|
|
||||||
|
class TestCodeResult:
|
||||||
|
def test_defaults(self):
|
||||||
|
r = CodeResult()
|
||||||
|
assert r.text is None
|
||||||
|
assert r.data is None
|
||||||
|
assert r.stdout == ""
|
||||||
|
assert r.stderr == ""
|
||||||
|
assert r.error is None
|
||||||
|
|
||||||
|
def test_with_values(self):
|
||||||
|
r = CodeResult(
|
||||||
|
text="84",
|
||||||
|
data={"text/plain": "84"},
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
error=None,
|
||||||
|
)
|
||||||
|
assert r.text == "84"
|
||||||
|
assert r.data["text/plain"] == "84"
|
||||||
|
|
||||||
|
def test_error_result(self):
|
||||||
|
r = CodeResult(error="ZeroDivisionError: division by zero\n...")
|
||||||
|
assert r.error is not None
|
||||||
|
assert "ZeroDivisionError" in r.error
|
||||||
|
|
||||||
|
|
||||||
|
class TestDeprecationWarnings:
|
||||||
|
def test_import_sandbox_from_wrenn_warns(self):
|
||||||
|
import importlib
|
||||||
|
import sys
|
||||||
|
import warnings
|
||||||
|
|
||||||
|
# Clear cached attribute
|
||||||
|
if "Sandbox" in dir(sys.modules.get("wrenn", object())):
|
||||||
|
delattr(sys.modules["wrenn"], "Sandbox")
|
||||||
|
|
||||||
|
with warnings.catch_warnings(record=True) as w:
|
||||||
|
warnings.simplefilter("always")
|
||||||
|
from wrenn import Sandbox
|
||||||
|
|
||||||
|
assert Sandbox is Capsule
|
||||||
|
fw = [x for x in w if issubclass(x.category, FutureWarning)]
|
||||||
|
assert len(fw) >= 1
|
||||||
|
assert "Sandbox" in str(fw[0].message)
|
||||||
264
tests/test_client.py
Normal file
264
tests/test_client.py
Normal file
@ -0,0 +1,264 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import respx
|
||||||
|
|
||||||
|
from wrenn.client import AsyncWrennClient, WrennClient
|
||||||
|
from wrenn.exceptions import (
|
||||||
|
WrennAgentError,
|
||||||
|
WrennAuthenticationError,
|
||||||
|
WrennConflictError,
|
||||||
|
WrennInternalError,
|
||||||
|
WrennNotFoundError,
|
||||||
|
WrennValidationError,
|
||||||
|
)
|
||||||
|
from wrenn.models import (
|
||||||
|
Capsule,
|
||||||
|
Status,
|
||||||
|
Template,
|
||||||
|
)
|
||||||
|
|
||||||
|
BASE = "https://app.wrenn.dev/api"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def client():
|
||||||
|
with WrennClient(api_key="wrn_test1234567890abcdef12345678") as c:
|
||||||
|
yield c
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def async_client():
|
||||||
|
return AsyncWrennClient(api_key="wrn_test1234567890abcdef12345678")
|
||||||
|
|
||||||
|
|
||||||
|
class TestCapsules:
|
||||||
|
@respx.mock
|
||||||
|
def test_create(self, client):
|
||||||
|
respx.post(f"{BASE}/v1/capsules").respond(
|
||||||
|
201,
|
||||||
|
json={
|
||||||
|
"id": "sb-1",
|
||||||
|
"status": "pending",
|
||||||
|
"template": "base-python",
|
||||||
|
"vcpus": 2,
|
||||||
|
"memory_mb": 1024,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
resp = client.capsules.create(template="base-python", vcpus=2, memory_mb=1024)
|
||||||
|
assert isinstance(resp, Capsule)
|
||||||
|
assert resp.id == "sb-1"
|
||||||
|
assert resp.status == Status.pending
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_create_defaults(self, client):
|
||||||
|
respx.post(f"{BASE}/v1/capsules").respond(
|
||||||
|
201, json={"id": "sb-2", "status": "pending"}
|
||||||
|
)
|
||||||
|
resp = client.capsules.create()
|
||||||
|
assert resp.id == "sb-2"
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_list(self, client):
|
||||||
|
respx.get(f"{BASE}/v1/capsules").respond(
|
||||||
|
200, json=[{"id": "sb-1", "status": "running"}]
|
||||||
|
)
|
||||||
|
boxes = client.capsules.list()
|
||||||
|
assert len(boxes) == 1
|
||||||
|
assert boxes[0].status == Status.running
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_get(self, client):
|
||||||
|
respx.get(f"{BASE}/v1/capsules/sb-1").respond(
|
||||||
|
200, json={"id": "sb-1", "status": "running"}
|
||||||
|
)
|
||||||
|
resp = client.capsules.get("sb-1")
|
||||||
|
assert resp.id == "sb-1"
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_destroy(self, client):
|
||||||
|
route = respx.delete(f"{BASE}/v1/capsules/sb-1").respond(204)
|
||||||
|
client.capsules.destroy("sb-1")
|
||||||
|
assert route.called
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_pause(self, client):
|
||||||
|
respx.post(f"{BASE}/v1/capsules/sb-1/pause").respond(
|
||||||
|
200, json={"id": "sb-1", "status": "paused"}
|
||||||
|
)
|
||||||
|
resp = client.capsules.pause("sb-1")
|
||||||
|
assert resp.status == Status.paused
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_resume(self, client):
|
||||||
|
respx.post(f"{BASE}/v1/capsules/sb-1/resume").respond(
|
||||||
|
200, json={"id": "sb-1", "status": "running"}
|
||||||
|
)
|
||||||
|
resp = client.capsules.resume("sb-1")
|
||||||
|
assert resp.status == Status.running
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_ping(self, client):
|
||||||
|
route = respx.post(f"{BASE}/v1/capsules/sb-1/ping").respond(204)
|
||||||
|
client.capsules.ping("sb-1")
|
||||||
|
assert route.called
|
||||||
|
|
||||||
|
|
||||||
|
class TestSnapshots:
|
||||||
|
@respx.mock
|
||||||
|
def test_create(self, client):
|
||||||
|
respx.post(f"{BASE}/v1/snapshots").respond(
|
||||||
|
201,
|
||||||
|
json={"name": "snap-1", "type": "snapshot", "vcpus": 1},
|
||||||
|
)
|
||||||
|
resp = client.snapshots.create(capsule_id="sb-1", name="snap-1")
|
||||||
|
assert isinstance(resp, Template)
|
||||||
|
assert resp.name == "snap-1"
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_create_with_overwrite(self, client):
|
||||||
|
route = respx.post(f"{BASE}/v1/snapshots").respond(
|
||||||
|
201, json={"name": "snap-1", "type": "snapshot"}
|
||||||
|
)
|
||||||
|
client.snapshots.create(capsule_id="sb-1", overwrite=True)
|
||||||
|
req = route.calls[0].request
|
||||||
|
assert "overwrite=true" in str(req.url)
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_list(self, client):
|
||||||
|
respx.get(f"{BASE}/v1/snapshots").respond(
|
||||||
|
200, json=[{"name": "base-python", "type": "base"}]
|
||||||
|
)
|
||||||
|
snaps = client.snapshots.list()
|
||||||
|
assert len(snaps) == 1
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_list_with_filter(self, client):
|
||||||
|
route = respx.get(f"{BASE}/v1/snapshots").respond(200, json=[])
|
||||||
|
client.snapshots.list(type="snapshot")
|
||||||
|
req = route.calls[0].request
|
||||||
|
assert "type=snapshot" in str(req.url)
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_delete(self, client):
|
||||||
|
route = respx.delete(f"{BASE}/v1/snapshots/snap-1").respond(204)
|
||||||
|
client.snapshots.delete("snap-1")
|
||||||
|
assert route.called
|
||||||
|
|
||||||
|
|
||||||
|
class TestErrorHandling:
|
||||||
|
@respx.mock
|
||||||
|
def test_validation_error(self, client):
|
||||||
|
respx.post(f"{BASE}/v1/capsules").respond(
|
||||||
|
400,
|
||||||
|
json={"error": {"code": "invalid_request", "message": "bad input"}},
|
||||||
|
)
|
||||||
|
with pytest.raises(WrennValidationError) as exc_info:
|
||||||
|
client.capsules.create()
|
||||||
|
assert exc_info.value.code == "invalid_request"
|
||||||
|
assert exc_info.value.status_code == 400
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_auth_error(self, client):
|
||||||
|
respx.get(f"{BASE}/v1/capsules").respond(
|
||||||
|
401,
|
||||||
|
json={"error": {"code": "unauthorized", "message": "bad key"}},
|
||||||
|
)
|
||||||
|
with pytest.raises(WrennAuthenticationError):
|
||||||
|
client.capsules.list()
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_not_found_error(self, client):
|
||||||
|
respx.get(f"{BASE}/v1/capsules/nope").respond(
|
||||||
|
404,
|
||||||
|
json={"error": {"code": "not_found", "message": "capsule not found"}},
|
||||||
|
)
|
||||||
|
with pytest.raises(WrennNotFoundError):
|
||||||
|
client.capsules.get("nope")
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_conflict_error(self, client):
|
||||||
|
respx.get(f"{BASE}/v1/capsules/sb-1").respond(
|
||||||
|
409,
|
||||||
|
json={"error": {"code": "invalid_state", "message": "not running"}},
|
||||||
|
)
|
||||||
|
with pytest.raises(WrennConflictError):
|
||||||
|
client.capsules.get("sb-1")
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_agent_error(self, client):
|
||||||
|
respx.post(f"{BASE}/v1/capsules").respond(
|
||||||
|
502,
|
||||||
|
json={"error": {"code": "agent_error", "message": "host agent failed"}},
|
||||||
|
)
|
||||||
|
with pytest.raises(WrennAgentError):
|
||||||
|
client.capsules.create()
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_internal_error(self, client):
|
||||||
|
respx.get(f"{BASE}/v1/capsules/sb-1").respond(
|
||||||
|
500,
|
||||||
|
json={"error": {"code": "internal_error", "message": "oops"}},
|
||||||
|
)
|
||||||
|
with pytest.raises(WrennInternalError):
|
||||||
|
client.capsules.get("sb-1")
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_unknown_error_code_falls_back(self, client):
|
||||||
|
respx.get(f"{BASE}/v1/capsules/sb-1").respond(
|
||||||
|
418,
|
||||||
|
json={"error": {"code": "teapot", "message": "I'm a teapot"}},
|
||||||
|
)
|
||||||
|
from wrenn.exceptions import WrennError
|
||||||
|
|
||||||
|
with pytest.raises(WrennError) as exc_info:
|
||||||
|
client.capsules.get("sb-1")
|
||||||
|
assert exc_info.value.code == "teapot"
|
||||||
|
|
||||||
|
|
||||||
|
class TestAuthModes:
|
||||||
|
def test_api_key_header(self):
|
||||||
|
with WrennClient(api_key="wrn_test1234567890abcdef12345678") as c:
|
||||||
|
assert c._http.headers["X-API-Key"] == "wrn_test1234567890abcdef12345678"
|
||||||
|
|
||||||
|
def test_no_auth_raises(self):
|
||||||
|
with pytest.raises(ValueError, match="No API key"):
|
||||||
|
WrennClient()
|
||||||
|
|
||||||
|
def test_env_var_fallback(self, monkeypatch):
|
||||||
|
monkeypatch.setenv("WRENN_API_KEY", "wrn_from_env")
|
||||||
|
with WrennClient() as c:
|
||||||
|
assert c._http.headers["X-API-Key"] == "wrn_from_env"
|
||||||
|
|
||||||
|
|
||||||
|
class TestAsyncClient:
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@respx.mock
|
||||||
|
async def test_async_capsules_create(self, async_client):
|
||||||
|
async with async_client:
|
||||||
|
respx.post(f"{BASE}/v1/capsules").respond(
|
||||||
|
201, json={"id": "sb-1", "status": "pending"}
|
||||||
|
)
|
||||||
|
resp = await async_client.capsules.create(template="base-python")
|
||||||
|
assert resp.id == "sb-1"
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@respx.mock
|
||||||
|
async def test_async_capsules_list(self, async_client):
|
||||||
|
async with async_client:
|
||||||
|
respx.get(f"{BASE}/v1/capsules").respond(
|
||||||
|
200, json=[{"id": "sb-1"}]
|
||||||
|
)
|
||||||
|
boxes = await async_client.capsules.list()
|
||||||
|
assert len(boxes) == 1
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@respx.mock
|
||||||
|
async def test_async_error_handling(self, async_client):
|
||||||
|
async with async_client:
|
||||||
|
respx.get(f"{BASE}/v1/capsules/nope").respond(
|
||||||
|
404,
|
||||||
|
json={"error": {"code": "not_found", "message": "not found"}},
|
||||||
|
)
|
||||||
|
with pytest.raises(WrennNotFoundError):
|
||||||
|
await async_client.capsules.get("nope")
|
||||||
495
tests/test_filesystem_pty.py
Normal file
495
tests/test_filesystem_pty.py
Normal file
@ -0,0 +1,495 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import json
|
||||||
|
from unittest.mock import AsyncMock, MagicMock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import respx
|
||||||
|
|
||||||
|
from wrenn.capsule import Capsule
|
||||||
|
from wrenn.models import FileEntry
|
||||||
|
from wrenn.pty import (
|
||||||
|
AsyncPtySession,
|
||||||
|
PtyEventType,
|
||||||
|
PtySession,
|
||||||
|
_parse_pty_event,
|
||||||
|
)
|
||||||
|
|
||||||
|
BASE = "https://app.wrenn.dev/api"
|
||||||
|
|
||||||
|
|
||||||
|
def _make_capsule(cap_id: str = "cl-abc") -> Capsule:
|
||||||
|
respx.post(f"{BASE}/v1/capsules").respond(
|
||||||
|
201, json={"id": cap_id, "status": "running"}
|
||||||
|
)
|
||||||
|
return Capsule(api_key="wrn_test1234567890abcdef12345678")
|
||||||
|
|
||||||
|
|
||||||
|
class TestFilesRead:
|
||||||
|
@respx.mock
|
||||||
|
def test_read_returns_string(self):
|
||||||
|
cap = _make_capsule()
|
||||||
|
content = b"file contents here"
|
||||||
|
respx.post(f"{BASE}/v1/capsules/cl-abc/files/read").respond(
|
||||||
|
200, content=content
|
||||||
|
)
|
||||||
|
data = cap.files.read("/app/main.py")
|
||||||
|
assert data == "file contents here"
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_read_bytes(self):
|
||||||
|
cap = _make_capsule()
|
||||||
|
content = b"\x00\x01\x02"
|
||||||
|
respx.post(f"{BASE}/v1/capsules/cl-abc/files/read").respond(
|
||||||
|
200, content=content
|
||||||
|
)
|
||||||
|
data = cap.files.read_bytes("/bin/binary")
|
||||||
|
assert data == b"\x00\x01\x02"
|
||||||
|
|
||||||
|
|
||||||
|
class TestFilesWrite:
|
||||||
|
@respx.mock
|
||||||
|
def test_write_string(self):
|
||||||
|
cap = _make_capsule()
|
||||||
|
route = respx.post(f"{BASE}/v1/capsules/cl-abc/files/write").respond(204)
|
||||||
|
cap.files.write("/app/main.py", "print('hello')")
|
||||||
|
assert route.called
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_write_bytes(self):
|
||||||
|
cap = _make_capsule()
|
||||||
|
route = respx.post(f"{BASE}/v1/capsules/cl-abc/files/write").respond(204)
|
||||||
|
cap.files.write("/app/data.bin", b"\x00\x01\x02")
|
||||||
|
assert route.called
|
||||||
|
|
||||||
|
|
||||||
|
class TestFilesList:
|
||||||
|
@respx.mock
|
||||||
|
def test_list_returns_entries(self):
|
||||||
|
cap = _make_capsule()
|
||||||
|
respx.post(f"{BASE}/v1/capsules/cl-abc/files/list").respond(
|
||||||
|
200,
|
||||||
|
json={
|
||||||
|
"entries": [
|
||||||
|
{
|
||||||
|
"name": "main.py",
|
||||||
|
"path": "/home/user/main.py",
|
||||||
|
"type": "file",
|
||||||
|
"size": 1024,
|
||||||
|
"mode": 33188,
|
||||||
|
"permissions": "-rw-r--r--",
|
||||||
|
"owner": "root",
|
||||||
|
"group": "root",
|
||||||
|
"modified_at": 1712899200,
|
||||||
|
"symlink_target": None,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "config",
|
||||||
|
"path": "/home/user/config",
|
||||||
|
"type": "directory",
|
||||||
|
"size": 4096,
|
||||||
|
"mode": 16877,
|
||||||
|
"permissions": "drwxr-xr-x",
|
||||||
|
"owner": "root",
|
||||||
|
"group": "root",
|
||||||
|
"modified_at": 1712899100,
|
||||||
|
"symlink_target": None,
|
||||||
|
},
|
||||||
|
]
|
||||||
|
},
|
||||||
|
)
|
||||||
|
entries = cap.files.list("/home/user")
|
||||||
|
assert len(entries) == 2
|
||||||
|
assert isinstance(entries[0], FileEntry)
|
||||||
|
assert entries[0].name == "main.py"
|
||||||
|
assert entries[0].type == "file"
|
||||||
|
assert entries[1].name == "config"
|
||||||
|
assert entries[1].type == "directory"
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_list_with_depth(self):
|
||||||
|
cap = _make_capsule()
|
||||||
|
route = respx.post(f"{BASE}/v1/capsules/cl-abc/files/list").respond(
|
||||||
|
200, json={"entries": []}
|
||||||
|
)
|
||||||
|
cap.files.list("/home/user", depth=3)
|
||||||
|
body = json.loads(route.calls[0].request.content)
|
||||||
|
assert body["depth"] == 3
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_list_empty(self):
|
||||||
|
cap = _make_capsule()
|
||||||
|
respx.post(f"{BASE}/v1/capsules/cl-abc/files/list").respond(
|
||||||
|
200, json={"entries": []}
|
||||||
|
)
|
||||||
|
entries = cap.files.list("/empty")
|
||||||
|
assert entries == []
|
||||||
|
|
||||||
|
|
||||||
|
class TestFilesMakeDir:
|
||||||
|
@respx.mock
|
||||||
|
def test_make_dir_returns_entry(self):
|
||||||
|
cap = _make_capsule()
|
||||||
|
respx.post(f"{BASE}/v1/capsules/cl-abc/files/mkdir").respond(
|
||||||
|
200,
|
||||||
|
json={
|
||||||
|
"entry": {
|
||||||
|
"name": "data",
|
||||||
|
"path": "/home/user/data",
|
||||||
|
"type": "directory",
|
||||||
|
"size": 4096,
|
||||||
|
"mode": 16877,
|
||||||
|
"permissions": "drwxr-xr-x",
|
||||||
|
"owner": "root",
|
||||||
|
"group": "root",
|
||||||
|
"modified_at": 1712899200,
|
||||||
|
"symlink_target": None,
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
entry = cap.files.make_dir("/home/user/data")
|
||||||
|
assert isinstance(entry, FileEntry)
|
||||||
|
assert entry.name == "data"
|
||||||
|
assert entry.type == "directory"
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_make_dir_existing_returns_gracefully(self):
|
||||||
|
cap = _make_capsule()
|
||||||
|
respx.post(f"{BASE}/v1/capsules/cl-abc/files/mkdir").respond(
|
||||||
|
409,
|
||||||
|
json={"error": {"code": "conflict", "message": "already exists"}},
|
||||||
|
)
|
||||||
|
respx.post(f"{BASE}/v1/capsules/cl-abc/files/list").respond(
|
||||||
|
200,
|
||||||
|
json={
|
||||||
|
"entries": [
|
||||||
|
{
|
||||||
|
"name": "data",
|
||||||
|
"path": "/home/user/data",
|
||||||
|
"type": "directory",
|
||||||
|
"size": 4096,
|
||||||
|
"mode": 16877,
|
||||||
|
"permissions": "drwxr-xr-x",
|
||||||
|
"owner": "root",
|
||||||
|
"group": "root",
|
||||||
|
"modified_at": 1712899200,
|
||||||
|
"symlink_target": None,
|
||||||
|
}
|
||||||
|
]
|
||||||
|
},
|
||||||
|
)
|
||||||
|
entry = cap.files.make_dir("/home/user/data")
|
||||||
|
assert entry.name == "data"
|
||||||
|
|
||||||
|
|
||||||
|
class TestFilesRemove:
|
||||||
|
@respx.mock
|
||||||
|
def test_remove_succeeds(self):
|
||||||
|
cap = _make_capsule()
|
||||||
|
route = respx.post(f"{BASE}/v1/capsules/cl-abc/files/remove").respond(204)
|
||||||
|
cap.files.remove("/home/user/old_data")
|
||||||
|
assert route.called
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_remove_sends_path(self):
|
||||||
|
cap = _make_capsule()
|
||||||
|
route = respx.post(f"{BASE}/v1/capsules/cl-abc/files/remove").respond(204)
|
||||||
|
cap.files.remove("/tmp/test.txt")
|
||||||
|
body = json.loads(route.calls[0].request.content)
|
||||||
|
assert body["path"] == "/tmp/test.txt"
|
||||||
|
|
||||||
|
|
||||||
|
class TestFilesExists:
|
||||||
|
@respx.mock
|
||||||
|
def test_exists_true(self):
|
||||||
|
cap = _make_capsule()
|
||||||
|
respx.post(f"{BASE}/v1/capsules/cl-abc/files/list").respond(
|
||||||
|
200,
|
||||||
|
json={
|
||||||
|
"entries": [
|
||||||
|
{"name": "hello.txt", "path": "/tmp/hello.txt", "type": "file"}
|
||||||
|
]
|
||||||
|
},
|
||||||
|
)
|
||||||
|
assert cap.files.exists("/tmp/hello.txt") is True
|
||||||
|
|
||||||
|
@respx.mock
|
||||||
|
def test_exists_false(self):
|
||||||
|
cap = _make_capsule()
|
||||||
|
respx.post(f"{BASE}/v1/capsules/cl-abc/files/list").respond(
|
||||||
|
200, json={"entries": []}
|
||||||
|
)
|
||||||
|
assert cap.files.exists("/tmp/nope.txt") is False
|
||||||
|
|
||||||
|
|
||||||
|
class TestPtyEventParsing:
|
||||||
|
def test_started_event(self):
|
||||||
|
raw = {"type": "started", "tag": "pty-a1b2c3d4", "pid": 42}
|
||||||
|
event = _parse_pty_event(raw)
|
||||||
|
assert event.type == PtyEventType.started
|
||||||
|
assert event.pid == 42
|
||||||
|
assert event.tag == "pty-a1b2c3d4"
|
||||||
|
|
||||||
|
def test_output_event_base64(self):
|
||||||
|
encoded = base64.b64encode(b"ls -la\n").decode()
|
||||||
|
raw = {"type": "output", "data": encoded}
|
||||||
|
event = _parse_pty_event(raw)
|
||||||
|
assert event.type == PtyEventType.output
|
||||||
|
assert event.data == b"ls -la\n"
|
||||||
|
|
||||||
|
def test_output_event_empty(self):
|
||||||
|
raw = {"type": "output", "data": ""}
|
||||||
|
event = _parse_pty_event(raw)
|
||||||
|
assert event.data == b""
|
||||||
|
|
||||||
|
def test_exit_event(self):
|
||||||
|
raw = {"type": "exit", "exit_code": 0}
|
||||||
|
event = _parse_pty_event(raw)
|
||||||
|
assert event.type == PtyEventType.exit
|
||||||
|
assert event.exit_code == 0
|
||||||
|
|
||||||
|
def test_error_event(self):
|
||||||
|
raw = {"type": "error", "data": "process not found", "fatal": True}
|
||||||
|
event = _parse_pty_event(raw)
|
||||||
|
assert event.type == PtyEventType.error
|
||||||
|
assert event.data == "process not found"
|
||||||
|
assert event.fatal is True
|
||||||
|
|
||||||
|
def test_ping_event(self):
|
||||||
|
raw = {"type": "ping"}
|
||||||
|
event = _parse_pty_event(raw)
|
||||||
|
assert event.type == PtyEventType.ping
|
||||||
|
|
||||||
|
|
||||||
|
class TestPtySessionWrite:
|
||||||
|
def test_write_sends_base64_input(self):
|
||||||
|
ws = MagicMock()
|
||||||
|
session = PtySession(ws, "cl-abc")
|
||||||
|
session.write(b"ls -la\n")
|
||||||
|
sent = json.loads(ws.send_text.call_args[0][0])
|
||||||
|
assert sent["type"] == "input"
|
||||||
|
assert base64.b64decode(sent["data"]) == b"ls -la\n"
|
||||||
|
|
||||||
|
|
||||||
|
class TestPtySessionResize:
|
||||||
|
def test_resize_sends_dimensions(self):
|
||||||
|
ws = MagicMock()
|
||||||
|
session = PtySession(ws, "cl-abc")
|
||||||
|
session.resize(120, 40)
|
||||||
|
sent = json.loads(ws.send_text.call_args[0][0])
|
||||||
|
assert sent["type"] == "resize"
|
||||||
|
assert sent["cols"] == 120
|
||||||
|
assert sent["rows"] == 40
|
||||||
|
|
||||||
|
def test_resize_zero_raises(self):
|
||||||
|
ws = MagicMock()
|
||||||
|
session = PtySession(ws, "cl-abc")
|
||||||
|
with pytest.raises(ValueError, match="greater than 0"):
|
||||||
|
session.resize(0, 40)
|
||||||
|
with pytest.raises(ValueError, match="greater than 0"):
|
||||||
|
session.resize(80, 0)
|
||||||
|
|
||||||
|
|
||||||
|
class TestPtySessionKill:
|
||||||
|
def test_kill_sends_message(self):
|
||||||
|
ws = MagicMock()
|
||||||
|
session = PtySession(ws, "cl-abc")
|
||||||
|
session.kill()
|
||||||
|
sent = json.loads(ws.send_text.call_args[0][0])
|
||||||
|
assert sent["type"] == "kill"
|
||||||
|
|
||||||
|
|
||||||
|
class TestPtySessionIteration:
|
||||||
|
def test_iter_yields_events_until_exit(self):
|
||||||
|
ws = MagicMock()
|
||||||
|
messages = [
|
||||||
|
json.dumps({"type": "started", "tag": "pty-abc12345", "pid": 1}),
|
||||||
|
json.dumps(
|
||||||
|
{"type": "output", "data": base64.b64encode(b"hello").decode()}
|
||||||
|
),
|
||||||
|
json.dumps({"type": "exit", "exit_code": 0}),
|
||||||
|
]
|
||||||
|
ws.receive_text.side_effect = messages
|
||||||
|
session = PtySession(ws, "cl-abc")
|
||||||
|
events = list(session)
|
||||||
|
assert len(events) == 2
|
||||||
|
assert events[0].type == PtyEventType.started
|
||||||
|
assert session.tag == "pty-abc12345"
|
||||||
|
assert session.pid == 1
|
||||||
|
assert events[1].type == PtyEventType.output
|
||||||
|
assert events[1].data == b"hello"
|
||||||
|
|
||||||
|
def test_iter_stops_on_fatal_error(self):
|
||||||
|
ws = MagicMock()
|
||||||
|
messages = [
|
||||||
|
json.dumps({"type": "error", "data": "fatal", "fatal": True}),
|
||||||
|
]
|
||||||
|
ws.receive_text.side_effect = messages
|
||||||
|
session = PtySession(ws, "cl-abc")
|
||||||
|
events = list(session)
|
||||||
|
assert len(events) == 1
|
||||||
|
assert events[0].type == PtyEventType.error
|
||||||
|
|
||||||
|
def test_iter_stops_on_disconnect(self):
|
||||||
|
import httpx_ws
|
||||||
|
|
||||||
|
ws = MagicMock()
|
||||||
|
ws.receive_text.side_effect = httpx_ws.WebSocketDisconnect()
|
||||||
|
session = PtySession(ws, "cl-abc")
|
||||||
|
events = list(session)
|
||||||
|
assert events == []
|
||||||
|
|
||||||
|
|
||||||
|
class TestPtySessionContextManager:
|
||||||
|
def test_exit_kills_and_closes(self):
|
||||||
|
ws = MagicMock()
|
||||||
|
session = PtySession(ws, "cl-abc")
|
||||||
|
with session:
|
||||||
|
pass
|
||||||
|
ws.send_text.assert_called()
|
||||||
|
ws.close.assert_called()
|
||||||
|
|
||||||
|
def test_exit_ignores_errors(self):
|
||||||
|
ws = MagicMock()
|
||||||
|
ws.send_text.side_effect = Exception("already closed")
|
||||||
|
session = PtySession(ws, "cl-abc")
|
||||||
|
with session:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class TestPtySessionSendStart:
|
||||||
|
def test_send_start_with_defaults(self):
|
||||||
|
ws = MagicMock()
|
||||||
|
session = PtySession(ws, "cl-abc")
|
||||||
|
session._send_start()
|
||||||
|
sent = json.loads(ws.send_text.call_args[0][0])
|
||||||
|
assert sent["type"] == "start"
|
||||||
|
assert sent["cmd"] == "/bin/bash"
|
||||||
|
assert sent["cols"] == 80
|
||||||
|
assert sent["rows"] == 24
|
||||||
|
|
||||||
|
def test_send_start_with_all_params(self):
|
||||||
|
ws = MagicMock()
|
||||||
|
session = PtySession(ws, "cl-abc")
|
||||||
|
session._send_start(
|
||||||
|
cmd="/bin/zsh",
|
||||||
|
args=["-l"],
|
||||||
|
cols=120,
|
||||||
|
rows=40,
|
||||||
|
envs={"TERM": "xterm-256color"},
|
||||||
|
cwd="/home/user",
|
||||||
|
)
|
||||||
|
sent = json.loads(ws.send_text.call_args[0][0])
|
||||||
|
assert sent["cmd"] == "/bin/zsh"
|
||||||
|
assert sent["args"] == ["-l"]
|
||||||
|
assert sent["cols"] == 120
|
||||||
|
|
||||||
|
|
||||||
|
class TestPtySessionSendConnect:
|
||||||
|
def test_send_connect(self):
|
||||||
|
ws = MagicMock()
|
||||||
|
session = PtySession(ws, "cl-abc")
|
||||||
|
session._send_connect("pty-abc12345")
|
||||||
|
sent = json.loads(ws.send_text.call_args[0][0])
|
||||||
|
assert sent["type"] == "connect"
|
||||||
|
assert sent["tag"] == "pty-abc12345"
|
||||||
|
|
||||||
|
|
||||||
|
class TestAsyncPtySession:
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_write_sends_base64(self):
|
||||||
|
ws = AsyncMock()
|
||||||
|
session = AsyncPtySession(ws, "cl-abc")
|
||||||
|
await session.write(b"hello")
|
||||||
|
sent = json.loads(ws.send_text.call_args[0][0])
|
||||||
|
assert sent["type"] == "input"
|
||||||
|
assert base64.b64decode(sent["data"]) == b"hello"
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_resize(self):
|
||||||
|
ws = AsyncMock()
|
||||||
|
session = AsyncPtySession(ws, "cl-abc")
|
||||||
|
await session.resize(100, 30)
|
||||||
|
sent = json.loads(ws.send_text.call_args[0][0])
|
||||||
|
assert sent["type"] == "resize"
|
||||||
|
assert sent["cols"] == 100
|
||||||
|
assert sent["rows"] == 30
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_resize_zero_raises(self):
|
||||||
|
ws = AsyncMock()
|
||||||
|
session = AsyncPtySession(ws, "cl-abc")
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
await session.resize(0, 10)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_kill(self):
|
||||||
|
ws = AsyncMock()
|
||||||
|
session = AsyncPtySession(ws, "cl-abc")
|
||||||
|
await session.kill()
|
||||||
|
sent = json.loads(ws.send_text.call_args[0][0])
|
||||||
|
assert sent["type"] == "kill"
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_context_manager(self):
|
||||||
|
ws = AsyncMock()
|
||||||
|
session = AsyncPtySession(ws, "cl-abc")
|
||||||
|
async with session:
|
||||||
|
pass
|
||||||
|
ws.send_text.assert_called()
|
||||||
|
ws.close.assert_called()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_send_start(self):
|
||||||
|
ws = AsyncMock()
|
||||||
|
session = AsyncPtySession(ws, "cl-abc")
|
||||||
|
await session._send_start(cmd="/bin/zsh", cols=100, rows=30)
|
||||||
|
sent = json.loads(ws.send_text.call_args[0][0])
|
||||||
|
assert sent["type"] == "start"
|
||||||
|
assert sent["cmd"] == "/bin/zsh"
|
||||||
|
assert sent["cols"] == 100
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_iteration(self):
|
||||||
|
ws = AsyncMock()
|
||||||
|
messages = [
|
||||||
|
json.dumps({"type": "started", "tag": "pty-xyz", "pid": 5}),
|
||||||
|
json.dumps(
|
||||||
|
{"type": "output", "data": base64.b64encode(b"hi").decode()}
|
||||||
|
),
|
||||||
|
json.dumps({"type": "exit", "exit_code": 0}),
|
||||||
|
]
|
||||||
|
ws.receive_text.side_effect = messages
|
||||||
|
session = AsyncPtySession(ws, "cl-abc")
|
||||||
|
events = []
|
||||||
|
async for event in session:
|
||||||
|
events.append(event)
|
||||||
|
assert len(events) == 2
|
||||||
|
assert events[0].type == PtyEventType.started
|
||||||
|
assert session.tag == "pty-xyz"
|
||||||
|
assert session.pid == 5
|
||||||
|
|
||||||
|
|
||||||
|
class TestExports:
|
||||||
|
def test_file_entry_importable(self):
|
||||||
|
from wrenn import FileEntry as FE
|
||||||
|
|
||||||
|
assert FE is not None
|
||||||
|
|
||||||
|
def test_pty_session_importable(self):
|
||||||
|
from wrenn import PtySession as PS
|
||||||
|
|
||||||
|
assert PS is not None
|
||||||
|
|
||||||
|
def test_async_pty_session_importable(self):
|
||||||
|
from wrenn import AsyncPtySession as APS
|
||||||
|
|
||||||
|
assert APS is not None
|
||||||
|
|
||||||
|
def test_pty_event_importable(self):
|
||||||
|
from wrenn import PtyEvent as PE
|
||||||
|
from wrenn import PtyEventType as PET
|
||||||
|
|
||||||
|
assert PE is not None
|
||||||
|
assert PET is not None
|
||||||
568
tests/test_integration.py
Normal file
568
tests/test_integration.py
Normal file
@ -0,0 +1,568 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
from typing import Generator
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from wrenn.client import AsyncWrennClient, WrennClient
|
||||||
|
from wrenn.exceptions import WrennNotFoundError, WrennValidationError
|
||||||
|
from wrenn.pty import PtyEventType
|
||||||
|
|
||||||
|
WRENN_API_KEY = os.environ.get("WRENN_API_KEY")
|
||||||
|
WRENN_TOKEN = os.environ.get("WRENN_TOKEN")
|
||||||
|
WRENN_BASE_URL = os.environ.get("WRENN_BASE_URL", "http://localhost:8080")
|
||||||
|
WRENN_TEST_EMAIL = os.environ.get("WRENN_TEST_EMAIL")
|
||||||
|
WRENN_TEST_PASSWORD = os.environ.get("WRENN_TEST_PASSWORD")
|
||||||
|
|
||||||
|
|
||||||
|
def _has_auth() -> bool:
|
||||||
|
return bool(WRENN_API_KEY or WRENN_TOKEN)
|
||||||
|
|
||||||
|
|
||||||
|
requires_auth = pytest.mark.skipif(
|
||||||
|
not _has_auth(),
|
||||||
|
reason="Set WRENN_API_KEY or WRENN_TOKEN to run integration tests",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def client() -> Generator[WrennClient, None, None]:
|
||||||
|
with WrennClient(
|
||||||
|
api_key=WRENN_API_KEY,
|
||||||
|
token=WRENN_TOKEN,
|
||||||
|
base_url=WRENN_BASE_URL,
|
||||||
|
) as c:
|
||||||
|
yield c
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def async_client() -> AsyncWrennClient:
|
||||||
|
return AsyncWrennClient(
|
||||||
|
api_key=WRENN_API_KEY,
|
||||||
|
token=WRENN_TOKEN,
|
||||||
|
base_url=WRENN_BASE_URL,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def bearer_client() -> Generator[WrennClient, None, None]:
|
||||||
|
if WRENN_TOKEN:
|
||||||
|
with WrennClient(token=WRENN_TOKEN, base_url=WRENN_BASE_URL) as c:
|
||||||
|
yield c
|
||||||
|
elif WRENN_TEST_EMAIL and WRENN_TEST_PASSWORD:
|
||||||
|
with WrennClient(
|
||||||
|
api_key=WRENN_API_KEY, token=WRENN_TOKEN, base_url=WRENN_BASE_URL
|
||||||
|
) as c:
|
||||||
|
resp = c.auth.login(WRENN_TEST_EMAIL, WRENN_TEST_PASSWORD)
|
||||||
|
with WrennClient(token=resp.token, base_url=WRENN_BASE_URL) as c:
|
||||||
|
yield c
|
||||||
|
else:
|
||||||
|
pytest.skip(
|
||||||
|
"Set WRENN_TOKEN or WRENN_TEST_EMAIL+WRENN_TEST_PASSWORD for bearer-auth tests"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestCapsuleLifecycle:
|
||||||
|
def test_create_exec_destroy(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
result = cap.exec("echo", args=["hello"])
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "hello" in result.stdout
|
||||||
|
|
||||||
|
def test_exec_with_args(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
result = cap.exec("echo", args=["hello", "world"])
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "hello world" in result.stdout
|
||||||
|
|
||||||
|
def test_exec_nonzero_exit(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
result = cap.exec("sh", args=["-c", "exit 42"])
|
||||||
|
assert result.exit_code == 42
|
||||||
|
|
||||||
|
def test_exec_stderr(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
result = cap.exec("sh", args=["-c", "echo err>&2"])
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "err" in result.stderr
|
||||||
|
|
||||||
|
def test_context_manager_cleanup(self, client):
|
||||||
|
cap = client.capsules.create(template="minimal", timeout_sec=120)
|
||||||
|
cap_id = cap.id
|
||||||
|
|
||||||
|
with cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
|
||||||
|
fetched = client.capsules.get(cap_id)
|
||||||
|
assert fetched.status in ("stopped", "destroyed")
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestFileIO:
|
||||||
|
def test_upload_and_download(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
content = b"Hello from integration test!"
|
||||||
|
cap.upload("/tmp/test_file.txt", content)
|
||||||
|
downloaded = cap.download("/tmp/test_file.txt")
|
||||||
|
assert downloaded == content
|
||||||
|
|
||||||
|
def test_download_nonexistent_file(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
with pytest.raises(Exception):
|
||||||
|
cap.download("/tmp/no_such_file_12345")
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestPauseResume:
|
||||||
|
def test_pause_and_resume(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
cap.pause()
|
||||||
|
assert cap.status == "paused"
|
||||||
|
|
||||||
|
cap.resume()
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
|
||||||
|
result = cap.exec("echo", args=["resumed"])
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "resumed" in result.stdout
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestPing:
|
||||||
|
def test_ping_resets_timer(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
cap.ping()
|
||||||
|
result = cap.exec("echo", args=["still_alive"])
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "still_alive" in result.stdout
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestProxy:
|
||||||
|
def test_get_url(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
url = cap.get_url(8888)
|
||||||
|
assert cap.id in url
|
||||||
|
assert "8888" in url
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestListAndGet:
|
||||||
|
def test_list_capsules(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
boxes = client.capsules.list()
|
||||||
|
ids = [b.id for b in boxes]
|
||||||
|
assert cap.id in ids
|
||||||
|
|
||||||
|
def test_get_existing_capsule(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
fetched = client.capsules.get(cap.id)
|
||||||
|
assert fetched.id == cap.id
|
||||||
|
assert fetched.status == "running"
|
||||||
|
|
||||||
|
def test_get_nonexistent_capsule(self, client):
|
||||||
|
with pytest.raises((WrennNotFoundError, WrennValidationError)):
|
||||||
|
client.capsules.get("cl-nonexistent00000000000000000")
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestSnapshots:
|
||||||
|
def test_list_templates(self, client):
|
||||||
|
templates = client.snapshots.list()
|
||||||
|
assert isinstance(templates, list)
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestAPIKeys:
|
||||||
|
def test_create_list_delete(self, bearer_client):
|
||||||
|
key_resp = bearer_client.api_keys.create(name="integration-test-key")
|
||||||
|
assert key_resp.name == "integration-test-key"
|
||||||
|
assert key_resp.key is not None
|
||||||
|
assert key_resp.id is not None
|
||||||
|
|
||||||
|
try:
|
||||||
|
keys = bearer_client.api_keys.list()
|
||||||
|
ids = [k.id for k in keys]
|
||||||
|
assert key_resp.id in ids
|
||||||
|
finally:
|
||||||
|
bearer_client.api_keys.delete(key_resp.id)
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestRunCode:
|
||||||
|
def test_basic_execution(self, client):
|
||||||
|
with client.capsules.create(
|
||||||
|
template="python-interpreter-v0-beta", timeout_sec=120
|
||||||
|
) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
|
||||||
|
r = cap.run_code("x = 42")
|
||||||
|
assert r.error is None
|
||||||
|
|
||||||
|
r = cap.run_code("x * 2")
|
||||||
|
assert r.text == "84"
|
||||||
|
|
||||||
|
def test_state_persists(self, client):
|
||||||
|
with client.capsules.create(
|
||||||
|
template="python-interpreter-v0-beta", timeout_sec=120
|
||||||
|
) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
|
||||||
|
cap.run_code("def greet(name): return f'hello {name}'")
|
||||||
|
r = cap.run_code("greet('capsule')")
|
||||||
|
assert "hello capsule" in (r.text or "")
|
||||||
|
|
||||||
|
def test_error_traceback(self, client):
|
||||||
|
with client.capsules.create(
|
||||||
|
template="python-interpreter-v0-beta", timeout_sec=120
|
||||||
|
) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
|
||||||
|
r = cap.run_code("1/0")
|
||||||
|
assert r.error is not None
|
||||||
|
assert "ZeroDivisionError" in r.error
|
||||||
|
|
||||||
|
def test_stdout_capture(self, client):
|
||||||
|
with client.capsules.create(
|
||||||
|
template="python-interpreter-v0-beta", timeout_sec=120
|
||||||
|
) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
|
||||||
|
r = cap.run_code("print('hello from kernel')")
|
||||||
|
assert "hello from kernel" in r.stdout
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestAsyncCapsuleLifecycle:
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_create_exec_destroy(self, async_client):
|
||||||
|
async with async_client:
|
||||||
|
cap = await async_client.capsules.create(
|
||||||
|
template="minimal", timeout_sec=120
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
await cap.async_wait_ready(timeout=60, interval=1)
|
||||||
|
result = await cap.async_exec("echo", args=["async_hello"])
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "async_hello" in result.stdout
|
||||||
|
finally:
|
||||||
|
await cap.async_destroy()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_upload_download(self, async_client):
|
||||||
|
async with async_client:
|
||||||
|
cap = await async_client.capsules.create(
|
||||||
|
template="minimal", timeout_sec=120
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
await cap.async_wait_ready(timeout=60, interval=1)
|
||||||
|
content = b"Async upload test"
|
||||||
|
await cap.async_upload("/tmp/async_test.txt", content)
|
||||||
|
downloaded = await cap.async_download("/tmp/async_test.txt")
|
||||||
|
assert downloaded == content
|
||||||
|
finally:
|
||||||
|
await cap.async_destroy()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_run_code(self, async_client):
|
||||||
|
async with async_client:
|
||||||
|
cap = await async_client.capsules.create(
|
||||||
|
template="python-interpreter-v0-beta", timeout_sec=120
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
await cap.async_wait_ready(timeout=60, interval=1)
|
||||||
|
r = await cap.async_run_code("42 * 2")
|
||||||
|
assert r.text == "84"
|
||||||
|
finally:
|
||||||
|
await cap.async_destroy()
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestFilesystemListDir:
|
||||||
|
def test_list_dir_root(self, client: WrennClient):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
cap.mkdir("/tmp/ls_test_root")
|
||||||
|
cap.upload("/tmp/ls_test_root/hello.txt", b"hello")
|
||||||
|
entries = cap.list_dir("/tmp/ls_test_root")
|
||||||
|
assert isinstance(entries, list)
|
||||||
|
names = [e.name for e in entries]
|
||||||
|
assert "hello.txt" in names
|
||||||
|
|
||||||
|
def test_list_dir_after_mkdir(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
cap.mkdir("/tmp/fs_test_dir")
|
||||||
|
entries = cap.list_dir("/tmp")
|
||||||
|
names = [e.name for e in entries]
|
||||||
|
assert "fs_test_dir" in names
|
||||||
|
|
||||||
|
def test_list_dir_file_metadata(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
cap.upload("/tmp/meta_test.txt", b"hello world")
|
||||||
|
entries = cap.list_dir("/tmp")
|
||||||
|
match = [e for e in entries if e.name == "meta_test.txt"]
|
||||||
|
assert len(match) == 1
|
||||||
|
f = match[0]
|
||||||
|
assert f.type == "file"
|
||||||
|
assert f.size == 11
|
||||||
|
assert f.permissions is not None
|
||||||
|
assert f.owner is not None
|
||||||
|
assert f.group is not None
|
||||||
|
assert f.modified_at is not None
|
||||||
|
|
||||||
|
def test_list_dir_depth(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
cap.mkdir("/tmp/depth_a/depth_b")
|
||||||
|
cap.upload("/tmp/depth_a/depth_b/nested.txt", b"deep")
|
||||||
|
entries = cap.list_dir("/tmp/depth_a", depth=2)
|
||||||
|
paths = [e.path for e in entries]
|
||||||
|
assert any("nested.txt" in p for p in paths)
|
||||||
|
|
||||||
|
def test_list_dir_empty_directory(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
cap.mkdir("/tmp/empty_dir_test")
|
||||||
|
entries = cap.list_dir("/tmp/empty_dir_test")
|
||||||
|
assert entries == []
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestFilesystemMkdir:
|
||||||
|
def test_mkdir_creates_directory(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
entry = cap.mkdir("/tmp/mkdir_test")
|
||||||
|
assert entry.name == "mkdir_test"
|
||||||
|
assert entry.type == "directory"
|
||||||
|
assert entry.path == "/tmp/mkdir_test"
|
||||||
|
|
||||||
|
def test_mkdir_creates_parents(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
entry = cap.mkdir("/tmp/a/b/c/d")
|
||||||
|
assert entry.type == "directory"
|
||||||
|
|
||||||
|
def test_mkdir_already_exists(self, client: WrennClient):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
cap.mkdir("/tmp/exist_test")
|
||||||
|
entry = cap.mkdir("/tmp/exist_test")
|
||||||
|
assert entry.type == "directory"
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestFilesystemRemove:
|
||||||
|
def test_remove_file(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
cap.upload("/tmp/rm_test.txt", b"delete me")
|
||||||
|
entries_before = cap.list_dir("/tmp")
|
||||||
|
assert any(e.name == "rm_test.txt" for e in entries_before)
|
||||||
|
cap.remove("/tmp/rm_test.txt")
|
||||||
|
entries_after = cap.list_dir("/tmp")
|
||||||
|
assert not any(e.name == "rm_test.txt" for e in entries_after)
|
||||||
|
|
||||||
|
def test_remove_directory(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
cap.mkdir("/tmp/rm_dir_test")
|
||||||
|
cap.upload("/tmp/rm_dir_test/file.txt", b"inside")
|
||||||
|
cap.remove("/tmp/rm_dir_test")
|
||||||
|
entries = cap.list_dir("/tmp")
|
||||||
|
assert not any(e.name == "rm_dir_test" for e in entries)
|
||||||
|
|
||||||
|
def test_upload_download_remove_roundtrip(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
content = b"round trip test data " * 100
|
||||||
|
cap.upload("/tmp/rt.txt", content)
|
||||||
|
downloaded = cap.download("/tmp/rt.txt")
|
||||||
|
assert downloaded == content
|
||||||
|
cap.remove("/tmp/rt.txt")
|
||||||
|
with pytest.raises(Exception):
|
||||||
|
cap.download("/tmp/rt.txt")
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestStreamUploadDownload:
|
||||||
|
def test_stream_upload_and_download(self, client: WrennClient):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
chunks = [b"chunk0_", b"chunk1_", b"chunk2"]
|
||||||
|
|
||||||
|
def data_gen():
|
||||||
|
yield from chunks
|
||||||
|
|
||||||
|
cap.stream_upload("/tmp/stream_test.bin", data_gen())
|
||||||
|
downloaded = cap.download("/tmp/stream_test.bin")
|
||||||
|
assert downloaded == b"chunk0_chunk1_chunk2"
|
||||||
|
|
||||||
|
def test_stream_download_large(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
content = b"x" * 65536 * 3
|
||||||
|
cap.upload("/tmp/large.bin", content)
|
||||||
|
collected = b""
|
||||||
|
for chunk in cap.stream_download("/tmp/large.bin"):
|
||||||
|
collected += chunk
|
||||||
|
assert collected == content
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestPty:
|
||||||
|
def test_pty_basic_output(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
with cap.pty(cmd="/bin/sh", cwd="/tmp") as term:
|
||||||
|
term.write(b"echo pty_hello\n")
|
||||||
|
output = b""
|
||||||
|
for event in term:
|
||||||
|
if event.type == PtyEventType.output:
|
||||||
|
output += event.data
|
||||||
|
elif event.type == PtyEventType.exit:
|
||||||
|
break
|
||||||
|
if b"pty_hello" in output:
|
||||||
|
term.write(b"exit\n")
|
||||||
|
assert b"pty_hello" in output
|
||||||
|
|
||||||
|
def test_pty_tag_and_pid(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
with cap.pty(cmd="/bin/sh") as term:
|
||||||
|
started = False
|
||||||
|
for event in term:
|
||||||
|
if event.type == PtyEventType.started:
|
||||||
|
started = True
|
||||||
|
assert term.tag is not None
|
||||||
|
assert term.pid is not None
|
||||||
|
assert term.tag.startswith("pty-")
|
||||||
|
elif event.type == PtyEventType.output:
|
||||||
|
term.write(b"exit\n")
|
||||||
|
elif event.type == PtyEventType.exit:
|
||||||
|
break
|
||||||
|
assert started
|
||||||
|
|
||||||
|
def test_pty_exit_on_command_exit(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
with cap.pty(cmd="/bin/echo", args=["immediate"]) as term:
|
||||||
|
events = list(term)
|
||||||
|
types = [e.type for e in events]
|
||||||
|
assert PtyEventType.started in types
|
||||||
|
assert PtyEventType.output in types or PtyEventType.exit in types
|
||||||
|
|
||||||
|
def test_pty_resize(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
with cap.pty(cmd="/bin/sh", cols=80, rows=24) as term:
|
||||||
|
for event in term:
|
||||||
|
if event.type == PtyEventType.started:
|
||||||
|
term.resize(120, 40)
|
||||||
|
term.write(b"exit\n")
|
||||||
|
elif event.type == PtyEventType.exit:
|
||||||
|
break
|
||||||
|
|
||||||
|
def test_pty_envs(self, client):
|
||||||
|
with client.capsules.create(template="minimal", timeout_sec=120) as cap:
|
||||||
|
cap.wait_ready(timeout=60, interval=1)
|
||||||
|
with cap.pty(cmd="/bin/sh", envs={"MY_VAR": "hello_env"}) as term:
|
||||||
|
output = b""
|
||||||
|
for event in term:
|
||||||
|
if event.type == PtyEventType.started:
|
||||||
|
term.write(b"echo $MY_VAR\n")
|
||||||
|
elif event.type == PtyEventType.output:
|
||||||
|
output += event.data
|
||||||
|
if b"hello_env" in output:
|
||||||
|
term.write(b"exit\n")
|
||||||
|
elif event.type == PtyEventType.exit:
|
||||||
|
break
|
||||||
|
assert b"hello_env" in output
|
||||||
|
|
||||||
|
|
||||||
|
@requires_auth
|
||||||
|
class TestAsyncFilesystem:
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_list_dir(self, async_client):
|
||||||
|
async with async_client:
|
||||||
|
cap = await async_client.capsules.create(
|
||||||
|
template="minimal", timeout_sec=120
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
await cap.async_wait_ready(timeout=60, interval=1)
|
||||||
|
await cap.async_mkdir("/tmp/async_ls_test")
|
||||||
|
await cap.async_upload("/tmp/async_ls_test/file.txt", b"data")
|
||||||
|
entries = await cap.async_list_dir("/tmp/async_ls_test")
|
||||||
|
assert isinstance(entries, list)
|
||||||
|
assert any(e.name == "file.txt" for e in entries)
|
||||||
|
finally:
|
||||||
|
await cap.async_destroy()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_mkdir(self, async_client):
|
||||||
|
async with async_client:
|
||||||
|
cap = await async_client.capsules.create(
|
||||||
|
template="minimal", timeout_sec=120
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
await cap.async_wait_ready(timeout=60, interval=1)
|
||||||
|
entry = await cap.async_mkdir("/tmp/async_mkdir_test")
|
||||||
|
assert entry.type == "directory"
|
||||||
|
assert entry.name == "async_mkdir_test"
|
||||||
|
finally:
|
||||||
|
await cap.async_destroy()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_remove(self, async_client):
|
||||||
|
async with async_client:
|
||||||
|
cap = await async_client.capsules.create(
|
||||||
|
template="minimal", timeout_sec=120
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
await cap.async_wait_ready(timeout=60, interval=1)
|
||||||
|
await cap.async_upload("/tmp/async_rm.txt", b"bye")
|
||||||
|
entries = await cap.async_list_dir("/tmp")
|
||||||
|
assert any(e.name == "async_rm.txt" for e in entries)
|
||||||
|
await cap.async_remove("/tmp/async_rm.txt")
|
||||||
|
entries = await cap.async_list_dir("/tmp")
|
||||||
|
assert not any(e.name == "async_rm.txt" for e in entries)
|
||||||
|
finally:
|
||||||
|
await cap.async_destroy()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_full_filesystem_roundtrip(self, async_client):
|
||||||
|
async with async_client:
|
||||||
|
cap = await async_client.capsules.create(
|
||||||
|
template="minimal", timeout_sec=120
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
await cap.async_wait_ready(timeout=60, interval=1)
|
||||||
|
|
||||||
|
await cap.async_mkdir("/tmp/async_rt")
|
||||||
|
await cap.async_upload("/tmp/async_rt/file.txt", b"async content")
|
||||||
|
entries = await cap.async_list_dir("/tmp/async_rt")
|
||||||
|
assert any(e.name == "file.txt" for e in entries)
|
||||||
|
|
||||||
|
data = await cap.async_download("/tmp/async_rt/file.txt")
|
||||||
|
assert data == b"async content"
|
||||||
|
|
||||||
|
await cap.async_remove("/tmp/async_rt/file.txt")
|
||||||
|
entries = await cap.async_list_dir("/tmp/async_rt")
|
||||||
|
assert not any(e.name == "file.txt" for e in entries)
|
||||||
|
finally:
|
||||||
|
await cap.async_destroy()
|
||||||
89
uv.lock
generated
89
uv.lock
generated
@ -1,5 +1,5 @@
|
|||||||
version = 1
|
version = 1
|
||||||
revision = 2
|
revision = 3
|
||||||
requires-python = ">=3.13"
|
requires-python = ">=3.13"
|
||||||
resolution-markers = [
|
resolution-markers = [
|
||||||
"python_full_version >= '3.14'",
|
"python_full_version >= '3.14'",
|
||||||
@ -112,6 +112,33 @@ wheels = [
|
|||||||
{ url = "https://files.pythonhosted.org/packages/ed/3a/7f169ffc7a2d69a4f9158b1ac083f685b7f4a1a8a1db5d1e4abbb4e741b7/datamodel_code_generator-0.56.0-py3-none-any.whl", hash = "sha256:a0559683fbe90cdf2ce9b6637e3adae3e3a8056a8d0516df581d486e2834ead2", size = 256545, upload-time = "2026-04-04T09:46:17.582Z" },
|
{ url = "https://files.pythonhosted.org/packages/ed/3a/7f169ffc7a2d69a4f9158b1ac083f685b7f4a1a8a1db5d1e4abbb4e741b7/datamodel_code_generator-0.56.0-py3-none-any.whl", hash = "sha256:a0559683fbe90cdf2ce9b6637e3adae3e3a8056a8d0516df581d486e2834ead2", size = 256545, upload-time = "2026-04-04T09:46:17.582Z" },
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[package.optional-dependencies]
|
||||||
|
ruff = [
|
||||||
|
{ name = "ruff" },
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "dnspython"
|
||||||
|
version = "2.8.0"
|
||||||
|
source = { registry = "https://pypi.org/simple" }
|
||||||
|
sdist = { url = "https://files.pythonhosted.org/packages/8c/8b/57666417c0f90f08bcafa776861060426765fdb422eb10212086fb811d26/dnspython-2.8.0.tar.gz", hash = "sha256:181d3c6996452cb1189c4046c61599b84a5a86e099562ffde77d26984ff26d0f", size = 368251, upload-time = "2025-09-07T18:58:00.022Z" }
|
||||||
|
wheels = [
|
||||||
|
{ url = "https://files.pythonhosted.org/packages/ba/5a/18ad964b0086c6e62e2e7500f7edc89e3faa45033c71c1893d34eed2b2de/dnspython-2.8.0-py3-none-any.whl", hash = "sha256:01d9bbc4a2d76bf0db7c1f729812ded6d912bd318d3b1cf81d30c0f845dbf3af", size = 331094, upload-time = "2025-09-07T18:57:58.071Z" },
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "email-validator"
|
||||||
|
version = "2.3.0"
|
||||||
|
source = { registry = "https://pypi.org/simple" }
|
||||||
|
dependencies = [
|
||||||
|
{ name = "dnspython" },
|
||||||
|
{ name = "idna" },
|
||||||
|
]
|
||||||
|
sdist = { url = "https://files.pythonhosted.org/packages/f5/22/900cb125c76b7aaa450ce02fd727f452243f2e91a61af068b40adba60ea9/email_validator-2.3.0.tar.gz", hash = "sha256:9fc05c37f2f6cf439ff414f8fc46d917929974a82244c20eb10231ba60c54426", size = 51238, upload-time = "2025-08-26T13:09:06.831Z" }
|
||||||
|
wheels = [
|
||||||
|
{ url = "https://files.pythonhosted.org/packages/de/15/545e2b6cf2e3be84bc1ed85613edd75b8aea69807a71c26f4ca6a9258e82/email_validator-2.3.0-py3-none-any.whl", hash = "sha256:80f13f623413e6b197ae73bb10bf4eb0908faf509ad8362c5edeb0be7fd450b4", size = 35604, upload-time = "2025-08-26T13:09:05.858Z" },
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "genson"
|
name = "genson"
|
||||||
version = "1.3.0"
|
version = "1.3.0"
|
||||||
@ -158,6 +185,21 @@ wheels = [
|
|||||||
{ url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" },
|
{ url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" },
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "httpx-ws"
|
||||||
|
version = "0.9.0"
|
||||||
|
source = { registry = "https://pypi.org/simple" }
|
||||||
|
dependencies = [
|
||||||
|
{ name = "anyio" },
|
||||||
|
{ name = "httpcore" },
|
||||||
|
{ name = "httpx" },
|
||||||
|
{ name = "wsproto" },
|
||||||
|
]
|
||||||
|
sdist = { url = "https://files.pythonhosted.org/packages/cd/cd/ca91a07ae446451f7476bf3fcc909e98cb942ff032ebfda0e3fe449aca7b/httpx_ws-0.9.0.tar.gz", hash = "sha256:797373326f70eec1ae96f6e43ae9f12002fd7d73aee139a4985eaab964338a08", size = 107105, upload-time = "2026-03-28T14:11:10.781Z" }
|
||||||
|
wheels = [
|
||||||
|
{ url = "https://files.pythonhosted.org/packages/98/f8/a6bc80313a9e93c888fa10534dfce2ad76ff86911b6f485777ce6de6a073/httpx_ws-0.9.0-py3-none-any.whl", hash = "sha256:71640d2fb1bf9a225775015b33cd755cfd4c5f7e21c885192fe3adc4c387b248", size = 15759, upload-time = "2026-03-28T14:11:11.887Z" },
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "idna"
|
name = "idna"
|
||||||
version = "3.11"
|
version = "3.11"
|
||||||
@ -504,6 +546,15 @@ wheels = [
|
|||||||
{ url = "https://files.pythonhosted.org/packages/e5/35/f8b19922b6a25bc0880171a2f1a003eaeb93657475193ab516fd87cac9da/pytest_asyncio-1.3.0-py3-none-any.whl", hash = "sha256:611e26147c7f77640e6d0a92a38ed17c3e9848063698d5c93d5aa7aa11cebff5", size = 15075, upload-time = "2025-11-10T16:07:45.537Z" },
|
{ url = "https://files.pythonhosted.org/packages/e5/35/f8b19922b6a25bc0880171a2f1a003eaeb93657475193ab516fd87cac9da/pytest_asyncio-1.3.0-py3-none-any.whl", hash = "sha256:611e26147c7f77640e6d0a92a38ed17c3e9848063698d5c93d5aa7aa11cebff5", size = 15075, upload-time = "2025-11-10T16:07:45.537Z" },
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "python-dotenv"
|
||||||
|
version = "1.2.2"
|
||||||
|
source = { registry = "https://pypi.org/simple" }
|
||||||
|
sdist = { url = "https://files.pythonhosted.org/packages/82/ed/0301aeeac3e5353ef3d94b6ec08bbcabd04a72018415dcb29e588514bba8/python_dotenv-1.2.2.tar.gz", hash = "sha256:2c371a91fbd7ba082c2c1dc1f8bf89ca22564a087c2c287cd9b662adde799cf3", size = 50135, upload-time = "2026-03-01T16:00:26.196Z" }
|
||||||
|
wheels = [
|
||||||
|
{ url = "https://files.pythonhosted.org/packages/0b/d7/1959b9648791274998a9c3526f6d0ec8fd2233e4d4acce81bbae76b44b2a/python_dotenv-1.2.2-py3-none-any.whl", hash = "sha256:1d8214789a24de455a8b8bd8ae6fe3c6b69a5e3d64aa8a8e5d68e694bbcb285a", size = 22101, upload-time = "2026-03-01T16:00:25.09Z" },
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "pytokens"
|
name = "pytokens"
|
||||||
version = "0.4.1"
|
version = "0.4.1"
|
||||||
@ -564,6 +615,18 @@ wheels = [
|
|||||||
{ url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" },
|
{ url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" },
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "respx"
|
||||||
|
version = "0.23.1"
|
||||||
|
source = { registry = "https://pypi.org/simple" }
|
||||||
|
dependencies = [
|
||||||
|
{ name = "httpx" },
|
||||||
|
]
|
||||||
|
sdist = { url = "https://files.pythonhosted.org/packages/43/98/4e55c9c486404ec12373708d015ebce157966965a5ebe7f28ff2c784d41b/respx-0.23.1.tar.gz", hash = "sha256:242dcc6ce6b5b9bf621f5870c82a63997e8e82bc7c947f9ffe272b8f3dd5a780", size = 29243, upload-time = "2026-04-08T14:37:16.008Z" }
|
||||||
|
wheels = [
|
||||||
|
{ url = "https://files.pythonhosted.org/packages/1d/4a/221da6ca167db45693d8d26c7dc79ccfc978a440251bf6721c9aaf251ac0/respx-0.23.1-py2.py3-none-any.whl", hash = "sha256:b18004b029935384bccfa6d7d9d74b4ec9af73a081cc28600fffc0447f4b8c1a", size = 25557, upload-time = "2026-04-08T14:37:14.613Z" },
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "ruff"
|
name = "ruff"
|
||||||
version = "0.15.10"
|
version = "0.15.10"
|
||||||
@ -627,30 +690,50 @@ name = "wrenn"
|
|||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
source = { editable = "." }
|
source = { editable = "." }
|
||||||
dependencies = [
|
dependencies = [
|
||||||
|
{ name = "email-validator" },
|
||||||
{ name = "httpx" },
|
{ name = "httpx" },
|
||||||
|
{ name = "httpx-ws" },
|
||||||
{ name = "pydantic" },
|
{ name = "pydantic" },
|
||||||
|
{ name = "python-dotenv" },
|
||||||
]
|
]
|
||||||
|
|
||||||
[package.dev-dependencies]
|
[package.dev-dependencies]
|
||||||
dev = [
|
dev = [
|
||||||
{ name = "datamodel-code-generator" },
|
{ name = "datamodel-code-generator", extra = ["ruff"] },
|
||||||
{ name = "mypy" },
|
{ name = "mypy" },
|
||||||
{ name = "pytest" },
|
{ name = "pytest" },
|
||||||
{ name = "pytest-asyncio" },
|
{ name = "pytest-asyncio" },
|
||||||
|
{ name = "respx" },
|
||||||
{ name = "ruff" },
|
{ name = "ruff" },
|
||||||
]
|
]
|
||||||
|
|
||||||
[package.metadata]
|
[package.metadata]
|
||||||
requires-dist = [
|
requires-dist = [
|
||||||
|
{ name = "email-validator", specifier = ">=2.3.0" },
|
||||||
{ name = "httpx", specifier = ">=0.28.1" },
|
{ name = "httpx", specifier = ">=0.28.1" },
|
||||||
|
{ name = "httpx-ws", specifier = ">=0.9.0" },
|
||||||
{ name = "pydantic", specifier = ">=2.12.5" },
|
{ name = "pydantic", specifier = ">=2.12.5" },
|
||||||
|
{ name = "python-dotenv", specifier = ">=1.2.2" },
|
||||||
]
|
]
|
||||||
|
|
||||||
[package.metadata.requires-dev]
|
[package.metadata.requires-dev]
|
||||||
dev = [
|
dev = [
|
||||||
{ name = "datamodel-code-generator", specifier = ">=0.56.0" },
|
{ name = "datamodel-code-generator", extras = ["ruff"], specifier = ">=0.56.0" },
|
||||||
{ name = "mypy", specifier = ">=1.20.0" },
|
{ name = "mypy", specifier = ">=1.20.0" },
|
||||||
{ name = "pytest", specifier = ">=9.0.3" },
|
{ name = "pytest", specifier = ">=9.0.3" },
|
||||||
{ name = "pytest-asyncio", specifier = ">=1.3.0" },
|
{ name = "pytest-asyncio", specifier = ">=1.3.0" },
|
||||||
|
{ name = "respx", specifier = ">=0.23.1" },
|
||||||
{ name = "ruff", specifier = ">=0.15.10" },
|
{ name = "ruff", specifier = ">=0.15.10" },
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "wsproto"
|
||||||
|
version = "1.3.2"
|
||||||
|
source = { registry = "https://pypi.org/simple" }
|
||||||
|
dependencies = [
|
||||||
|
{ name = "h11" },
|
||||||
|
]
|
||||||
|
sdist = { url = "https://files.pythonhosted.org/packages/c7/79/12135bdf8b9c9367b8701c2c19a14c913c120b882d50b014ca0d38083c2c/wsproto-1.3.2.tar.gz", hash = "sha256:b86885dcf294e15204919950f666e06ffc6c7c114ca900b060d6e16293528294", size = 50116, upload-time = "2025-11-20T18:18:01.871Z" }
|
||||||
|
wheels = [
|
||||||
|
{ url = "https://files.pythonhosted.org/packages/a4/f5/10b68b7b1544245097b2a1b8238f66f2fc6dcaeb24ba5d917f52bd2eed4f/wsproto-1.3.2-py3-none-any.whl", hash = "sha256:61eea322cdf56e8cc904bd3ad7573359a242ba65688716b0710a5eb12beab584", size = 24405, upload-time = "2025-11-20T18:18:00.454Z" },
|
||||||
|
]
|
||||||
|
|||||||
Reference in New Issue
Block a user