From e5e4e1a85b7cf5aacb826a11bd6964e402bdefca Mon Sep 17 00:00:00 2001 From: pptx704 Date: Sat, 16 May 2026 17:57:20 +0600 Subject: [PATCH] fix: update SDK for v0.2.0 API compatibility Sync OpenAPI spec to v0.2.0, fix type annotation shadowing by using builtins.list in annotated signatures, guard poll interval lookup against None status, and reorder capsule ID assignment to validate before storing. --- .gitignore | 1 + AGENTS.md | 56 ---------------- api/openapi.yaml | 115 +++++++++++++++++++++++++++++---- docs/reference.md | 23 +++---- src/wrenn/async_capsule.py | 29 +++++++-- src/wrenn/capsule.py | 33 +++++++--- src/wrenn/client.py | 4 +- src/wrenn/exceptions.py | 3 + src/wrenn/models/_generated.py | 5 +- tests/test_capsule_features.py | 49 +++++++++----- tests/test_client.py | 20 +++--- 11 files changed, 212 insertions(+), 126 deletions(-) delete mode 100644 AGENTS.md diff --git a/.gitignore b/.gitignore index 619209d..3632361 100644 --- a/.gitignore +++ b/.gitignore @@ -181,3 +181,4 @@ CODE_EXECUTION.md .code-review-graph/ .claude .mcp.json +AGENTS.md diff --git a/AGENTS.md b/AGENTS.md deleted file mode 100644 index 53b599d..0000000 --- a/AGENTS.md +++ /dev/null @@ -1,56 +0,0 @@ -# AGENTS.md - -## Project - -Wrenn Python SDK — a client library for the Wrenn microVM platform. e2b drop-in replacement. -Package name: `wrenn`. Python 3.13+, managed with [uv](https://docs.astral.sh/uv/). - -## Commands - -```bash -uv sync # install deps -make lint # ruff check + format check (no auto-fix) -make test # unit tests only (tests/test_client.py) -make test-integration # all tests including integration (needs live server) -make generate # regenerate models from OpenAPI spec (fetches from remote) -make check # lint + unit test -``` - -- `make test` only runs `tests/test_client.py`, not all unit tests. To run a specific test file: `uv run pytest tests/test_capsule_features.py -v` -- No typecheck step in Makefile or CI. `mypy` is a dev dependency but not wired up — do not assume it runs. - -## Architecture - -- `src/wrenn/` — the library package - - `capsule.py` / `async_capsule.py` — high-level `Capsule` / `AsyncCapsule` (main user-facing classes) - - `client.py` — low-level `WrennClient` / `AsyncWrennClient` - - `commands.py` — command execution and streaming - - `files.py` — filesystem operations - - `pty.py` — interactive terminal (PTY) over WebSocket - - `exceptions.py` — typed error hierarchy (`WrennError` base) - - `models/_generated.py` — **auto-generated** from OpenAPI spec via `datamodel-codegen` (never edit directly; run `make generate`) - - `sandbox.py` — deprecated `Sandbox` alias for `Capsule` - - `code_interpreter/` — specialized capsule for stateful Jupyter kernel execution -- `tests/` — unit tests use `respx` to mock `httpx`; integration tests are in `tests/integration/` -- `api/openapi.yaml` — downloaded OpenAPI spec used for code generation - -## Key Conventions - -- Generated code lives in `src/wrenn/models/_generated.py`. Never edit it. Run `make generate` to update. -- `Sandbox` is a deprecated alias for `Capsule`. New code should use `Capsule` / `AsyncCapsule`. -- Dual sync/async API: every major class has an `Async` counterpart. -- Uses `httpx` for HTTP, `httpx-ws` for WebSockets, `pydantic` for models. -- `__init__.py` uses `__getattr__` for lazy deprecated aliases (`Sandbox`, `WrennHostHasSandboxesError`). - -## Testing - -- Unit tests mock HTTP via `respx` (httpx mocking library). -- Integration tests require env vars: `WRENN_API_KEY` (or `WRENN_TOKEN`), optionally `WRENN_BASE_URL`. -- Integration test fixtures in `tests/integration/conftest.py` create real capsules and clean them up. -- `pytest` marker: `@pytest.mark.integration` for tests needing a live server. - -## CI - -Woodpecker CI (`.woodpecker/check.yml`) runs on push to `main` and `dev`: -1. `make lint` -2. `make test` (unit tests only — integration tests are not in CI) diff --git a/api/openapi.yaml b/api/openapi.yaml index 8d3861c..dfc5c75 100644 --- a/api/openapi.yaml +++ b/api/openapi.yaml @@ -1,8 +1,8 @@ openapi: "3.1.0" info: title: Wrenn API - description: MicroVM-based code execution platform API. - version: "0.1.3" + description: AI agent execution platform API. + version: "0.2.0" servers: - url: http://localhost:8080 @@ -866,8 +866,8 @@ paths: schema: $ref: "#/components/schemas/CreateCapsuleRequest" responses: - "201": - description: Capsule created + "202": + description: Capsule creation initiated (status will be "starting") content: application/json: schema: @@ -988,8 +988,8 @@ paths: security: - apiKeyAuth: [] responses: - "204": - description: Capsule destroyed + "202": + description: Capsule destruction initiated /v1/capsules/{id}/exec: parameters: @@ -1260,8 +1260,8 @@ paths: destroys all running resources. The capsule exists only as files on disk and can be resumed later. responses: - "200": - description: Capsule paused (snapshot taken, resources released) + "202": + description: Capsule pause initiated (status will be "pausing") content: application/json: schema: @@ -1292,8 +1292,8 @@ paths: memory loading. Boots a fresh Firecracker process, sets up a new network slot, and waits for envd to become ready. responses: - "200": - description: Capsule resumed (new VM booted from snapshot) + "202": + description: Capsule resume initiated (status will be "resuming") content: application/json: schema: @@ -2035,6 +2035,51 @@ paths: schema: $ref: "#/components/schemas/Error" + /v1/hosts/sandbox-events: + post: + summary: Sandbox lifecycle event callback + operationId: sandboxEventCallback + tags: [hosts] + security: + - hostTokenAuth: [] + description: | + Receives autonomous lifecycle events from host agents (e.g. auto-pause + from the TTL reaper). The event is published to an internal Redis stream + for the control plane's event consumer to process. + requestBody: + required: true + content: + application/json: + schema: + type: object + required: [event, sandbox_id, host_id] + properties: + event: + type: string + enum: [sandbox.auto_paused] + sandbox_id: + type: string + host_id: + type: string + timestamp: + type: integer + format: int64 + responses: + "204": + description: Event accepted + "400": + description: Invalid request + content: + application/json: + schema: + $ref: "#/components/schemas/Error" + "403": + description: Host ID mismatch + content: + application/json: + schema: + $ref: "#/components/schemas/Error" + /v1/hosts/auth/refresh: post: summary: Refresh host JWT @@ -2346,6 +2391,54 @@ paths: schema: $ref: "#/components/schemas/Error" + /v1/admin/users/{id}/admin: + put: + summary: Grant or revoke platform admin + operationId: setUserAdmin + tags: [admin] + description: | + Sets the platform admin flag on a user. Cannot remove the last admin. + Requires platform admin access (JWT + is_admin). + The target user's JWT is not re-issued — their frontend will reflect the + change on next login or team switch. + security: + - bearerAuth: [] + parameters: + - name: id + in: path + required: true + schema: + type: string + example: "usr-a1b2c3d4" + requestBody: + required: true + content: + application/json: + schema: + type: object + required: [admin] + properties: + admin: + type: boolean + description: true to grant admin, false to revoke. + responses: + "204": + description: Admin status updated + "400": + $ref: "#/components/responses/BadRequest" + "403": + description: Caller is not a platform admin + content: + application/json: + schema: + $ref: "#/components/schemas/Error" + "404": + description: User not found + content: + application/json: + schema: + $ref: "#/components/schemas/Error" + components: securitySchemes: apiKeyAuth: @@ -2544,7 +2637,7 @@ components: type: string status: type: string - enum: [pending, starting, running, paused, hibernated, stopped, missing, error] + enum: [pending, starting, running, pausing, paused, resuming, stopping, hibernated, stopped, missing, error] template: type: string vcpus: diff --git a/docs/reference.md b/docs/reference.md index 7e32f6c..9a406df 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -1964,15 +1964,17 @@ inactivity TTL is set. #### wait\_ready ```python -async def wait_ready(timeout: float = 30, interval: float = 0.5) -> None +async def wait_ready(timeout: float = 30) -> None ``` Await until the capsule status is ``running``. +Polling interval adapts to the current transient status: +0.5 s for starting/resuming, 2 s for pausing, 1 s for stopping. + **Arguments**: - `timeout` _float_ - Maximum seconds to wait. Defaults to ``30``. -- `interval` _float_ - Polling interval in seconds. Defaults to ``0.5``. **Raises**: @@ -2534,15 +2536,17 @@ inactivity TTL is set. #### wait\_ready ```python -def wait_ready(timeout: float = 30, interval: float = 0.5) -> None +def wait_ready(timeout: float = 30) -> None ``` Block until the capsule status is ``running``. +Polling interval adapts to the current transient status: +0.5 s for starting/resuming, 2 s for pausing, 1 s for stopping. + **Arguments**: - `timeout` _float_ - Maximum seconds to wait. Defaults to ``30``. -- `interval` _float_ - Polling interval in seconds. Defaults to ``0.5``. **Raises**: @@ -2700,17 +2704,6 @@ Create a snapshot template from this capsule's current state. # wrenn.\_config - - -## ConnectionConfig Objects - -```python -@dataclass(frozen=True) -class ConnectionConfig() -``` - -Resolved credentials and base URL for Wrenn API calls. - # wrenn.\_git.\_auth diff --git a/src/wrenn/async_capsule.py b/src/wrenn/async_capsule.py index 1d72408..44a9e84 100644 --- a/src/wrenn/async_capsule.py +++ b/src/wrenn/async_capsule.py @@ -1,8 +1,8 @@ from __future__ import annotations import asyncio -import logging import builtins +import logging import time from collections.abc import AsyncIterator from contextlib import asynccontextmanager @@ -140,14 +140,19 @@ class AsyncCapsule: info = await client.capsules.get(capsule_id) if info.status == Status.paused: - info = await client.capsules.resume(capsule_id) + await client.capsules.resume(capsule_id) - return cls( + capsule = cls( _capsule_id=capsule_id, _client=client, _info=info, ) + if info.status != Status.running: + await capsule.wait_ready() + + return capsule + # ── Dual instance/static lifecycle ────────────────────────── destroy = _DualMethod("_instance_destroy", "_static_destroy") @@ -224,12 +229,21 @@ class AsyncCapsule: """ await self._client.capsules.ping(self._id) - async def wait_ready(self, timeout: float = 30, interval: float = 0.5) -> None: + _POLL_INTERVALS: dict[Status, float] = { + Status.starting: 0.5, + Status.resuming: 0.5, + Status.pausing: 2.0, + Status.stopping: 1.0, + } + + async def wait_ready(self, timeout: float = 30) -> None: """Await until the capsule status is ``running``. + Polling interval adapts to the current transient status: + 0.5 s for starting/resuming, 2 s for pausing, 1 s for stopping. + Args: timeout (float): Maximum seconds to wait. Defaults to ``30``. - interval (float): Polling interval in seconds. Defaults to ``0.5``. Raises: TimeoutError: If the capsule does not reach ``running`` state @@ -246,7 +260,10 @@ class AsyncCapsule: if info.status in (Status.error, Status.stopped): raise RuntimeError(f"Capsule entered {info.status} state while waiting") if info.status == Status.paused: - info = await self._client.capsules.resume(self._id) + await self._client.capsules.resume(self._id) + interval = ( + self._POLL_INTERVALS.get(info.status, 0.5) if info.status else 0.5 + ) await asyncio.sleep(interval) raise TimeoutError(f"Capsule {self._id} did not become ready within {timeout}s") diff --git a/src/wrenn/capsule.py b/src/wrenn/capsule.py index 29fe52f..24dd4a5 100644 --- a/src/wrenn/capsule.py +++ b/src/wrenn/capsule.py @@ -1,7 +1,7 @@ from __future__ import annotations -import logging import builtins +import logging import time from collections.abc import Iterator from contextlib import contextmanager @@ -112,9 +112,9 @@ class Capsule: memory_mb=memory_mb, timeout_sec=timeout, ) - self._id = self._info.id - if self._id is None: + if self._info.id is None: raise RuntimeError("API returned a capsule without an ID") + self._id = self._info.id except Exception: self._client.close() raise @@ -214,14 +214,19 @@ class Capsule: info = client.capsules.get(capsule_id) if info.status == Status.paused: - info = client.capsules.resume(capsule_id) + client.capsules.resume(capsule_id) - return cls( + capsule = cls( _capsule_id=capsule_id, _client=client, _info=info, ) + if info.status != Status.running: + capsule.wait_ready() + + return capsule + # ── Dual instance/static lifecycle ────────────────────────── destroy = _DualMethod("_instance_destroy", "_static_destroy") @@ -306,12 +311,21 @@ class Capsule: """ self._client.capsules.ping(self._id) - def wait_ready(self, timeout: float = 30, interval: float = 0.5) -> None: + _POLL_INTERVALS: dict[Status, float] = { + Status.starting: 0.5, + Status.resuming: 0.5, + Status.pausing: 2.0, + Status.stopping: 1.0, + } + + def wait_ready(self, timeout: float = 30) -> None: """Block until the capsule status is ``running``. + Polling interval adapts to the current transient status: + 0.5 s for starting/resuming, 2 s for pausing, 1 s for stopping. + Args: timeout (float): Maximum seconds to wait. Defaults to ``30``. - interval (float): Polling interval in seconds. Defaults to ``0.5``. Raises: TimeoutError: If the capsule does not reach ``running`` state @@ -328,7 +342,10 @@ class Capsule: if info.status in (Status.error, Status.stopped): raise RuntimeError(f"Capsule entered {info.status} state while waiting") if info.status == Status.paused: - info = self._client.capsules.resume(self._id) + self._client.capsules.resume(self._id) + interval = ( + self._POLL_INTERVALS.get(info.status, 0.5) if info.status else 0.5 + ) time.sleep(interval) raise TimeoutError(f"Capsule {self._id} did not become ready within {timeout}s") diff --git a/src/wrenn/client.py b/src/wrenn/client.py index c51b190..ceece27 100644 --- a/src/wrenn/client.py +++ b/src/wrenn/client.py @@ -111,7 +111,7 @@ class CapsulesResource: Raises: WrennNotFoundError: If no capsule with the given ID exists. """ - resp = self._http.post(f"/v1/capsules/{id}/pause", timeout=_LONG_TIMEOUT) + resp = self._http.post(f"/v1/capsules/{id}/pause") return CapsuleModel.model_validate(handle_response(resp)) def resume(self, id: str) -> CapsuleModel: @@ -227,7 +227,7 @@ class AsyncCapsulesResource: Raises: WrennNotFoundError: If no capsule with the given ID exists. """ - resp = await self._http.post(f"/v1/capsules/{id}/pause", timeout=_LONG_TIMEOUT) + resp = await self._http.post(f"/v1/capsules/{id}/pause") return CapsuleModel.model_validate(handle_response(resp)) async def resume(self, id: str) -> CapsuleModel: diff --git a/src/wrenn/exceptions.py b/src/wrenn/exceptions.py index af16f6c..65ac7e8 100644 --- a/src/wrenn/exceptions.py +++ b/src/wrenn/exceptions.py @@ -150,6 +150,9 @@ def handle_response(resp: httpx.Response) -> dict | list: if resp.status_code == 204: return {} + if not resp.content: + return {} + return resp.json() diff --git a/src/wrenn/models/_generated.py b/src/wrenn/models/_generated.py index 5542c2f..8bcec35 100644 --- a/src/wrenn/models/_generated.py +++ b/src/wrenn/models/_generated.py @@ -1,6 +1,6 @@ # generated by datamodel-codegen: # filename: openapi.yaml -# timestamp: 2026-04-22T20:21:34+00:00 +# timestamp: 2026-05-15T07:57:28+00:00 from __future__ import annotations from pydantic import AwareDatetime, BaseModel, EmailStr, Field @@ -133,7 +133,10 @@ class Status(StrEnum): pending = "pending" starting = "starting" running = "running" + pausing = "pausing" paused = "paused" + resuming = "resuming" + stopping = "stopping" hibernated = "hibernated" stopped = "stopped" missing = "missing" diff --git a/tests/test_capsule_features.py b/tests/test_capsule_features.py index 825eb52..229a907 100644 --- a/tests/test_capsule_features.py +++ b/tests/test_capsule_features.py @@ -1,5 +1,6 @@ from __future__ import annotations +import httpx import respx from wrenn.capsule import Capsule, _build_proxy_url @@ -30,9 +31,13 @@ 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"} + 202, json={"id": "cl-1", "status": "starting", "template": "minimal"} + ) + cap = Capsule( + template="minimal", + api_key="wrn_test1234567890abcdef12345678", + base_url=BASE, ) - cap = Capsule(template="minimal", api_key="wrn_test1234567890abcdef12345678", base_url=BASE) assert cap.capsule_id == "cl-1" assert hasattr(cap, "commands") assert hasattr(cap, "files") @@ -40,7 +45,7 @@ class TestCapsuleCreate: @respx.mock def test_capsule_create_classmethod(self): respx.post(f"{BASE}/v1/capsules").respond( - 201, json={"id": "cl-2", "status": "pending"} + 202, json={"id": "cl-2", "status": "starting"} ) cap = Capsule.create(api_key="wrn_test1234567890abcdef12345678", base_url=BASE) assert cap.capsule_id == "cl-2" @@ -48,9 +53,9 @@ class TestCapsuleCreate: @respx.mock def test_capsule_context_manager_kills(self): respx.post(f"{BASE}/v1/capsules").respond( - 201, json={"id": "cl-1", "status": "pending"} + 202, json={"id": "cl-1", "status": "starting"} ) - kill_route = respx.delete(f"{BASE}/v1/capsules/cl-1").respond(204) + kill_route = respx.delete(f"{BASE}/v1/capsules/cl-1").respond(202) with Capsule(api_key="wrn_test1234567890abcdef12345678", base_url=BASE) as cap: assert cap.capsule_id == "cl-1" assert kill_route.called @@ -59,7 +64,7 @@ class TestCapsuleCreate: 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"} + 202, json={"id": "cl-3", "status": "starting"} ) cap = Capsule(base_url=BASE) assert cap.capsule_id == "cl-3" @@ -68,17 +73,21 @@ class TestCapsuleCreate: 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", base_url=BASE) + route = respx.delete(f"{BASE}/v1/capsules/cl-1").respond(202) + Capsule._static_destroy( + "cl-1", api_key="wrn_test1234567890abcdef12345678", base_url=BASE + ) 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"} + 202, json={"id": "cl-1", "status": "pausing"} ) - info = Capsule._static_pause("cl-1", api_key="wrn_test1234567890abcdef12345678", base_url=BASE) - assert info.status.value == "paused" + info = Capsule._static_pause( + "cl-1", api_key="wrn_test1234567890abcdef12345678", base_url=BASE + ) + assert info.status.value == "pausing" @respx.mock def test_static_list(self): @@ -106,18 +115,24 @@ class TestCapsuleConnect: 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", base_url=BASE) + cap = Capsule.connect( + "cl-1", api_key="wrn_test1234567890abcdef12345678", base_url=BASE + ) 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"} - ) + get_route = respx.get(f"{BASE}/v1/capsules/cl-1") + get_route.side_effect = [ + httpx.Response(200, json={"id": "cl-1", "status": "paused"}), + httpx.Response(200, json={"id": "cl-1", "status": "running"}), + ] respx.post(f"{BASE}/v1/capsules/cl-1/resume").respond( - 200, json={"id": "cl-1", "status": "running"} + 202, json={"id": "cl-1", "status": "resuming"} + ) + cap = Capsule.connect( + "cl-1", api_key="wrn_test1234567890abcdef12345678", base_url=BASE ) - cap = Capsule.connect("cl-1", api_key="wrn_test1234567890abcdef12345678", base_url=BASE) assert cap.capsule_id == "cl-1" diff --git a/tests/test_client.py b/tests/test_client.py index 36adce9..1269233 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -36,10 +36,10 @@ class TestCapsules: @respx.mock def test_create(self, client): respx.post(f"{BASE}/v1/capsules").respond( - 201, + 202, json={ "id": "sb-1", - "status": "pending", + "status": "starting", "template": "base-python", "vcpus": 2, "memory_mb": 1024, @@ -48,12 +48,12 @@ class TestCapsules: 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 + assert resp.status == Status.starting @respx.mock def test_create_defaults(self, client): respx.post(f"{BASE}/v1/capsules").respond( - 201, json={"id": "sb-2", "status": "pending"} + 202, json={"id": "sb-2", "status": "starting"} ) resp = client.capsules.create() assert resp.id == "sb-2" @@ -77,25 +77,25 @@ class TestCapsules: @respx.mock def test_destroy(self, client): - route = respx.delete(f"{BASE}/v1/capsules/sb-1").respond(204) + route = respx.delete(f"{BASE}/v1/capsules/sb-1").respond(202) 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"} + 202, json={"id": "sb-1", "status": "pausing"} ) resp = client.capsules.pause("sb-1") - assert resp.status == Status.paused + assert resp.status == Status.pausing @respx.mock def test_resume(self, client): respx.post(f"{BASE}/v1/capsules/sb-1/resume").respond( - 200, json={"id": "sb-1", "status": "running"} + 202, json={"id": "sb-1", "status": "resuming"} ) resp = client.capsules.resume("sb-1") - assert resp.status == Status.running + assert resp.status == Status.resuming @respx.mock def test_ping(self, client): @@ -238,7 +238,7 @@ class TestAsyncClient: 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"} + 202, json={"id": "sb-1", "status": "starting"} ) resp = await async_client.capsules.create(template="base-python") assert resp.id == "sb-1"