From bf5914c0a8e7ba622b1fce060f7602f14dad5270 Mon Sep 17 00:00:00 2001 From: Tasnim Kabir Sadik Date: Mon, 13 Apr 2026 03:16:27 +0600 Subject: [PATCH] fix: renamed sandbox to capsule --- .woodpecker/check.yml | 50 +- api/openapi.yaml | 262 ++-- src/wrenn/__init__.py | 58 +- src/wrenn/capsule.py | 1171 ++++++++++++++++ src/wrenn/client.py | 95 +- src/wrenn/exceptions.py | 39 +- src/wrenn/models/__init__.py | 8 +- src/wrenn/models/_generated.py | 196 +-- src/wrenn/pty.py | 8 +- src/wrenn/sandbox.py | 1197 +---------------- ...x_features.py => test_capsule_features.py} | 119 +- tests/test_client.py | 88 +- tests/test_filesystem_pty.py | 75 +- tests/test_integration.py | 368 ++--- 14 files changed, 1929 insertions(+), 1805 deletions(-) create mode 100644 src/wrenn/capsule.py rename tests/{test_sandbox_features.py => test_capsule_features.py} (53%) diff --git a/.woodpecker/check.yml b/.woodpecker/check.yml index 7b4b167..83a35d7 100644 --- a/.woodpecker/check.yml +++ b/.woodpecker/check.yml @@ -1,42 +1,46 @@ -kind: pipeline -name: static-analysis - when: - - event: push - branch: - - main - - dev + event: push + branch: + - main + - dev variables: - - &python_image ghcr.io/astral-sh/uv:python3.13-bookworm-slim - - &uv_cache_dir /root/.cache/uv + - &python_image "ghcr.io/astral-sh/uv:python3.13-bookworm-slim" + - &uv_cache_dir "/root/.cache/uv" steps: - lint: + - 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" + UV_FROZEN: 1 commands: - uv sync --no-install-project - make lint - volumes: - - name: uv-cache - path: *uv_cache_dir - test: + - name: test image: *python_image environment: UV_CACHE_DIR: *uv_cache_dir - UV_FROZEN: "1" + UV_FROZEN: 1 commands: - uv sync --no-install-project - make test - volumes: - - name: uv-cache - path: *uv_cache_dir -volumes: - - name: uv-cache - host: - path: /var/lib/woodpecker/cache/uv + - name: rebuild-cache + image: woodpeckerci/plugin-cache + when: + - status: [success] + settings: + rebuild: true + cache_key: "uv-{{ checksum \"uv.lock\" }}" + mount: + - /root/.cache/uv diff --git a/api/openapi.yaml b/api/openapi.yaml index 0b56fe5..b6bd643 100644 --- a/api/openapi.yaml +++ b/api/openapi.yaml @@ -1,6 +1,6 @@ openapi: "3.1.0" info: - title: Wrenn Sandbox API + title: Wrenn API description: MicroVM-based code execution platform API. version: "0.1.0" @@ -393,7 +393,7 @@ paths: - bearerAuth: [] description: | Owner only. Soft-deletes the team and destroys all running/paused/starting - sandboxes. All DB records are preserved. The team slug is permanently reserved. + capsulees. All DB records are preserved. The team slug is permanently reserved. responses: "204": description: Team deleted @@ -570,11 +570,11 @@ paths: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes: + /v1/capsules: post: - summary: Create a sandbox - operationId: createSandbox - tags: [sandboxes] + summary: Create a capsule + operationId: createCapsule + tags: [capsules] security: - apiKeyAuth: [] requestBody: @@ -582,14 +582,14 @@ paths: content: application/json: schema: - $ref: "#/components/schemas/CreateSandboxRequest" + $ref: "#/components/schemas/CreateCapsuleRequest" responses: "201": - description: Sandbox created + description: Capsule created content: application/json: schema: - $ref: "#/components/schemas/Sandbox" + $ref: "#/components/schemas/Capsule" "502": description: Host agent error content: @@ -598,26 +598,26 @@ paths: $ref: "#/components/schemas/Error" get: - summary: List sandboxes for your team - operationId: listSandboxes - tags: [sandboxes] + summary: List capsulees for your team + operationId: listCapsules + tags: [capsules] security: - apiKeyAuth: [] responses: "200": - description: List of sandboxes + description: List of capsulees content: application/json: schema: type: array items: - $ref: "#/components/schemas/Sandbox" + $ref: "#/components/schemas/Capsule" - /v1/sandboxes/stats: + /v1/capsules/stats: get: - summary: Get sandbox usage stats for your team - operationId: getSandboxStats - tags: [sandboxes] + summary: Get capsule usage stats for your team + operationId: getCapsuleStats + tags: [capsules] security: - apiKeyAuth: [] parameters: @@ -631,15 +631,15 @@ paths: description: Time window for the time-series data. responses: "200": - description: Sandbox stats for the team + description: Capsule stats for the team content: application/json: schema: - $ref: "#/components/schemas/SandboxStats" + $ref: "#/components/schemas/CapsuleStats" "400": $ref: "#/components/responses/BadRequest" - /v1/sandboxes/{id}: + /v1/capsules/{id}: parameters: - name: id in: path @@ -648,36 +648,36 @@ paths: type: string get: - summary: Get sandbox details - operationId: getSandbox - tags: [sandboxes] + summary: Get capsule details + operationId: getCapsule + tags: [capsules] security: - apiKeyAuth: [] responses: "200": - description: Sandbox details + description: Capsule details content: application/json: schema: - $ref: "#/components/schemas/Sandbox" + $ref: "#/components/schemas/Capsule" "404": - description: Sandbox not found + description: Capsule not found content: application/json: schema: $ref: "#/components/schemas/Error" delete: - summary: Destroy a sandbox - operationId: destroySandbox - tags: [sandboxes] + summary: Destroy a capsule + operationId: destroyCapsule + tags: [capsules] security: - apiKeyAuth: [] responses: "204": - description: Sandbox destroyed + description: Capsule destroyed - /v1/sandboxes/{id}/exec: + /v1/capsules/{id}/exec: parameters: - name: id in: path @@ -688,7 +688,7 @@ paths: post: summary: Execute a command operationId: execCommand - tags: [sandboxes] + tags: [capsules] security: - apiKeyAuth: [] requestBody: @@ -705,19 +705,19 @@ paths: schema: $ref: "#/components/schemas/ExecResponse" "404": - description: Sandbox not found + description: Capsule not found content: application/json: schema: $ref: "#/components/schemas/Error" "409": - description: Sandbox not running + description: Capsule not running content: application/json: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes/{id}/ping: + /v1/capsules/{id}/ping: parameters: - name: id in: path @@ -726,32 +726,32 @@ paths: type: string post: - summary: Reset sandbox inactivity timer - operationId: pingSandbox - tags: [sandboxes] + summary: Reset capsule inactivity timer + operationId: pingCapsule + tags: [capsules] security: - apiKeyAuth: [] description: | - Resets the last_active_at timestamp for a running sandbox, preventing - the auto-pause TTL from expiring. Use this as a keepalive for sandboxes + Resets the last_active_at timestamp for a running capsule, preventing + the auto-pause TTL from expiring. Use this as a keepalive for capsulees that are idle but should remain running. responses: "204": description: Ping acknowledged, inactivity timer reset "404": - description: Sandbox not found + description: Capsule not found content: application/json: schema: $ref: "#/components/schemas/Error" "409": - description: Sandbox not running + description: Capsule not running content: application/json: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes/{id}/metrics: + /v1/capsules/{id}/metrics: parameters: - name: id in: path @@ -760,22 +760,22 @@ paths: type: string get: - summary: Get per-sandbox resource metrics - operationId: getSandboxMetrics - tags: [sandboxes] + summary: Get per-capsule resource metrics + operationId: getCapsuleMetrics + tags: [capsules] security: - apiKeyAuth: [] - bearerAuth: [] description: | - Returns time-series CPU, memory, and disk metrics for a sandbox. + Returns time-series CPU, memory, and disk metrics for a capsule. Three tiers are available with different granularity and retention: - `10m`: 500ms samples, last 10 minutes - `2h`: 30-second averages, last 2 hours - `24h`: 5-minute averages, last 24 hours - For running sandboxes, data comes from the host agent's in-memory - ring buffer. For paused sandboxes, data is read from persisted - snapshots in the database. Stopped/destroyed sandboxes return 404. + For running capsulees, data comes from the host agent's in-memory + ring buffer. For paused capsulees, data is read from persisted + snapshots in the database. Stopped/destroyed capsulees return 404. parameters: - name: range in: query @@ -791,7 +791,7 @@ paths: content: application/json: schema: - $ref: "#/components/schemas/SandboxMetrics" + $ref: "#/components/schemas/CapsuleMetrics" "400": description: Invalid range parameter content: @@ -799,13 +799,13 @@ paths: schema: $ref: "#/components/schemas/Error" "404": - description: Sandbox not found or metrics not available + description: Capsule not found or metrics not available content: application/json: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes/{id}/pause: + /v1/capsules/{id}/pause: parameters: - name: id in: path @@ -814,30 +814,30 @@ paths: type: string post: - summary: Pause a running sandbox - operationId: pauseSandbox - tags: [sandboxes] + summary: Pause a running capsule + operationId: pauseCapsule + tags: [capsules] security: - apiKeyAuth: [] description: | - Takes a snapshot of the sandbox (VM state + memory + rootfs), then - destroys all running resources. The sandbox exists only as files on + Takes a snapshot of the capsule (VM state + memory + rootfs), then + destroys all running resources. The capsule exists only as files on disk and can be resumed later. responses: "200": - description: Sandbox paused (snapshot taken, resources released) + description: Capsule paused (snapshot taken, resources released) content: application/json: schema: - $ref: "#/components/schemas/Sandbox" + $ref: "#/components/schemas/Capsule" "409": - description: Sandbox not running + description: Capsule not running content: application/json: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes/{id}/resume: + /v1/capsules/{id}/resume: parameters: - name: id in: path @@ -846,24 +846,24 @@ paths: type: string post: - summary: Resume a paused sandbox - operationId: resumeSandbox - tags: [sandboxes] + summary: Resume a paused capsule + operationId: resumeCapsule + tags: [capsules] security: - apiKeyAuth: [] description: | - Restores a paused sandbox from its snapshot using UFFD for lazy + Restores a paused capsule from its snapshot using UFFD for lazy memory loading. Boots a fresh Firecracker process, sets up a new network slot, and waits for envd to become ready. responses: "200": - description: Sandbox resumed (new VM booted from snapshot) + description: Capsule resumed (new VM booted from snapshot) content: application/json: schema: - $ref: "#/components/schemas/Sandbox" + $ref: "#/components/schemas/Capsule" "409": - description: Sandbox not paused + description: Capsule not paused content: application/json: schema: @@ -877,9 +877,9 @@ paths: security: - apiKeyAuth: [] description: | - Pauses a running sandbox, takes a full snapshot, copies the snapshot + Pauses a running capsule, takes a full snapshot, copies the snapshot files to the images directory as a reusable template, then destroys - the sandbox. The template can be used to create new sandboxes. + the capsule. The template can be used to create new capsulees. parameters: - name: overwrite in: query @@ -902,7 +902,7 @@ paths: schema: $ref: "#/components/schemas/Template" "409": - description: Name already exists or sandbox not running + description: Name already exists or capsule not running content: application/json: schema: @@ -957,7 +957,7 @@ paths: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes/{id}/files/write: + /v1/capsules/{id}/files/write: parameters: - name: id in: path @@ -968,7 +968,7 @@ paths: post: summary: Upload a file operationId: uploadFile - tags: [sandboxes] + tags: [capsules] security: - apiKeyAuth: [] requestBody: @@ -981,7 +981,7 @@ paths: properties: path: type: string - description: Absolute destination path inside the sandbox + description: Absolute destination path inside the capsule file: type: string format: binary @@ -990,7 +990,7 @@ paths: "204": description: File uploaded "409": - description: Sandbox not running + description: Capsule not running content: application/json: schema: @@ -1002,7 +1002,7 @@ paths: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes/{id}/files/read: + /v1/capsules/{id}/files/read: parameters: - name: id in: path @@ -1013,7 +1013,7 @@ paths: post: summary: Download a file operationId: downloadFile - tags: [sandboxes] + tags: [capsules] security: - apiKeyAuth: [] requestBody: @@ -1031,13 +1031,13 @@ paths: type: string format: binary "404": - description: Sandbox or file not found + description: Capsule or file not found content: application/json: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes/{id}/files/list: + /v1/capsules/{id}/files/list: parameters: - name: id in: path @@ -1048,7 +1048,7 @@ paths: post: summary: List directory contents operationId: listDir - tags: [sandboxes] + tags: [capsules] security: - apiKeyAuth: [] requestBody: @@ -1065,19 +1065,19 @@ paths: schema: $ref: "#/components/schemas/ListDirResponse" "404": - description: Sandbox not found + description: Capsule not found content: application/json: schema: $ref: "#/components/schemas/Error" "409": - description: Sandbox not running + description: Capsule not running content: application/json: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes/{id}/files/mkdir: + /v1/capsules/{id}/files/mkdir: parameters: - name: id in: path @@ -1088,7 +1088,7 @@ paths: post: summary: Create a directory operationId: makeDir - tags: [sandboxes] + tags: [capsules] security: - apiKeyAuth: [] requestBody: @@ -1105,19 +1105,19 @@ paths: schema: $ref: "#/components/schemas/MakeDirResponse" "404": - description: Sandbox not found + description: Capsule not found content: application/json: schema: $ref: "#/components/schemas/Error" "409": - description: Sandbox not running + description: Capsule not running content: application/json: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes/{id}/files/remove: + /v1/capsules/{id}/files/remove: parameters: - name: id in: path @@ -1128,7 +1128,7 @@ paths: post: summary: Remove a file or directory operationId: removePath - tags: [sandboxes] + tags: [capsules] security: - apiKeyAuth: [] requestBody: @@ -1141,19 +1141,19 @@ paths: "204": description: File or directory removed "404": - description: Sandbox not found + description: Capsule not found content: application/json: schema: $ref: "#/components/schemas/Error" "409": - description: Sandbox not running + description: Capsule not running content: application/json: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes/{id}/exec/stream: + /v1/capsules/{id}/exec/stream: parameters: - name: id in: path @@ -1164,7 +1164,7 @@ paths: get: summary: Stream command execution via WebSocket operationId: execStream - tags: [sandboxes] + tags: [capsules] security: - apiKeyAuth: [] description: | @@ -1194,19 +1194,19 @@ paths: "101": description: WebSocket upgrade "404": - description: Sandbox not found + description: Capsule not found content: application/json: schema: $ref: "#/components/schemas/Error" "409": - description: Sandbox not running + description: Capsule not running content: application/json: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes/{id}/pty: + /v1/capsules/{id}/pty: parameters: - name: id in: path @@ -1217,7 +1217,7 @@ paths: get: summary: Interactive PTY session via WebSocket operationId: ptySession - tags: [sandboxes] + tags: [capsules] security: - apiKeyAuth: [] description: | @@ -1266,25 +1266,25 @@ paths: Sessions have a 120-second inactivity timeout (reset on input/resize). Sessions persist across WebSocket disconnections — the process keeps - running in the sandbox. Use the `tag` from the "started" response to + running in the capsule. Use the `tag` from the "started" response to reconnect later. responses: "101": description: WebSocket upgrade "404": - description: Sandbox not found + description: Capsule not found content: application/json: schema: $ref: "#/components/schemas/Error" "409": - description: Sandbox not running + description: Capsule not running content: application/json: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes/{id}/files/stream/write: + /v1/capsules/{id}/files/stream/write: parameters: - name: id in: path @@ -1295,11 +1295,11 @@ paths: post: summary: Upload a file (streaming) operationId: streamUploadFile - tags: [sandboxes] + tags: [capsules] security: - apiKeyAuth: [] description: | - Streams file content to the sandbox without buffering in memory. + Streams file content to the capsule without buffering in memory. Suitable for large files. Uses the same multipart/form-data format as the non-streaming upload endpoint. requestBody: @@ -1312,7 +1312,7 @@ paths: properties: path: type: string - description: Absolute destination path inside the sandbox + description: Absolute destination path inside the capsule file: type: string format: binary @@ -1321,19 +1321,19 @@ paths: "204": description: File uploaded "404": - description: Sandbox not found + description: Capsule not found content: application/json: schema: $ref: "#/components/schemas/Error" "409": - description: Sandbox not running + description: Capsule not running content: application/json: schema: $ref: "#/components/schemas/Error" - /v1/sandboxes/{id}/files/stream/read: + /v1/capsules/{id}/files/stream/read: parameters: - name: id in: path @@ -1344,11 +1344,11 @@ paths: post: summary: Download a file (streaming) operationId: streamDownloadFile - tags: [sandboxes] + tags: [capsules] security: - apiKeyAuth: [] description: | - Streams file content from the sandbox without buffering in memory. + Streams file content from the capsule without buffering in memory. Suitable for large files. Returns raw bytes with chunked transfer encoding. requestBody: required: true @@ -1365,13 +1365,13 @@ paths: type: string format: binary "404": - description: Sandbox or file not found + description: Capsule or file not found content: application/json: schema: $ref: "#/components/schemas/Error" "409": - description: Sandbox not running + description: Capsule not running content: application/json: schema: @@ -1469,14 +1469,14 @@ paths: description: | Admins can delete any host. Team owners and admins can delete BYOC hosts belonging to their team. Without `?force=true`, returns 409 if the host - has active sandboxes. With `?force=true`, destroys all sandboxes first. + has active capsulees. With `?force=true`, destroys all capsulees first. parameters: - name: force in: query required: false schema: type: boolean - description: If true, destroy all sandboxes on the host before deleting. + description: If true, destroy all capsulees on the host before deleting. responses: "204": description: Host deleted @@ -1487,11 +1487,11 @@ paths: schema: $ref: "#/components/schemas/Error" "409": - description: Host has active sandboxes (only when force is not set) + description: Host has active capsulees (only when force is not set) content: application/json: schema: - $ref: "#/components/schemas/HostHasSandboxesError" + $ref: "#/components/schemas/HostHasCapsulesError" /v1/hosts/{id}/token: parameters: @@ -1644,7 +1644,7 @@ paths: security: - bearerAuth: [] description: | - Returns the list of sandbox IDs that would be destroyed if the host + Returns the list of capsule IDs that would be destroyed if the host were deleted with `?force=true`. No state is modified. responses: "200": @@ -1917,7 +1917,7 @@ components: type: apiKey in: header name: X-API-Key - description: API key for sandbox lifecycle operations. Create via POST /v1/api-keys. + description: API key for capsule lifecycle operations. Create via POST /v1/api-keys. bearerAuth: type: http @@ -2002,7 +2002,7 @@ components: description: Full plaintext key. Only returned on creation, never again. nullable: true - CreateSandboxRequest: + CreateCapsuleRequest: type: object properties: template: @@ -2018,11 +2018,11 @@ components: type: integer default: 0 description: > - Auto-pause TTL in seconds. The sandbox is automatically paused + Auto-pause TTL in seconds. The capsule is automatically paused after this duration of inactivity (no exec or ping). 0 means no auto-pause. - SandboxStats: + CapsuleStats: type: object properties: range: @@ -2073,7 +2073,7 @@ components: items: type: integer - Sandbox: + Capsule: type: object properties: id: @@ -2114,7 +2114,7 @@ components: properties: sandbox_id: type: string - description: ID of the running sandbox to snapshot. + description: ID of the running capsule to snapshot. name: type: string description: Name for the snapshot template. Auto-generated if omitted. @@ -2180,7 +2180,7 @@ components: properties: path: type: string - description: Absolute file path inside the sandbox + description: Absolute file path inside the capsule ListDirRequest: type: object @@ -2188,7 +2188,7 @@ components: properties: path: type: string - description: Directory path inside the sandbox + description: Directory path inside the capsule depth: type: integer default: 1 @@ -2238,7 +2238,7 @@ components: properties: path: type: string - description: Directory path to create inside the sandbox + description: Directory path to create inside the capsule MakeDirResponse: type: object @@ -2252,7 +2252,7 @@ components: properties: path: type: string - description: Path to remove inside the sandbox + description: Path to remove inside the capsule CreateHostRequest: type: object @@ -2390,9 +2390,9 @@ components: type: array items: type: string - description: IDs of sandboxes that would be destroyed on force-delete. + description: IDs of capsulees that would be destroyed on force-delete. - HostHasSandboxesError: + HostHasCapsulesError: type: object properties: error: @@ -2407,7 +2407,7 @@ components: type: array items: type: string - description: IDs of active sandboxes blocking deletion. + description: IDs of active capsulees blocking deletion. AddTagRequest: type: object @@ -2471,7 +2471,7 @@ components: items: $ref: "#/components/schemas/TeamMember" - SandboxMetrics: + CapsuleMetrics: type: object properties: sandbox_id: diff --git a/src/wrenn/__init__.py b/src/wrenn/__init__.py index d478216..c25aaf8 100644 --- a/src/wrenn/__init__.py +++ b/src/wrenn/__init__.py @@ -1,22 +1,7 @@ -from wrenn.client import AsyncWrennClient, WrennClient -from wrenn.exceptions import ( - WrennAgentError, - WrennAuthenticationError, - WrennConflictError, - WrennError, - WrennForbiddenError, - WrennHostHasSandboxesError, - WrennHostUnavailableError, - WrennInternalError, - WrennNotFoundError, - WrennValidationError, -) -from wrenn.models import FileEntry -from wrenn.pty import AsyncPtySession, PtyEvent, PtyEventType, PtySession -from wrenn.sandbox import ( +from wrenn.capsule import ( + Capsule, CodeResult, ExecResult, - Sandbox, StreamErrorEvent, StreamEvent, StreamExitEvent, @@ -24,6 +9,21 @@ from wrenn.sandbox import ( StreamStderrEvent, StreamStdoutEvent, ) +from wrenn.client import AsyncWrennClient, WrennClient +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" @@ -31,6 +31,7 @@ __all__ = [ "__version__", "AsyncPtySession", "AsyncWrennClient", + "Capsule", "CodeResult", "ExecResult", "FileEntry", @@ -50,9 +51,32 @@ __all__ = [ "WrennConflictError", "WrennError", "WrennForbiddenError", + "WrennHostHasCapsulesError", "WrennHostHasSandboxesError", "WrennHostUnavailableError", "WrennInternalError", "WrennNotFoundError", "WrennValidationError", ] + + +def __getattr__(name: str) -> type: + if name == "Sandbox": + import warnings + + warnings.warn( + "'Sandbox' is deprecated, use 'Capsule' instead", + DeprecationWarning, + stacklevel=2, + ) + return Capsule + if name == "WrennHostHasSandboxesError": + import warnings + + warnings.warn( + "'WrennHostHasSandboxesError' is deprecated, use 'WrennHostHasCapsulesError' instead", + DeprecationWarning, + stacklevel=2, + ) + return WrennHostHasCapsulesError + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/wrenn/capsule.py b/src/wrenn/capsule.py new file mode 100644 index 0000000..17fec62 --- /dev/null +++ b/src/wrenn/capsule.py @@ -0,0 +1,1171 @@ +from __future__ import annotations + +import asyncio +import base64 +import json +import os +import time +import uuid +import warnings +from collections.abc import AsyncIterator, Iterator +from contextlib import asynccontextmanager, contextmanager +from typing import Any + +import httpx +import httpx_ws + +from wrenn.exceptions import handle_response +from wrenn.models import Capsule as CapsuleModel +from wrenn.models import ( + ExecResponse, + FileEntry, + ListDirResponse, + MakeDirResponse, + Status, +) +from wrenn.pty import AsyncPtySession, PtySession + + +class ExecResult: + """Typed result from a synchronous exec call.""" + + __slots__ = ("stdout", "stderr", "exit_code", "duration_ms", "encoding") + + def __init__( + self, + stdout: str, + stderr: str, + exit_code: int, + duration_ms: int | None, + encoding: str | None, + ) -> None: + self.stdout = stdout + self.stderr = stderr + self.exit_code = exit_code + self.duration_ms = duration_ms + self.encoding = encoding + + +class CodeResult: + """Typed result from stateful code execution (``run_code``). + + 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. + """ + + __slots__ = ("text", "data", "stdout", "stderr", "error") + + def __init__( + self, + text: str | None = None, + data: dict[str, str] | None = None, + stdout: str = "", + stderr: str = "", + error: str | None = None, + ) -> None: + self.text = text + self.data = data + self.stdout = stdout + self.stderr = stderr + self.error = error + + +class StreamEvent: + """Base class for streaming exec events.""" + + __slots__ = ("type",) + + def __init__(self, type: str) -> None: + self.type = type + + +class StreamStartEvent(StreamEvent): + """Process started.""" + + __slots__ = ("pid",) + + def __init__(self, pid: int) -> None: + super().__init__("start") + self.pid = pid + + +class StreamStdoutEvent(StreamEvent): + """Stdout data received.""" + + __slots__ = ("data",) + + def __init__(self, data: str) -> None: + super().__init__("stdout") + self.data = data + + +class StreamStderrEvent(StreamEvent): + """Stderr data received.""" + + __slots__ = ("data",) + + def __init__(self, data: str) -> None: + super().__init__("stderr") + self.data = data + + +class StreamExitEvent(StreamEvent): + """Process exited.""" + + __slots__ = ("exit_code",) + + def __init__(self, exit_code: int) -> None: + super().__init__("exit") + self.exit_code = exit_code + + +class StreamErrorEvent(StreamEvent): + """Error occurred.""" + + __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 _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 Capsule(CapsuleModel): + """Developer-facing capsule interface wrapping the generated Capsule model. + + Provides data-plane methods (exec, file I/O, lifecycle), capsule proxy + helpers, and context-manager support for automatic cleanup. + """ + + _http: httpx.Client | None + _async_http: httpx.AsyncClient | None + _base_url: str + _api_key: str | None + _token: str | None + _proxy_client: httpx.Client | None + _async_proxy_client: httpx.AsyncClient | None + _kernel_id: str | None + _jupyter_ws: Any + _async_jupyter_ws: Any + + def _bind( + self, + http: httpx.Client | httpx.AsyncClient, + base_url: str, + api_key: str | None = None, + token: str | None = None, + ) -> None: + self._base_url = base_url + self._api_key = api_key + self._token = token + self._proxy_client = None + self._async_proxy_client = None + self._kernel_id = None + self._jupyter_ws = None + self._async_jupyter_ws = None + if isinstance(http, httpx.Client): + self._http = http + self._async_http = None + else: + self._http = None # type: ignore[assignment] + self._async_http = http + + def _proxy_headers(self) -> dict[str, str]: + headers: dict[str, str] = {} + if self._api_key: + headers["X-API-Key"] = self._api_key + if self._token: + headers["Authorization"] = f"Bearer {self._token}" + return headers + + def _clear_content_type(self) -> dict[str, str]: + assert self._http is not None + headers = dict(self._http.headers) + headers.pop("Content-Type", None) + return headers + + def _async_clear_content_type(self) -> dict[str, str]: + assert self._async_http is not None + headers = dict(self._async_http.headers) + headers.pop("Content-Type", None) + return headers + + def get_url(self, port: int) -> str: + """Construct the proxy URL for a port inside this capsule. + + Args: + port: Port number of the service running inside the capsule. + + Returns: + A URL string like ``http://8888-cl-abc123.api.wrenn.dev``. + """ + return _build_proxy_url(self._base_url, self.id, port) + + @property + def http_client(self) -> httpx.Client: + """A pre-configured ``httpx.Client`` targeting the capsule proxy on port 8888. + + The client has auth headers set and ``base_url`` pointing to + the proxy URL for port 8888. Closed automatically when the capsule exits. + """ + if self._proxy_client is None: + url = ( + _build_proxy_url(self._base_url, self.id, 8888) + .replace("ws://", "http://") + .replace("wss://", "https://") + ) + self._proxy_client = httpx.Client( + base_url=url, + headers=self._proxy_headers(), + ) + return self._proxy_client + + def wait_ready(self, timeout: float = 30, interval: float = 0.5) -> None: + """Block until the capsule status is ``running``. + + Args: + timeout: Maximum seconds to wait. + interval: Seconds between polls. + + Raises: + TimeoutError: If the capsule does not become ready in time. + """ + assert self._http is not None + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + resp = self._http.get(f"/v1/capsules/{self.id}") + data = resp.json() + status = data.get("status") + if status == Status.running: + self.status = Status.running + return + if status in (Status.error, Status.stopped): + raise RuntimeError(f"Capsule entered {status} state while waiting") + time.sleep(interval) + raise TimeoutError(f"Capsule {self.id} did not become ready within {timeout}s") + + async def async_wait_ready( + self, timeout: float = 30, interval: float = 0.5 + ) -> None: + """Async version of ``wait_ready``.""" + assert self._async_http is not None + import asyncio + + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + resp = await self._async_http.get(f"/v1/capsules/{self.id}") + data = resp.json() + status = data.get("status") + if status == Status.running: + self.status = Status.running + return + if status in (Status.error, Status.stopped): + raise RuntimeError(f"Capsule entered {status} state while waiting") + await asyncio.sleep(interval) + raise TimeoutError(f"Capsule {self.id} did not become ready within {timeout}s") + + def exec( + self, + cmd: str, + args: list[str] | None = None, + timeout_sec: int | None = 30, + ) -> ExecResult: + """Execute a command synchronously inside the capsule. + + Args: + cmd: Command to run. + args: Optional positional arguments. + timeout_sec: Execution timeout in seconds. + + Returns: + An ``ExecResult`` with ``stdout``, ``stderr``, ``exit_code``, ``duration_ms``. + """ + assert self._http is not None + payload: dict = {"cmd": cmd} + if args is not None: + payload["args"] = args + if timeout_sec is not None: + payload["timeout_sec"] = timeout_sec + resp = self._http.post(f"/v1/capsules/{self.id}/exec", json=payload) + resp.raise_for_status() + er = ExecResponse.model_validate(resp.json()) + stdout = er.stdout or "" + stderr = er.stderr or "" + if er.encoding == "base64": + stdout = base64.b64decode(stdout).decode("utf-8", errors="replace") + if stderr: + stderr = base64.b64decode(stderr).decode("utf-8", errors="replace") + return ExecResult( + stdout=stdout, + stderr=stderr, + exit_code=er.exit_code if er.exit_code is not None else -1, + duration_ms=er.duration_ms, + encoding=er.encoding, + ) + + async def async_exec( + self, + cmd: str, + args: list[str] | None = None, + timeout_sec: int | None = 30, + ) -> ExecResult: + """Async version of ``exec``.""" + assert self._async_http is not None + payload: dict = {"cmd": cmd} + if args is not None: + payload["args"] = args + if timeout_sec is not None: + payload["timeout_sec"] = timeout_sec + resp = await self._async_http.post(f"/v1/capsules/{self.id}/exec", json=payload) + resp.raise_for_status() + er = ExecResponse.model_validate(resp.json()) + stdout = er.stdout or "" + stderr = er.stderr or "" + if er.encoding == "base64": + stdout = base64.b64decode(stdout).decode("utf-8", errors="replace") + if stderr: + stderr = base64.b64decode(stderr).decode("utf-8", errors="replace") + return ExecResult( + stdout=stdout, + stderr=stderr, + exit_code=er.exit_code if er.exit_code is not None else -1, + duration_ms=er.duration_ms, + encoding=er.encoding, + ) + + def exec_stream( + self, + cmd: str, + args: list[str] | None = None, + ) -> Iterator[StreamEvent]: + """Execute a command via WebSocket, yielding ``StreamEvent`` objects. + + Args: + cmd: Command to run. + args: Optional positional arguments. + + Yields: + ``StreamStartEvent``, ``StreamStdoutEvent``, ``StreamStderrEvent``, + ``StreamExitEvent``, or ``StreamErrorEvent``. + """ + assert self._http is not None + ws: httpx_ws.WebSocketSession + with httpx_ws.connect_ws( # type: ignore[attr-defined] + f"/v1/capsules/{self.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_data: dict = ws.receive_json() # type: ignore[assignment] + event = _parse_stream_event(raw_data) + yield event + + if event.type in ("exit", "error"): + break + + except httpx_ws.WebSocketDisconnect: + break + + async def async_exec_stream( + self, cmd: str, args: list[str] | None = None + ) -> AsyncIterator[StreamEvent]: + """Async version of ``exec_stream``.""" + assert self._async_http is not None + ws: httpx_ws.AsyncWebSocketSession + async with httpx_ws.aconnect_ws( # type: ignore[attr-defined, var-annotated] + f"/v1/capsules/{self.id}/exec/stream", self._async_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_data = await ws.receive_json() + event = _parse_stream_event(raw_data) + yield event + + if event.type in ("exit", "error"): + break + except httpx_ws.WebSocketDisconnect: + pass + + def upload(self, path: str, data: bytes) -> None: + """Upload a small file to the capsule. + + Args: + path: Absolute destination path inside the capsule. + data: File contents as bytes. + """ + assert self._http is not None + resp = self._http.post( + f"/v1/capsules/{self.id}/files/write", + files={"file": ("upload", data)}, + data={"path": path}, + ) + + resp.raise_for_status() + + async def async_upload(self, path: str, data: bytes) -> None: + """Async version of ``upload``.""" + assert self._async_http is not None + resp = await self._async_http.post( + f"/v1/capsules/{self.id}/files/write", + files={"file": ("upload", data)}, + data={"path": path}, + ) + resp.raise_for_status() + + def download(self, path: str) -> bytes: + """Download a small file from the capsule. + + Args: + path: Absolute file path inside the capsule. + + Returns: + File contents as bytes. + """ + assert self._http is not None + resp = self._http.post( + f"/v1/capsules/{self.id}/files/read", + json={"path": path}, + ) + resp.raise_for_status() + return resp.content + + async def async_download(self, path: str) -> bytes: + """Async version of ``download``.""" + assert self._async_http is not None + resp = await self._async_http.post( + f"/v1/capsules/{self.id}/files/read", + json={"path": path}, + ) + resp.raise_for_status() + return resp.content + + def stream_upload(self, path: str, stream: Iterator[bytes]) -> None: + """Streaming upload for large files. + + Args: + path: Absolute destination path inside the capsule. + stream: An iterator yielding byte chunks. + """ + assert self._http is not None + + boundary = os.urandom(16).hex().encode("utf-8") + + def _multipart_stream() -> 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" + + headers = { + "Content-Type": f"multipart/form-data; boundary={boundary.decode('utf-8')}" + } + + resp = self._http.post( + f"/v1/capsules/{self.id}/files/stream/write", + content=_multipart_stream(), + headers=headers, + ) + resp.raise_for_status() + + async def async_stream_upload( + self, path: str, stream: AsyncIterator[bytes] + ) -> None: + """Async version of ``stream_upload``.""" + assert self._async_http is not None + + boundary = os.urandom(16).hex().encode("utf-8") + + async def _async_multipart_stream() -> 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" + + headers = { + "Content-Type": f"multipart/form-data; boundary={boundary.decode('utf-8')}" + } + + resp = await self._async_http.post( + f"/v1/capsules/{self.id}/files/stream/write", + content=_async_multipart_stream(), + headers=headers, + ) + resp.raise_for_status() + + def stream_download(self, path: str) -> Iterator[bytes]: + """Streaming download for large files. + + Args: + path: Absolute file path inside the capsule. + + Yields: + Byte chunks. + """ + assert self._http is not None + with self._http.stream( + "POST", + f"/v1/capsules/{self.id}/files/stream/read", + json={"path": path}, + ) as resp: + resp.raise_for_status() + yield from resp.iter_bytes() + + async def async_stream_download(self, path: str) -> AsyncIterator[bytes]: + """Async version of ``stream_download``.""" + assert self._async_http is not None + async with self._async_http.stream( + "POST", + f"/v1/capsules/{self.id}/files/stream/read", + json={"path": path}, + ) as resp: + resp.raise_for_status() + async for chunk in resp.aiter_bytes(): + yield chunk + + def list_dir(self, path: str, depth: int = 1) -> list[FileEntry]: + """List directory contents inside the capsule. + + Args: + path: Absolute directory path. + depth: Recursion depth. 1 = immediate children only. + + Returns: + List of FileEntry objects with full metadata. + + Raises: + WrennValidationError: Invalid path. + WrennNotFoundError: Capsule or directory not found. + WrennConflictError: Capsule is not running. + WrennAgentError: Agent error. + WrennHostUnavailableError: Host agent not reachable. + """ + assert self._http is not None + resp = self._http.post( + f"/v1/capsules/{self.id}/files/list", + json={"path": path, "depth": depth}, + ) + data = handle_response(resp) + parsed = ListDirResponse.model_validate(data) + return parsed.entries or [] + + async def async_list_dir(self, path: str, depth: int = 1) -> list[FileEntry]: + """Async version of ``list_dir``.""" + assert self._async_http is not None + resp = await self._async_http.post( + f"/v1/capsules/{self.id}/files/list", + json={"path": path, "depth": depth}, + ) + data = handle_response(resp) + parsed = ListDirResponse.model_validate(data) + return parsed.entries or [] + + def mkdir(self, path: str) -> FileEntry: + """Create a directory inside the capsule (with parents). + + Args: + path: Absolute directory path to create. + + Returns: + FileEntry for the created directory. + + Raises: + WrennValidationError: Path exists and is not a directory. + WrennConflictError: Directory already exists (returns existing entry). + Capsule is not running. + WrennNotFoundError: Capsule not found. + WrennAgentError: Agent error. + WrennHostUnavailableError: Host agent not reachable. + """ + assert self._http is not None + resp = self._http.post( + f"/v1/capsules/{self.id}/files/mkdir", + json={"path": path}, + ) + if resp.status_code == 409: + try: + body = resp.json() + err = body.get("error", {}) + if err.get("code") == "conflict": + parent_dir = os.path.dirname(path) + dir_name = os.path.basename(path) + + listing = self.list_dir(parent_dir, depth=0) + for entry in listing: + if entry.name == dir_name: + return entry + except Exception: + pass + data = handle_response(resp) + parsed = MakeDirResponse.model_validate(data) + if parsed.entry is None: + raise RuntimeError("mkdir response missing entry") + return parsed.entry + + async def async_mkdir(self, path: str) -> FileEntry: + """Async version of ``mkdir``.""" + assert self._async_http is not None + resp = await self._async_http.post( + f"/v1/capsules/{self.id}/files/mkdir", + json={"path": path}, + ) + if resp.status_code == 409: + try: + body = resp.json() + err = body.get("error", {}) + if err.get("code") == "conflict": + listing = await self.async_list_dir(path, depth=0) + parent_dir = os.path.dirname(path) + dir_name = os.path.basename(path) + + listing = self.list_dir(parent_dir, depth=0) + for entry in listing: + if entry.name == dir_name: + return entry + except Exception: + pass + data = handle_response(resp) + parsed = MakeDirResponse.model_validate(data) + 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 inside the capsule. + + Removes recursively. No confirmation or dry-run. Equivalent to rm -rf. + + Args: + path: Absolute path to remove. + + Raises: + WrennValidationError: Invalid path. + WrennNotFoundError: Capsule not found. + WrennConflictError: Capsule is not running. + WrennAgentError: Agent error. + WrennHostUnavailableError: Host agent not reachable. + """ + assert self._http is not None + resp = self._http.post( + f"/v1/capsules/{self.id}/files/remove", + json={"path": path}, + ) + handle_response(resp) + + async def async_remove(self, path: str) -> None: + """Async version of ``remove``.""" + assert self._async_http is not None + resp = await self._async_http.post( + f"/v1/capsules/{self.id}/files/remove", + json={"path": path}, + ) + handle_response(resp) + + @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. + + Args: + cmd: Command to run. Defaults to /bin/bash. + args: Command arguments. + cols: Terminal columns. Defaults to 80. + rows: Terminal rows. Defaults to 24. + envs: Environment variables. + cwd: Working directory. + + Returns: + A PtySession context manager. Use with a ``with`` statement. + """ + assert self._http is not None + assert self.id is not None + with httpx_ws.connect_ws( # type: ignore[attr-defined] + f"/v1/capsules/{self.id}/pty", client=self._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. + + Args: + tag: Session tag from a previous PtySession. + + Returns: + A PtySession context manager. + """ + assert self._http is not None + assert self.id is not None + with httpx_ws.connect_ws( + f"/v1/capsules/{self.id}/pty", client=self._http + ) as ws: + session = PtySession(ws, self.id) + session._send_connect(tag) + yield session + + @asynccontextmanager + async def async_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 version of ``pty``.""" + assert self._async_http is not None + assert self.id is not None + async with httpx_ws.aconnect_ws( # type: ignore[attr-defined, misc] + f"/v1/capsules/{self.id}/pty", client=self._async_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 async_pty_connect(self, tag: str) -> AsyncIterator[AsyncPtySession]: + """Async version of ``pty_connect``.""" + assert self._async_http is not None + assert self.id is not None + async with httpx_ws.aconnect_ws( # type: ignore[attr-defined, misc] + f"/v1/capsules/{self.id}/pty", client=self._async_http + ) as ws: + session = AsyncPtySession(ws, self.id) + await session._send_connect(tag) + yield session + + def ping(self) -> None: + """Reset the capsule inactivity timer.""" + assert self._http is not None + resp = self._http.post(f"/v1/capsules/{self.id}/ping") + resp.raise_for_status() + + async def async_ping(self) -> None: + """Async version of ``ping``.""" + assert self._async_http is not None + resp = await self._async_http.post(f"/v1/capsules/{self.id}/ping") + resp.raise_for_status() + + def pause(self) -> Capsule: + """Pause the capsule (snapshot and release resources). + + Returns: + Updated ``Capsule`` with new status. + """ + assert self._http is not None + resp = self._http.post(f"/v1/capsules/{self.id}/pause") + resp.raise_for_status() + updated = Capsule.model_validate(resp.json()) + self.status = updated.status + return self + + async def async_pause(self) -> Capsule: + """Async version of ``pause``.""" + assert self._async_http is not None + resp = await self._async_http.post(f"/v1/capsules/{self.id}/pause") + resp.raise_for_status() + updated = Capsule.model_validate(resp.json()) + self.status = updated.status + return self + + def resume(self) -> Capsule: + """Resume a paused capsule from its snapshot. + + Returns: + Updated ``Capsule`` with new status. + """ + assert self._http is not None + resp = self._http.post(f"/v1/capsules/{self.id}/resume") + resp.raise_for_status() + updated = Capsule.model_validate(resp.json()) + self.status = updated.status + return self + + async def async_resume(self) -> Capsule: + """Async version of ``resume``.""" + assert self._async_http is not None + resp = await self._async_http.post(f"/v1/capsules/{self.id}/resume") + resp.raise_for_status() + updated = Capsule.model_validate(resp.json()) + self.status = updated.status + return self + + def destroy(self) -> None: + """Tear down the capsule.""" + assert self._http is not None + resp = self._http.delete(f"/v1/capsules/{self.id}") + resp.raise_for_status() + + async def async_destroy(self) -> None: + """Async version of ``destroy``.""" + assert self._async_http is not None + resp = await self._async_http.delete(f"/v1/capsules/{self.id}") + resp.raise_for_status() + + def _ensure_kernel(self, jupyter_timeout: float = 30) -> str: + """Ensure a Jupyter kernel is running, creating one if needed. + + Polls the Jupyter server until it responds, then creates a kernel. + + Args: + jupyter_timeout: Maximum seconds to wait for Jupyter to become available. + + Returns: + The kernel ID. + + Raises: + TimeoutError: If Jupyter doesn't respond within the timeout. + """ + current_kernel = self._kernel_id + if current_kernel is not None: + return current_kernel + deadline = time.monotonic() + jupyter_timeout + last_exc: Exception | None = None + while time.monotonic() < deadline: + try: + resp = self.http_client.post("/api/kernels") + if resp.status_code < 500: + resp.raise_for_status() + data = resp.json() + self._kernel_id = data["id"] + return str(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}" + ) + + async def _async_ensure_kernel(self, jupyter_timeout: float = 30) -> str: + """Async version of ``_ensure_kernel``.""" + import asyncio + + current_kernel = self._kernel_id + if current_kernel is not None: + return current_kernel + + if self._async_proxy_client is None: + url = ( + _build_proxy_url(self._base_url, self.id, 8888) + .replace("ws://", "http://") + .replace("wss://", "https://") + ) + self._async_proxy_client = httpx.AsyncClient( + base_url=url, + headers=self._proxy_headers(), + ) + + deadline = time.monotonic() + jupyter_timeout + last_exc: Exception | None = None + while time.monotonic() < deadline: + try: + resp = await self._async_proxy_client.post("/api/kernels") + if resp.status_code < 500: + resp.raise_for_status() + data = resp.json() + self._kernel_id = data["id"] + return str(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._base_url, self.id, 8888) + return f"{proxy}/api/kernels/{kernel_id}/channels" + + def _jupyter_execute_request(self, 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 kernel inside the capsule. + + 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``. + """ + assert self._http is not None + 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 = self._proxy_headers() + + with httpx_ws.connect_ws(ws_url, headers=headers) as ws: # type: ignore[attr-defined, var-annotated] + 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", {}) + result.text = bundle.get("text/plain") + 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 + + return result + + async def async_run_code( + self, + code: str, + language: str = "python", + timeout: float = 30, + jupyter_timeout: float = 30, + ) -> CodeResult: + """Async version of ``run_code``.""" + assert self._async_http is not None + kernel_id = await self._async_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 = self._proxy_headers() + + async with httpx_ws.aconnect_ws(ws_url, headers=headers) as ws: # type: ignore[attr-defined, var-annotated] + 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) # type: ignore[misc] + 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", {}) + result.text = bundle.get("text/plain") + 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 + + return result + + def _cleanup(self) -> None: + if self._proxy_client is not None: + try: + self._proxy_client.close() + except Exception: + pass + self._proxy_client = None + + async def _async_cleanup(self) -> None: + if self._async_proxy_client is not None: + try: + await self._async_proxy_client.aclose() + except Exception: + pass + self._async_proxy_client = None + + def __enter__(self) -> Capsule: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: object, + ) -> None: + try: + self.destroy() + except Exception: + pass + self._cleanup() + + async def __aenter__(self) -> Capsule: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: object, + ) -> None: + try: + await self.async_destroy() + except Exception: + pass + await self._async_cleanup() + + +def __getattr__(name: str) -> type: + if name == "Sandbox": + warnings.warn( + "'Sandbox' is deprecated, use 'Capsule' instead", + DeprecationWarning, + stacklevel=2, + ) + return Capsule + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/wrenn/client.py b/src/wrenn/client.py index bd7fb69..4c06b35 100644 --- a/src/wrenn/client.py +++ b/src/wrenn/client.py @@ -1,10 +1,12 @@ from __future__ import annotations import builtins +import warnings from typing import cast import httpx +from wrenn.capsule import Capsule from wrenn.exceptions import handle_response from wrenn.models import ( APIKeyResponse, @@ -14,9 +16,8 @@ from wrenn.models import ( Template, ) from wrenn.models import ( - Sandbox as SandboxModel, + Capsule as CapsuleModel, ) -from wrenn.sandbox import Sandbox DEFAULT_BASE_URL = "https://api.wrenn.dev" @@ -112,8 +113,8 @@ class AsyncAPIKeysResource: handle_response(resp) -class SandboxesResource: - """Sync sandbox control-plane operations.""" +class CapsulesResource: + """Sync capsule control-plane operations.""" def __init__( self, @@ -133,7 +134,7 @@ class SandboxesResource: vcpus: int | None = None, memory_mb: int | None = None, timeout_sec: int | None = None, - ) -> Sandbox: + ) -> Capsule: payload: dict = {} if template is not None: payload["template"] = template @@ -143,27 +144,27 @@ class SandboxesResource: payload["memory_mb"] = memory_mb if timeout_sec is not None: payload["timeout_sec"] = timeout_sec - resp = self._http.post("/v1/sandboxes", json=payload) - model = SandboxModel.model_validate(handle_response(resp)) - sb = Sandbox.model_validate(model.model_dump()) - sb._bind(self._http, self._base_url, self._api_key, self._token) - return sb + resp = self._http.post("/v1/capsules", json=payload) + model = CapsuleModel.model_validate(handle_response(resp)) + cap = Capsule.model_validate(model.model_dump()) + cap._bind(self._http, self._base_url, self._api_key, self._token) + return cap - def list(self) -> list[SandboxModel]: - resp = self._http.get("/v1/sandboxes") - return [SandboxModel.model_validate(item) for item in 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) -> SandboxModel: - resp = self._http.get(f"/v1/sandboxes/{id}") - return SandboxModel.model_validate(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/sandboxes/{id}") + resp = self._http.delete(f"/v1/capsules/{id}") handle_response(resp) -class AsyncSandboxesResource: - """Async sandbox control-plane operations.""" +class AsyncCapsulesResource: + """Async capsule control-plane operations.""" def __init__( self, @@ -183,7 +184,7 @@ class AsyncSandboxesResource: vcpus: int | None = None, memory_mb: int | None = None, timeout_sec: int | None = None, - ) -> Sandbox: + ) -> Capsule: payload: dict = {} if template is not None: payload["template"] = template @@ -193,22 +194,22 @@ class AsyncSandboxesResource: payload["memory_mb"] = memory_mb if timeout_sec is not None: payload["timeout_sec"] = timeout_sec - resp = await self._http.post("/v1/sandboxes", json=payload) - model = SandboxModel.model_validate(handle_response(resp)) - sb = Sandbox.model_validate(model.model_dump()) - sb._bind(self._http, self._base_url, self._api_key, self._token) - return sb + resp = await self._http.post("/v1/capsules", json=payload) + model = CapsuleModel.model_validate(handle_response(resp)) + cap = Capsule.model_validate(model.model_dump()) + cap._bind(self._http, self._base_url, self._api_key, self._token) + return cap - async def list(self) -> list[SandboxModel]: - resp = await self._http.get("/v1/sandboxes") - return [SandboxModel.model_validate(item) for item in 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) -> SandboxModel: - resp = await self._http.get(f"/v1/sandboxes/{id}") - return SandboxModel.model_validate(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/sandboxes/{id}") + resp = await self._http.delete(f"/v1/capsules/{id}") handle_response(resp) @@ -220,11 +221,11 @@ class SnapshotsResource: def create( self, - sandbox_id: str, + capsule_id: str, name: str | None = None, overwrite: bool = False, ) -> Template: - payload: dict = {"sandbox_id": sandbox_id} + payload: dict = {"sandbox_id": capsule_id} if name is not None: payload["name"] = name params: dict = {} @@ -253,11 +254,11 @@ class AsyncSnapshotsResource: async def create( self, - sandbox_id: str, + capsule_id: str, name: str | None = None, overwrite: bool = False, ) -> Template: - payload: dict = {"sandbox_id": sandbox_id} + payload: dict = {"sandbox_id": capsule_id} if name is not None: payload["name"] = name params: dict = {} @@ -410,10 +411,19 @@ class WrennClient: self.auth = AuthResource(self._http) self.api_keys = APIKeysResource(self._http) - self.sandboxes = SandboxesResource(self._http, base_url, api_key, token) + self.capsules = CapsulesResource(self._http, base_url, api_key, token) self.snapshots = SnapshotsResource(self._http) self.hosts = HostsResource(self._http) + @property + def sandboxes(self) -> CapsulesResource: + warnings.warn( + "'client.sandboxes' is deprecated, use 'client.capsules' instead", + DeprecationWarning, + stacklevel=2, + ) + return self.capsules + def close(self) -> None: """Close the underlying HTTP connection pool.""" self._http.close() @@ -458,10 +468,19 @@ class AsyncWrennClient: self.auth = AsyncAuthResource(self._http) self.api_keys = AsyncAPIKeysResource(self._http) - self.sandboxes = AsyncSandboxesResource(self._http, base_url, api_key, token) + self.capsules = AsyncCapsulesResource(self._http, base_url, api_key, token) self.snapshots = AsyncSnapshotsResource(self._http) self.hosts = AsyncHostsResource(self._http) + @property + def sandboxes(self) -> AsyncCapsulesResource: + warnings.warn( + "'client.sandboxes' is deprecated, use 'client.capsules' instead", + DeprecationWarning, + stacklevel=2, + ) + return self.capsules + async def aclose(self) -> None: """Close the underlying async HTTP connection pool.""" await self._http.aclose() diff --git a/src/wrenn/exceptions.py b/src/wrenn/exceptions.py index 713aff7..c4b39d8 100644 --- a/src/wrenn/exceptions.py +++ b/src/wrenn/exceptions.py @@ -1,5 +1,7 @@ from __future__ import annotations +import warnings + import httpx @@ -33,15 +35,24 @@ class WrennConflictError(WrennError): """409 — State conflict (e.g. invalid_state).""" -class WrennHostHasSandboxesError(WrennConflictError): - """409 — Host still has running sandboxes.""" +class WrennHostHasCapsulesError(WrennConflictError): + """409 — Host still has running capsules.""" def __init__( - self, code: str, message: str, status_code: int, sandbox_ids: list[str] + self, code: str, message: str, status_code: int, capsule_ids: list[str] ) -> None: - self.sandbox_ids = sandbox_ids + 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.""" @@ -62,7 +73,8 @@ _ERROR_MAP: dict[str, type[WrennError]] = { "not_found": WrennNotFoundError, "invalid_state": WrennConflictError, "conflict": WrennConflictError, - "host_has_sandboxes": WrennHostHasSandboxesError, + "host_has_sandboxes": WrennHostHasCapsulesError, + "host_has_capsules": WrennHostHasCapsulesError, "host_unavailable": WrennHostUnavailableError, "agent_error": WrennAgentError, "internal_error": WrennInternalError, @@ -83,12 +95,12 @@ def handle_response(resp: httpx.Response) -> dict | list: exc_cls = _ERROR_MAP.get(code, WrennError) - if exc_cls is WrennHostHasSandboxesError: - raise WrennHostHasSandboxesError( + if exc_cls is WrennHostHasCapsulesError: + raise WrennHostHasCapsulesError( code=code, message=message, status_code=resp.status_code, - sandbox_ids=body.get("sandbox_ids", []), + capsule_ids=body.get("sandbox_ids", []), ) raise exc_cls( @@ -101,3 +113,14 @@ def handle_response(resp: httpx.Response) -> dict | list: 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}") diff --git a/src/wrenn/models/__init__.py b/src/wrenn/models/__init__.py index 7e51557..5628e11 100644 --- a/src/wrenn/models/__init__.py +++ b/src/wrenn/models/__init__.py @@ -1,10 +1,11 @@ from wrenn.models._generated import ( APIKeyResponse, AuthResponse, + Capsule, CreateAPIKeyRequest, + CreateCapsuleRequest, CreateHostRequest, CreateHostResponse, - CreateSandboxRequest, CreateSnapshotRequest, Encoding, Error, @@ -22,7 +23,6 @@ from wrenn.models._generated import ( RegisterHostRequest, RegisterHostResponse, RemoveRequest, - Sandbox, SignupRequest, Status, Status1, @@ -38,7 +38,7 @@ __all__ = [ "CreateAPIKeyRequest", "CreateHostRequest", "CreateHostResponse", - "CreateSandboxRequest", + "CreateCapsuleRequest", "CreateSnapshotRequest", "Encoding", "Error", @@ -56,7 +56,7 @@ __all__ = [ "RegisterHostRequest", "RegisterHostResponse", "RemoveRequest", - "Sandbox", + "Capsule", "SignupRequest", "Status", "Status1", diff --git a/src/wrenn/models/_generated.py b/src/wrenn/models/_generated.py index a211a9b..55a5742 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-11T15:00:55+00:00 +# timestamp: 2026-04-12T20:56:29+00:00 from __future__ import annotations @@ -22,7 +22,7 @@ class LoginRequest(BaseModel): class AuthResponse(BaseModel): - token: Annotated[str | None, Field(description="JWT token (valid for 6 hours)")] = ( + token: Annotated[str | None, Field(description='JWT token (valid for 6 hours)')] = ( None ) user_id: str | None = None @@ -32,7 +32,7 @@ class AuthResponse(BaseModel): class CreateAPIKeyRequest(BaseModel): - name: str | None = "Unnamed API Key" + name: str | None = 'Unnamed API Key' class APIKeyResponse(BaseModel): @@ -47,29 +47,29 @@ class APIKeyResponse(BaseModel): key: Annotated[ str | None, Field( - description="Full plaintext key. Only returned on creation, never again." + description='Full plaintext key. Only returned on creation, never again.' ), ] = None -class CreateSandboxRequest(BaseModel): - template: str | None = "minimal" +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 sandbox is automatically paused after this duration of inactivity (no exec or ping). 0 means no auto-pause.\n" + 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" + field_5m = '5m' + field_1h = '1h' + field_6h = '6h' + field_24h = '24h' + field_30d = '30d' class Current(BaseModel): @@ -100,29 +100,29 @@ class Series(BaseModel): memory_mb: list[int] | None = None -class SandboxStats(BaseModel): +class CapsuleStats(BaseModel): range: Range | None = None current: Current | None = None peaks: Annotated[ - Peaks | None, Field(description="Maximum values over the last 30 days.") + Peaks | None, Field(description='Maximum values over the last 30 days.') ] = None series: Annotated[ - Series | None, Field(description="Parallel arrays for chart rendering.") + 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" + pending = 'pending' + starting = 'starting' + running = 'running' + paused = 'paused' + hibernated = 'hibernated' + stopped = 'stopped' + missing = 'missing' + error = 'error' -class Sandbox(BaseModel): +class Capsule(BaseModel): id: str | None = None status: Status | None = None template: str | None = None @@ -139,17 +139,17 @@ class Sandbox(BaseModel): class CreateSnapshotRequest(BaseModel): sandbox_id: Annotated[ - str, Field(description="ID of the running sandbox to snapshot.") + 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."), + Field(description='Name for the snapshot template. Auto-generated if omitted.'), ] = None class Type(StrEnum): - base = "base" - snapshot = "snapshot" + base = 'base' + snapshot = 'snapshot' class Template(BaseModel): @@ -172,8 +172,8 @@ class Encoding(StrEnum): Output encoding. "base64" when stdout/stderr contain binary data. """ - utf_8 = "utf-8" - base64 = "base64" + utf_8 = 'utf-8' + base64 = 'base64' class ExecResponse(BaseModel): @@ -192,23 +192,23 @@ class ExecResponse(BaseModel): class ReadFileRequest(BaseModel): - path: Annotated[str, Field(description="Absolute file path inside the sandbox")] + path: Annotated[str, Field(description='Absolute file path inside the capsule')] class ListDirRequest(BaseModel): - path: Annotated[str, Field(description="Directory path inside the sandbox")] + path: Annotated[str, Field(description='Directory path inside the capsule')] depth: Annotated[ int | None, Field( - description="Recursion depth (0 = non-recursive, 1 = immediate children)" + description='Recursion depth (0 = non-recursive, 1 = immediate children)' ), ] = 1 class Type1(StrEnum): - file = "file" - directory = "directory" - symlink = "symlink" + file = 'file' + directory = 'directory' + symlink = 'symlink' class FileEntry(BaseModel): @@ -223,14 +223,14 @@ class FileEntry(BaseModel): owner: str | None = None group: str | None = None modified_at: Annotated[ - int | None, Field(description="Unix timestamp (seconds)") + 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 sandbox") + str, Field(description='Directory path to create inside the capsule') ] @@ -239,7 +239,7 @@ class MakeDirResponse(BaseModel): class RemoveRequest(BaseModel): - path: Annotated[str, Field(description="Path to remove inside the sandbox")] + path: Annotated[str, Field(description='Path to remove inside the capsule')] class Type2(StrEnum): @@ -247,51 +247,51 @@ class Type2(StrEnum): Host type. Regular hosts are shared; BYOC hosts belong to a team. """ - regular = "regular" - byoc = "byoc" + regular = 'regular' + byoc = 'byoc' class CreateHostRequest(BaseModel): type: Annotated[ Type2, Field( - description="Host type. Regular hosts are shared; BYOC hosts belong to a team." + 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 + 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)."), + 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).") + 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.") + str, Field(description='One-time registration token from POST /v1/hosts.') ] arch: Annotated[ - str | None, Field(description="CPU architecture (e.g. x86_64, aarch64).") + 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).")] + address: Annotated[str, Field(description='Host agent address (ip:port).')] class Type3(StrEnum): - regular = "regular" - byoc = "byoc" + regular = 'regular' + byoc = 'byoc' class Status1(StrEnum): - pending = "pending" - online = "online" - offline = "offline" - draining = "draining" - unreachable = "unreachable" + pending = 'pending' + online = 'online' + offline = 'offline' + draining = 'draining' + unreachable = 'unreachable' class Host(BaseModel): @@ -316,7 +316,7 @@ class RefreshHostTokenRequest(BaseModel): refresh_token: Annotated[ str, Field( - description="Refresh token obtained from registration or a previous refresh." + description='Refresh token obtained from registration or a previous refresh.' ), ] @@ -324,12 +324,12 @@ class RefreshHostTokenRequest(BaseModel): class RefreshHostTokenResponse(BaseModel): host: Host | None = None token: Annotated[ - str | None, Field(description="New host JWT. Valid for 7 days.") + 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." + description='New refresh token. Valid for 60 days; old token is revoked.' ), ] = None @@ -338,20 +338,20 @@ class HostDeletePreview(BaseModel): host: Host | None = None sandbox_ids: Annotated[ list[str] | None, - Field(description="IDs of sandboxes that would be destroyed on force-delete."), + 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 + 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 sandboxes blocking deletion."), + Field(description='IDs of active capsulees blocking deletion.'), ] = None -class HostHasSandboxesError(BaseModel): +class HostHasCapsulesError(BaseModel): error: Error | None = None @@ -368,15 +368,15 @@ 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)") + 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" + owner = 'owner' + admin = 'admin' + member = 'member' class TeamWithRole(Team): @@ -396,13 +396,13 @@ class TeamDetail(BaseModel): class Range1(StrEnum): - field_5m = "5m" - field_10m = "10m" - field_1h = "1h" - field_2h = "2h" - field_6h = "6h" - field_12h = "12h" - field_24h = "24h" + field_5m = '5m' + field_10m = '10m' + field_1h = '1h' + field_2h = '2h' + field_6h = '6h' + field_12h = '12h' + field_24h = '24h' class MetricPoint(BaseModel): @@ -410,41 +410,41 @@ class MetricPoint(BaseModel): cpu_pct: Annotated[ float | None, Field( - description="CPU utilization percentage (0-100), normalized to vCPU count" + 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)"), + 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") + 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" + 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" + 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.")] + name: Annotated[str, Field(description='Unique channel name within the team.')] provider: Provider config: Annotated[ dict[str, str], @@ -460,7 +460,7 @@ class TestChannelRequest(BaseModel): config: Annotated[ dict[str, str], Field( - description="Provider-specific configuration fields (same as CreateChannelRequest.config)." + description='Provider-specific configuration fields (same as CreateChannelRequest.config).' ), ] @@ -489,7 +489,7 @@ class ChannelResponse(BaseModel): updated_at: AwareDatetime | None = None secret: Annotated[ str | None, - Field(description="Webhook secret. Only returned on creation, never again."), + Field(description='Webhook secret. Only returned on creation, never again.'), ] = None @@ -511,7 +511,7 @@ class CreateHostResponse(BaseModel): registration_token: Annotated[ str | None, Field( - description="One-time registration token for the host agent. Expires in 1 hour." + description='One-time registration token for the host agent. Expires in 1 hour.' ), ] = None @@ -520,17 +520,17 @@ class RegisterHostResponse(BaseModel): host: Host | None = None token: Annotated[ str | None, - Field(description="Host JWT for X-Host-Token header. Valid for 7 days."), + 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." + description='Refresh token for obtaining new JWTs. Valid for 60 days; rotated on each use.' ), ] = None -class SandboxMetrics(BaseModel): +class CapsuleMetrics(BaseModel): sandbox_id: str | None = None range: Range1 | None = None points: list[MetricPoint] | None = None diff --git a/src/wrenn/pty.py b/src/wrenn/pty.py index cde476c..83ee871 100644 --- a/src/wrenn/pty.py +++ b/src/wrenn/pty.py @@ -66,9 +66,9 @@ class PtySession: break """ - def __init__(self, ws: httpx_ws.WebSocketSession, sandbox_id: str) -> None: + def __init__(self, ws: httpx_ws.WebSocketSession, capsule_id: str) -> None: self._ws = ws - self._sandbox_id = sandbox_id + self._capsule_id = capsule_id self._tag: str | None = None self._pid: int | None = None self._done = False @@ -192,9 +192,9 @@ class AsyncPtySession: break """ - def __init__(self, ws: httpx_ws.AsyncWebSocketSession, sandbox_id: str) -> None: + def __init__(self, ws: httpx_ws.AsyncWebSocketSession, capsule_id: str) -> None: self._ws = ws - self._sandbox_id = sandbox_id + self._capsule_id = capsule_id self._tag: str | None = None self._pid: int | None = None self._done = False diff --git a/src/wrenn/sandbox.py b/src/wrenn/sandbox.py index 09b40de..09126f8 100644 --- a/src/wrenn/sandbox.py +++ b/src/wrenn/sandbox.py @@ -1,1181 +1,26 @@ -from __future__ import annotations +import warnings as _warnings -import asyncio -import base64 -import json -import os -import time -import uuid -from collections.abc import AsyncIterator, Iterator -from contextlib import asynccontextmanager, contextmanager -from typing import Any - -import httpx -import httpx_ws - -from wrenn.exceptions import handle_response -from wrenn.models import ( - ExecResponse, - FileEntry, - ListDirResponse, - MakeDirResponse, - Status, +from wrenn.capsule import ( # noqa: F401 + CodeResult, + ExecResult, + StreamErrorEvent, + StreamEvent, + StreamExitEvent, + StreamStartEvent, + StreamStderrEvent, + StreamStdoutEvent, + _build_proxy_url, + _parse_stream_event, ) -from wrenn.models import Sandbox as SandboxModel -from wrenn.pty import AsyncPtySession, PtySession +from wrenn.capsule import Capsule -class _IterableReader: - """Internal adapter to make iterables/generators act like files with a . - read() method""" - - def __init__(self, iterable: Any) -> None: - self.iterator = iter(iterable) - self.buffer = b"" - - def read(self, size: int = -1) -> bytes: - if size == -1: - return self.buffer + b"".join( - chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") - for chunk in self.iterator - ) - - while len(self.buffer) < size: - try: - chunk = next(self.iterator) - self.buffer += ( - chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") - ) - except StopIteration: - break - - result = self.buffer[:size] - self.buffer = self.buffer[size:] - return result - - -class ExecResult: - """Typed result from a synchronous exec call.""" - - __slots__ = ("stdout", "stderr", "exit_code", "duration_ms", "encoding") - - def __init__( - self, - stdout: str, - stderr: str, - exit_code: int, - duration_ms: int | None, - encoding: str | None, - ) -> None: - self.stdout = stdout - self.stderr = stderr - self.exit_code = exit_code - self.duration_ms = duration_ms - self.encoding = encoding - - -class CodeResult: - """Typed result from stateful code execution (``run_code``). - - 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. - """ - - __slots__ = ("text", "data", "stdout", "stderr", "error") - - def __init__( - self, - text: str | None = None, - data: dict[str, str] | None = None, - stdout: str = "", - stderr: str = "", - error: str | None = None, - ) -> None: - self.text = text - self.data = data - self.stdout = stdout - self.stderr = stderr - self.error = error - - -class StreamEvent: - """Base class for streaming exec events.""" - - __slots__ = ("type",) - - def __init__(self, type: str) -> None: - self.type = type - - -class StreamStartEvent(StreamEvent): - """Process started.""" - - __slots__ = ("pid",) - - def __init__(self, pid: int) -> None: - super().__init__("start") - self.pid = pid - - -class StreamStdoutEvent(StreamEvent): - """Stdout data received.""" - - __slots__ = ("data",) - - def __init__(self, data: str) -> None: - super().__init__("stdout") - self.data = data - - -class StreamStderrEvent(StreamEvent): - """Stderr data received.""" - - __slots__ = ("data",) - - def __init__(self, data: str) -> None: - super().__init__("stderr") - self.data = data - - -class StreamExitEvent(StreamEvent): - """Process exited.""" - - __slots__ = ("exit_code",) - - def __init__(self, exit_code: int) -> None: - super().__init__("exit") - self.exit_code = exit_code - - -class StreamErrorEvent(StreamEvent): - """Error occurred.""" - - __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 _build_proxy_url(base_url: str, sandbox_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}-{sandbox_id}.{host}" - - -class Sandbox(SandboxModel): - """Developer-facing sandbox interface wrapping the generated Sandbox model. - - Provides data-plane methods (exec, file I/O, lifecycle), sandbox proxy - helpers, and context-manager support for automatic cleanup. - """ - - _http: httpx.Client | None - _async_http: httpx.AsyncClient | None - _base_url: str - _api_key: str | None - _token: str | None - _proxy_client: httpx.Client | None - _async_proxy_client: httpx.AsyncClient | None - _kernel_id: str | None - _jupyter_ws: Any - _async_jupyter_ws: Any - - def _bind( - self, - http: httpx.Client | httpx.AsyncClient, - base_url: str, - api_key: str | None = None, - token: str | None = None, - ) -> None: - self._base_url = base_url - self._api_key = api_key - self._token = token - self._proxy_client = None - self._async_proxy_client = None - self._kernel_id = None - self._jupyter_ws = None - self._async_jupyter_ws = None - if isinstance(http, httpx.Client): - self._http = http - self._async_http = None - else: - self._http = None # type: ignore[assignment] - self._async_http = http - - def _proxy_headers(self) -> dict[str, str]: - headers: dict[str, str] = {} - if self._api_key: - headers["X-API-Key"] = self._api_key - if self._token: - headers["Authorization"] = f"Bearer {self._token}" - return headers - - def _clear_content_type(self) -> dict[str, str]: - assert self._http is not None - headers = dict(self._http.headers) - headers.pop("Content-Type", None) - return headers - - def _async_clear_content_type(self) -> dict[str, str]: - assert self._async_http is not None - headers = dict(self._async_http.headers) - headers.pop("Content-Type", None) - return headers - - def get_url(self, port: int) -> str: - """Construct the proxy URL for a port inside this sandbox. - - Args: - port: Port number of the service running inside the sandbox. - - Returns: - A URL string like ``http://8888-cl-abc123.api.wrenn.dev``. - """ - return _build_proxy_url(self._base_url, self.id, port) - - @property - def http_client(self) -> httpx.Client: - """A pre-configured ``httpx.Client`` targeting the sandbox proxy on port 8888. - - The client has auth headers set and ``base_url`` pointing to - the proxy URL for port 8888. Closed automatically when the sandbox exits. - """ - if self._proxy_client is None: - url = ( - _build_proxy_url(self._base_url, self.id, 8888) - .replace("ws://", "http://") - .replace("wss://", "https://") - ) - self._proxy_client = httpx.Client( - base_url=url, - headers=self._proxy_headers(), - ) - return self._proxy_client - - def wait_ready(self, timeout: float = 30, interval: float = 0.5) -> None: - """Block until the sandbox status is ``running``. - - Args: - timeout: Maximum seconds to wait. - interval: Seconds between polls. - - Raises: - TimeoutError: If the sandbox does not become ready in time. - """ - assert self._http is not None - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - resp = self._http.get(f"/v1/sandboxes/{self.id}") - data = resp.json() - status = data.get("status") - if status == Status.running: - self.status = Status.running - return - if status in (Status.error, Status.stopped): - raise RuntimeError(f"Sandbox entered {status} state while waiting") - time.sleep(interval) - raise TimeoutError(f"Sandbox {self.id} did not become ready within {timeout}s") - - async def async_wait_ready( - self, timeout: float = 30, interval: float = 0.5 - ) -> None: - """Async version of ``wait_ready``.""" - assert self._async_http is not None - import asyncio - - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - resp = await self._async_http.get(f"/v1/sandboxes/{self.id}") - data = resp.json() - status = data.get("status") - if status == Status.running: - self.status = Status.running - return - if status in (Status.error, Status.stopped): - raise RuntimeError(f"Sandbox entered {status} state while waiting") - await asyncio.sleep(interval) - raise TimeoutError(f"Sandbox {self.id} did not become ready within {timeout}s") - - def exec( - self, - cmd: str, - args: list[str] | None = None, - timeout_sec: int | None = 30, - ) -> ExecResult: - """Execute a command synchronously inside the sandbox. - - Args: - cmd: Command to run. - args: Optional positional arguments. - timeout_sec: Execution timeout in seconds. - - Returns: - An ``ExecResult`` with ``stdout``, ``stderr``, ``exit_code``, ``duration_ms``. - """ - assert self._http is not None - payload: dict = {"cmd": cmd} - if args is not None: - payload["args"] = args - if timeout_sec is not None: - payload["timeout_sec"] = timeout_sec - resp = self._http.post(f"/v1/sandboxes/{self.id}/exec", json=payload) - resp.raise_for_status() - er = ExecResponse.model_validate(resp.json()) - stdout = er.stdout or "" - stderr = er.stderr or "" - if er.encoding == "base64": - stdout = base64.b64decode(stdout).decode("utf-8", errors="replace") - if stderr: - stderr = base64.b64decode(stderr).decode("utf-8", errors="replace") - return ExecResult( - stdout=stdout, - stderr=stderr, - exit_code=er.exit_code if er.exit_code is not None else -1, - duration_ms=er.duration_ms, - encoding=er.encoding, +def __getattr__(name: str) -> type: + if name == "Sandbox": + _warnings.warn( + "'Sandbox' is deprecated, use 'Capsule' instead", + DeprecationWarning, + stacklevel=2, ) - - async def async_exec( - self, - cmd: str, - args: list[str] | None = None, - timeout_sec: int | None = 30, - ) -> ExecResult: - """Async version of ``exec``.""" - assert self._async_http is not None - payload: dict = {"cmd": cmd} - if args is not None: - payload["args"] = args - if timeout_sec is not None: - payload["timeout_sec"] = timeout_sec - resp = await self._async_http.post( - f"/v1/sandboxes/{self.id}/exec", json=payload - ) - resp.raise_for_status() - er = ExecResponse.model_validate(resp.json()) - stdout = er.stdout or "" - stderr = er.stderr or "" - if er.encoding == "base64": - stdout = base64.b64decode(stdout).decode("utf-8", errors="replace") - if stderr: - stderr = base64.b64decode(stderr).decode("utf-8", errors="replace") - return ExecResult( - stdout=stdout, - stderr=stderr, - exit_code=er.exit_code if er.exit_code is not None else -1, - duration_ms=er.duration_ms, - encoding=er.encoding, - ) - - def exec_stream( - self, - cmd: str, - args: list[str] | None = None, - ) -> Iterator[StreamEvent]: - """Execute a command via WebSocket, yielding ``StreamEvent`` objects. - - Args: - cmd: Command to run. - args: Optional positional arguments. - - Yields: - ``StreamStartEvent``, ``StreamStdoutEvent``, ``StreamStderrEvent``, - ``StreamExitEvent``, or ``StreamErrorEvent``. - """ - assert self._http is not None - with httpx_ws.connect_ws( # type: ignore[attr-defined] - f"/v1/sandboxes/{self.id}/exec/stream", - self._http, - ) as ws: - start_msg: dict = {"type": "start", "cmd": cmd} - if args: - start_msg["args"] = args - ws.send(json.dumps(start_msg)) - for raw_msg in ws: - event = _parse_stream_event(json.loads(raw_msg)) - yield event - if event.type in ("exit", "error"): - break - - async def async_exec_stream( - self, cmd: str, args: list[str] | None = None - ) -> AsyncIterator[StreamEvent]: - """Async version of ``exec_stream``.""" - assert self._async_http is not None - async with httpx_ws.aconnect_ws( # type: ignore[attr-defined, var-annotated] - f"/v1/sandboxes/{self.id}/exec/stream", self._async_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_data = await ws.receive_json() - event = _parse_stream_event(raw_data) - yield event - - if event.type in ("exit", "error"): - break - except httpx_ws.WebSocketDisconnect: - pass - - def upload(self, path: str, data: bytes) -> None: - """Upload a small file to the sandbox. - - Args: - path: Absolute destination path inside the sandbox. - data: File contents as bytes. - """ - assert self._http is not None - resp = self._http.post( - f"/v1/sandboxes/{self.id}/files/write", - files={"file": ("upload", data)}, - data={"path": path}, - ) - - resp.raise_for_status() - - async def async_upload(self, path: str, data: bytes) -> None: - """Async version of ``upload``.""" - assert self._async_http is not None - resp = await self._async_http.post( - f"/v1/sandboxes/{self.id}/files/write", - files={"file": ("upload", data)}, - data={"path": path}, - ) - resp.raise_for_status() - - def download(self, path: str) -> bytes: - """Download a small file from the sandbox. - - Args: - path: Absolute file path inside the sandbox. - - Returns: - File contents as bytes. - """ - assert self._http is not None - resp = self._http.post( - f"/v1/sandboxes/{self.id}/files/read", - json={"path": path}, - ) - resp.raise_for_status() - return resp.content - - async def async_download(self, path: str) -> bytes: - """Async version of ``download``.""" - assert self._async_http is not None - resp = await self._async_http.post( - f"/v1/sandboxes/{self.id}/files/read", - json={"path": path}, - ) - resp.raise_for_status() - return resp.content - - def stream_upload(self, path: str, stream: Iterator[bytes]) -> None: - """Streaming upload for large files. - - Args: - path: Absolute destination path inside the sandbox. - stream: An iterator yielding byte chunks. - """ - assert self._http is not None - - boundary = os.urandom(16).hex().encode("utf-8") - - def _multipart_stream() -> 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" - - headers = { - "Content-Type": f"multipart/form-data; boundary={boundary.decode('utf-8')}" - } - - resp = self._http.post( - f"/v1/sandboxes/{self.id}/files/stream/write", - content=_multipart_stream(), - headers=headers, - ) - resp.raise_for_status() - - async def async_stream_upload( - self, path: str, stream: AsyncIterator[bytes] - ) -> None: - """Async version of ``stream_upload``.""" - assert self._async_http is not None - - boundary = os.urandom(16).hex().encode("utf-8") - - async def _async_multipart_stream() -> 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" - - headers = { - "Content-Type": f"multipart/form-data; boundary={boundary.decode('utf-8')}" - } - - # Use content= and headers= just like the sync version - resp = await self._async_http.post( - f"/v1/sandboxes/{self.id}/files/stream/write", - content=_async_multipart_stream(), - headers=headers, - ) - resp.raise_for_status() - - def stream_download(self, path: str) -> Iterator[bytes]: - """Streaming download for large files. - - Args: - path: Absolute file path inside the sandbox. - - Yields: - Byte chunks. - """ - assert self._http is not None - with self._http.stream( - "POST", - f"/v1/sandboxes/{self.id}/files/stream/read", - json={"path": path}, - ) as resp: - resp.raise_for_status() - yield from resp.iter_bytes() - - async def async_stream_download(self, path: str) -> AsyncIterator[bytes]: - """Async version of ``stream_download``.""" - assert self._async_http is not None - async with self._async_http.stream( - "POST", - f"/v1/sandboxes/{self.id}/files/stream/read", - json={"path": path}, - ) as resp: - resp.raise_for_status() - async for chunk in resp.aiter_bytes(): - yield chunk - - def list_dir(self, path: str, depth: int = 1) -> list[FileEntry]: - """List directory contents inside the sandbox. - - Args: - path: Absolute directory path. - depth: Recursion depth. 1 = immediate children only. - - Returns: - List of FileEntry objects with full metadata. - - Raises: - WrennValidationError: Invalid path. - WrennNotFoundError: Sandbox or directory not found. - WrennConflictError: Sandbox is not running. - WrennAgentError: Agent error. - WrennHostUnavailableError: Host agent not reachable. - """ - assert self._http is not None - resp = self._http.post( - f"/v1/sandboxes/{self.id}/files/list", - json={"path": path, "depth": depth}, - ) - data = handle_response(resp) - parsed = ListDirResponse.model_validate(data) - return parsed.entries or [] - - async def async_list_dir(self, path: str, depth: int = 1) -> list[FileEntry]: - """Async version of ``list_dir``.""" - assert self._async_http is not None - resp = await self._async_http.post( - f"/v1/sandboxes/{self.id}/files/list", - json={"path": path, "depth": depth}, - ) - data = handle_response(resp) - parsed = ListDirResponse.model_validate(data) - return parsed.entries or [] - - def mkdir(self, path: str) -> FileEntry: - """Create a directory inside the sandbox (with parents). - - Args: - path: Absolute directory path to create. - - Returns: - FileEntry for the created directory. - - Raises: - WrennValidationError: Path exists and is not a directory. - WrennConflictError: Directory already exists (returns existing entry). - Sandbox is not running. - WrennNotFoundError: Sandbox not found. - WrennAgentError: Agent error. - WrennHostUnavailableError: Host agent not reachable. - """ - assert self._http is not None - resp = self._http.post( - f"/v1/sandboxes/{self.id}/files/mkdir", - json={"path": path}, - ) - if resp.status_code == 409: - try: - body = resp.json() - err = body.get("error", {}) - if err.get("code") == "conflict": - parent_dir = os.path.dirname(path) - dir_name = os.path.basename(path) - - listing = self.list_dir(parent_dir, depth=0) - for entry in listing: - if entry.name == dir_name: - return entry - except Exception: - pass - data = handle_response(resp) - parsed = MakeDirResponse.model_validate(data) - entry = parsed.entry - if entry is None: - raise RuntimeError("mkdir response missing entry") - return entry - - async def async_mkdir(self, path: str) -> FileEntry: - """Async version of ``mkdir``.""" - assert self._async_http is not None - resp = await self._async_http.post( - f"/v1/sandboxes/{self.id}/files/mkdir", - json={"path": path}, - ) - if resp.status_code == 409: - try: - body = resp.json() - err = body.get("error", {}) - if err.get("code") == "conflict": - listing = await self.async_list_dir(path, depth=0) - parent_dir = os.path.dirname(path) - dir_name = os.path.basename(path) - - listing = self.list_dir(parent_dir, depth=0) - for entry in listing: - if entry.name == dir_name: - return entry - except Exception: - pass - data = handle_response(resp) - parsed = MakeDirResponse.model_validate(data) - entry = parsed.entry - if entry is None: - raise RuntimeError("mkdir response missing entry") - return entry - - def remove(self, path: str) -> None: - """Remove a file or directory inside the sandbox. - - Removes recursively. No confirmation or dry-run. Equivalent to rm -rf. - - Args: - path: Absolute path to remove. - - Raises: - WrennValidationError: Invalid path. - WrennNotFoundError: Sandbox not found. - WrennConflictError: Sandbox is not running. - WrennAgentError: Agent error. - WrennHostUnavailableError: Host agent not reachable. - """ - assert self._http is not None - resp = self._http.post( - f"/v1/sandboxes/{self.id}/files/remove", - json={"path": path}, - ) - handle_response(resp) - - async def async_remove(self, path: str) -> None: - """Async version of ``remove``.""" - assert self._async_http is not None - resp = await self._async_http.post( - f"/v1/sandboxes/{self.id}/files/remove", - json={"path": path}, - ) - handle_response(resp) - - @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, - ) -> PtySession: - """Open an interactive PTY session. - - Args: - cmd: Command to run. Defaults to /bin/bash. - args: Command arguments. - cols: Terminal columns. Defaults to 80. - rows: Terminal rows. Defaults to 24. - envs: Environment variables. - cwd: Working directory. - - Returns: - A PtySession context manager. Use with a ``with`` statement. - """ - assert self._http is not None - with httpx_ws.connect_ws( - f"/v1/sandboxes/{self.id}/pty", client=self._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) -> PtySession: - """Reconnect to an existing PTY session. - - Args: - tag: Session tag from a previous PtySession. - - Returns: - A PtySession context manager. - """ - assert self._http is not None - with httpx_ws.connect_ws( - f"/v1/sandboxes/{self.id}/pty", client=self._http - ) as ws: - session = PtySession(ws, self.id) - session._send_connect(tag) - yield session - - @asynccontextmanager - async def async_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, - ) -> AsyncPtySession: - """Async version of ``pty``.""" - assert self._async_http is not None - with await httpx_ws.aconnect_ws( - f"/v1/sandboxes/{self.id}/pty", client=self._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 async_pty_connect(self, tag: str) -> AsyncPtySession: - """Async version of ``pty_connect``.""" - assert self._async_http is not None - with await httpx_ws.aconnect_ws( - f"/v1/sandboxes/{self.id}/pty", client=self._http - ) as ws: - session = AsyncPtySession(ws, self.id) - await session._send_connect(tag) - yield session - - def ping(self) -> None: - """Reset the sandbox inactivity timer.""" - assert self._http is not None - resp = self._http.post(f"/v1/sandboxes/{self.id}/ping") - resp.raise_for_status() - - async def async_ping(self) -> None: - """Async version of ``ping``.""" - assert self._async_http is not None - resp = await self._async_http.post(f"/v1/sandboxes/{self.id}/ping") - resp.raise_for_status() - - def pause(self) -> Sandbox: - """Pause the sandbox (snapshot and release resources). - - Returns: - Updated ``Sandbox`` with new status. - """ - assert self._http is not None - resp = self._http.post(f"/v1/sandboxes/{self.id}/pause") - resp.raise_for_status() - updated = Sandbox.model_validate(resp.json()) - self.status = updated.status - return self - - async def async_pause(self) -> Sandbox: - """Async version of ``pause``.""" - assert self._async_http is not None - resp = await self._async_http.post(f"/v1/sandboxes/{self.id}/pause") - resp.raise_for_status() - updated = Sandbox.model_validate(resp.json()) - self.status = updated.status - return self - - def resume(self) -> Sandbox: - """Resume a paused sandbox from its snapshot. - - Returns: - Updated ``Sandbox`` with new status. - """ - assert self._http is not None - resp = self._http.post(f"/v1/sandboxes/{self.id}/resume") - resp.raise_for_status() - updated = Sandbox.model_validate(resp.json()) - self.status = updated.status - return self - - async def async_resume(self) -> Sandbox: - """Async version of ``resume``.""" - assert self._async_http is not None - resp = await self._async_http.post(f"/v1/sandboxes/{self.id}/resume") - resp.raise_for_status() - updated = Sandbox.model_validate(resp.json()) - self.status = updated.status - return self - - def destroy(self) -> None: - """Tear down the sandbox.""" - assert self._http is not None - resp = self._http.delete(f"/v1/sandboxes/{self.id}") - resp.raise_for_status() - - async def async_destroy(self) -> None: - """Async version of ``destroy``.""" - assert self._async_http is not None - resp = await self._async_http.delete(f"/v1/sandboxes/{self.id}") - resp.raise_for_status() - - def _ensure_kernel(self, jupyter_timeout: float = 30) -> str: - """Ensure a Jupyter kernel is running, creating one if needed. - - Polls the Jupyter server until it responds, then creates a kernel. - - Args: - jupyter_timeout: Maximum seconds to wait for Jupyter to become available. - - Returns: - The kernel ID. - - Raises: - TimeoutError: If Jupyter doesn't respond within the timeout. - """ - current_kernel = self._kernel_id - if current_kernel is not None: - return current_kernel - deadline = time.monotonic() + jupyter_timeout - last_exc: Exception | None = None - while time.monotonic() < deadline: - try: - resp = self.http_client.post("/api/kernels") - if resp.status_code < 500: - resp.raise_for_status() - data = resp.json() - self._kernel_id = data["id"] - return str(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}" - ) - - async def _async_ensure_kernel(self, jupyter_timeout: float = 30) -> str: - """Async version of ``_ensure_kernel``.""" - import asyncio - - current_kernel = self._kernel_id - if current_kernel is not None: - return current_kernel - - if self._async_proxy_client is None: - url = ( - _build_proxy_url(self._base_url, self.id, 8888) - .replace("ws://", "http://") - .replace("wss://", "https://") - ) - self._async_proxy_client = httpx.AsyncClient( - base_url=url, - headers=self._proxy_headers(), - ) - - deadline = time.monotonic() + jupyter_timeout - last_exc: Exception | None = None - while time.monotonic() < deadline: - try: - resp = await self._async_proxy_client.post("/api/kernels") - if resp.status_code < 500: - resp.raise_for_status() - data = resp.json() - self._kernel_id = data["id"] - return str(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._base_url, self.id, 8888) - return f"{proxy}/api/kernels/{kernel_id}/channels" - - def _jupyter_execute_request(self, 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 kernel inside the sandbox. - - 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``. - """ - assert self._http is not None - 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 = self._proxy_headers() - - with httpx_ws.connect_ws(ws_url, headers=headers) as ws: # type: ignore[attr-defined, var-annotated] - 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", {}) - result.text = bundle.get("text/plain") - 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 - - return result - - async def async_run_code( - self, - code: str, - language: str = "python", - timeout: float = 30, - jupyter_timeout: float = 30, - ) -> CodeResult: - """Async version of ``run_code``.""" - assert self._async_http is not None - kernel_id = await self._async_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 = self._proxy_headers() - - async with httpx_ws.aconnect_ws(ws_url, headers=headers) as ws: # type: ignore[attr-defined, var-annotated] - 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) # type: ignore[misc] - 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", {}) - result.text = bundle.get("text/plain") - 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 - - return result - - def _cleanup(self) -> None: - if self._proxy_client is not None: - try: - self._proxy_client.close() - except Exception: - pass - self._proxy_client = None - - async def _async_cleanup(self) -> None: - if self._async_proxy_client is not None: - try: - await self._async_proxy_client.aclose() - except Exception: - pass - self._async_proxy_client = None - - def __enter__(self) -> Sandbox: - return self - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc_val: BaseException | None, - exc_tb: object, - ) -> None: - try: - self.destroy() - except Exception: - pass - self._cleanup() - - async def __aenter__(self) -> Sandbox: - return self - - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc_val: BaseException | None, - exc_tb: object, - ) -> None: - try: - await self.async_destroy() - except Exception: - pass - await self._async_cleanup() + return Capsule + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/tests/test_sandbox_features.py b/tests/test_capsule_features.py similarity index 53% rename from tests/test_sandbox_features.py rename to tests/test_capsule_features.py index 7737b45..594a378 100644 --- a/tests/test_sandbox_features.py +++ b/tests/test_capsule_features.py @@ -1,11 +1,10 @@ from __future__ import annotations - import pytest import respx +from wrenn.capsule import Capsule, CodeResult, _build_proxy_url from wrenn.client import WrennClient -from wrenn.sandbox import CodeResult, Sandbox, _build_proxy_url @pytest.fixture @@ -32,14 +31,14 @@ class TestBuildProxyUrl: assert url == "ws://5000-sb-2.192.168.1.1" -class TestSandboxGetUrl: +class TestCapsuleGetUrl: @respx.mock def test_get_url_returns_proxy_url(self, client): - respx.post("https://api.wrenn.dev/v1/sandboxes").respond( + respx.post("https://api.wrenn.dev/v1/capsules").respond( 201, json={"id": "cl-abc", "status": "pending"} ) - sb = client.sandboxes.create(template="minimal") - url = sb.get_url(8888) + cap = client.capsules.create(template="minimal") + url = cap.get_url(8888) assert url == "wss://8888-cl-abc.api.wrenn.dev" @respx.mock @@ -48,22 +47,22 @@ class TestSandboxGetUrl: api_key="wrn_test1234567890abcdef12345678", base_url="http://localhost:8080", ) as c: - respx.post("http://localhost:8080/v1/sandboxes").respond( + respx.post("http://localhost:8080/v1/capsules").respond( 201, json={"id": "cl-xyz", "status": "pending"} ) - sb = c.sandboxes.create() - url = sb.get_url(3000) + cap = c.capsules.create() + url = cap.get_url(3000) assert url == "ws://3000-cl-xyz.localhost:8080" -class TestSandboxHttpClient: +class TestCapsuleHttpClient: @respx.mock def test_http_client_has_api_key_header(self, client): - respx.post("https://api.wrenn.dev/v1/sandboxes").respond( + respx.post("https://api.wrenn.dev/v1/capsules").respond( 201, json={"id": "cl-abc", "status": "pending"} ) - sb = client.sandboxes.create() - hc = sb.http_client + cap = client.capsules.create() + hc = cap.http_client assert hc.headers["X-API-Key"] == "wrn_test1234567890abcdef12345678" @respx.mock @@ -71,51 +70,51 @@ class TestSandboxHttpClient: route = respx.get("https://8888-cl-abc.api.wrenn.dev/api/kernels").respond( 200, json=[] ) - respx.post("https://api.wrenn.dev/v1/sandboxes").respond( + respx.post("https://api.wrenn.dev/v1/capsules").respond( 201, json={"id": "cl-abc", "status": "pending"} ) - sb = client.sandboxes.create() - resp = sb.http_client.get("/api/kernels") + cap = client.capsules.create() + resp = cap.http_client.get("/api/kernels") assert resp.status_code == 200 assert route.called def test_jwt_only_get_url_works(self): with WrennClient(token="jwt-abc") as c: - sb = Sandbox(id="cl-abc") - sb._bind(c._http, str(c._http.base_url), api_key=None, token="jwt-abc") - url = sb.get_url(8888) + cap = Capsule(id="cl-abc") + cap._bind(c._http, str(c._http.base_url), api_key=None, token="jwt-abc") + url = cap.get_url(8888) assert "8888-cl-abc" in url def test_jwt_only_http_client_has_bearer_header(self): with WrennClient(token="jwt-abc") as c: - sb = Sandbox(id="cl-abc") - sb._bind(c._http, str(c._http.base_url), api_key=None, token="jwt-abc") - hc = sb.http_client + cap = Capsule(id="cl-abc") + cap._bind(c._http, str(c._http.base_url), api_key=None, token="jwt-abc") + hc = cap.http_client assert hc.headers["Authorization"] == "Bearer jwt-abc" -class TestCreateReturnsBoundSandbox: +class TestCreateReturnsBoundCapsule: @respx.mock - def test_create_returns_sandbox_subclass(self, client): - respx.post("https://api.wrenn.dev/v1/sandboxes").respond( + def test_create_returns_capsule_subclass(self, client): + respx.post("https://api.wrenn.dev/v1/capsules").respond( 201, json={"id": "cl-1", "status": "pending", "template": "minimal"} ) - sb = client.sandboxes.create(template="minimal") - assert isinstance(sb, Sandbox) - assert sb.id == "cl-1" - assert hasattr(sb, "exec") - assert hasattr(sb, "run_code") - assert hasattr(sb, "get_url") + cap = client.capsules.create(template="minimal") + assert isinstance(cap, Capsule) + assert cap.id == "cl-1" + assert hasattr(cap, "exec") + assert hasattr(cap, "run_code") + assert hasattr(cap, "get_url") @respx.mock def test_create_context_manager(self, client): - route = respx.delete("https://api.wrenn.dev/v1/sandboxes/cl-1").respond(204) - respx.post("https://api.wrenn.dev/v1/sandboxes").respond( + route = respx.delete("https://api.wrenn.dev/v1/capsules/cl-1").respond(204) + respx.post("https://api.wrenn.dev/v1/capsules").respond( 201, json={"id": "cl-1", "status": "pending"} ) - sb = client.sandboxes.create() - with sb: - assert sb.id == "cl-1" + cap = client.capsules.create() + with cap: + assert cap.id == "cl-1" assert route.called @@ -147,8 +146,8 @@ class TestCodeResult: class TestJupyterMessageFormat: def test_execute_request_structure(self): - sb = Sandbox(id="test") - msg = sb._jupyter_execute_request("x = 42") + cap = Capsule(id="test") + msg = cap._jupyter_execute_request("x = 42") assert msg["msg_type"] == "execute_request" assert msg["content"]["code"] == "x = 42" assert msg["content"]["silent"] is False @@ -157,7 +156,45 @@ class TestJupyterMessageFormat: assert msg["header"]["msg_type"] == "execute_request" def test_execute_request_unique_ids(self): - sb = Sandbox(id="test") - m1 = sb._jupyter_execute_request("a") - m2 = sb._jupyter_execute_request("b") + cap = Capsule(id="test") + m1 = cap._jupyter_execute_request("a") + m2 = cap._jupyter_execute_request("b") assert m1["msg_id"] != m2["msg_id"] + + +class TestDeprecationWarnings: + def test_import_sandbox_from_capsule_warns(self): + import importlib + import warnings + + import wrenn.capsule as capsule_mod + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + klass = capsule_mod.Sandbox + assert klass is Capsule + assert len(w) == 1 + assert issubclass(w[0].category, DeprecationWarning) + assert "Sandbox" in str(w[0].message) + + def test_import_sandbox_from_wrenn_warns(self): + import warnings + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + from wrenn import Sandbox + + assert Sandbox is Capsule + assert any(issubclass(x.category, DeprecationWarning) for x in w) + + def test_client_sandboxes_property_warns(self): + import warnings + + with WrennClient(api_key="wrn_test1234567890abcdef12345678") as c: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + resource = c.sandboxes + assert resource is c.capsules + assert len(w) == 1 + assert issubclass(w[0].category, DeprecationWarning) + assert "sandboxes" in str(w[0].message) diff --git a/tests/test_client.py b/tests/test_client.py index b9adb02..17c3586 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -9,7 +9,7 @@ from wrenn.exceptions import ( WrennAuthenticationError, WrennConflictError, WrennForbiddenError, - WrennHostHasSandboxesError, + WrennHostHasCapsulesError, WrennInternalError, WrennNotFoundError, WrennValidationError, @@ -17,9 +17,9 @@ from wrenn.exceptions import ( from wrenn.models import ( APIKeyResponse, AuthResponse, + Capsule, CreateHostResponse, Host, - Sandbox, Status, Template, ) @@ -97,10 +97,10 @@ class TestAPIKeys: assert route.called -class TestSandboxes: +class TestCapsules: @respx.mock def test_create(self, client): - respx.post("https://api.wrenn.dev/v1/sandboxes").respond( + respx.post("https://api.wrenn.dev/v1/capsules").respond( 201, json={ "id": "sb-1", @@ -110,40 +110,40 @@ class TestSandboxes: "memory_mb": 1024, }, ) - resp = client.sandboxes.create(template="base-python", vcpus=2, memory_mb=1024) - assert isinstance(resp, Sandbox) + 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("https://api.wrenn.dev/v1/sandboxes").respond( + respx.post("https://api.wrenn.dev/v1/capsules").respond( 201, json={"id": "sb-2", "status": "pending"} ) - resp = client.sandboxes.create() + resp = client.capsules.create() assert resp.id == "sb-2" @respx.mock def test_list(self, client): - respx.get("https://api.wrenn.dev/v1/sandboxes").respond( + respx.get("https://api.wrenn.dev/v1/capsules").respond( 200, json=[{"id": "sb-1", "status": "running"}] ) - boxes = client.sandboxes.list() + boxes = client.capsules.list() assert len(boxes) == 1 assert boxes[0].status == Status.running @respx.mock def test_get(self, client): - respx.get("https://api.wrenn.dev/v1/sandboxes/sb-1").respond( + respx.get("https://api.wrenn.dev/v1/capsules/sb-1").respond( 200, json={"id": "sb-1", "status": "running"} ) - resp = client.sandboxes.get("sb-1") + resp = client.capsules.get("sb-1") assert resp.id == "sb-1" @respx.mock def test_destroy(self, client): - route = respx.delete("https://api.wrenn.dev/v1/sandboxes/sb-1").respond(204) - client.sandboxes.destroy("sb-1") + route = respx.delete("https://api.wrenn.dev/v1/capsules/sb-1").respond(204) + client.capsules.destroy("sb-1") assert route.called @@ -154,7 +154,7 @@ class TestSnapshots: 201, json={"name": "snap-1", "type": "snapshot", "vcpus": 1}, ) - resp = client.snapshots.create(sandbox_id="sb-1", name="snap-1") + resp = client.snapshots.create(capsule_id="sb-1", name="snap-1") assert isinstance(resp, Template) assert resp.name == "snap-1" @@ -163,7 +163,7 @@ class TestSnapshots: route = respx.post("https://api.wrenn.dev/v1/snapshots").respond( 201, json={"name": "snap-1", "type": "snapshot"} ) - client.snapshots.create(sandbox_id="sb-1", overwrite=True) + client.snapshots.create(capsule_id="sb-1", overwrite=True) req = route.calls[0].request assert "overwrite=true" in str(req.url) @@ -262,23 +262,23 @@ class TestHosts: class TestErrorHandling: @respx.mock def test_validation_error(self, client): - respx.post("https://api.wrenn.dev/v1/sandboxes").respond( + respx.post("https://api.wrenn.dev/v1/capsules").respond( 400, json={"error": {"code": "invalid_request", "message": "bad input"}}, ) with pytest.raises(WrennValidationError) as exc_info: - client.sandboxes.create() + 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("https://api.wrenn.dev/v1/sandboxes").respond( + respx.get("https://api.wrenn.dev/v1/capsules").respond( 401, json={"error": {"code": "unauthorized", "message": "bad key"}}, ) with pytest.raises(WrennAuthenticationError): - client.sandboxes.list() + client.capsules.list() @respx.mock def test_forbidden_error(self, client): @@ -291,66 +291,66 @@ class TestErrorHandling: @respx.mock def test_not_found_error(self, client): - respx.get("https://api.wrenn.dev/v1/sandboxes/nope").respond( + respx.get("https://api.wrenn.dev/v1/capsules/nope").respond( 404, - json={"error": {"code": "not_found", "message": "sandbox not found"}}, + json={"error": {"code": "not_found", "message": "capsule not found"}}, ) with pytest.raises(WrennNotFoundError): - client.sandboxes.get("nope") + client.capsules.get("nope") @respx.mock def test_conflict_error(self, client): - respx.get("https://api.wrenn.dev/v1/sandboxes/sb-1").respond( + respx.get("https://api.wrenn.dev/v1/capsules/sb-1").respond( 409, json={"error": {"code": "invalid_state", "message": "not running"}}, ) with pytest.raises(WrennConflictError): - client.sandboxes.get("sb-1") + client.capsules.get("sb-1") @respx.mock - def test_host_has_sandboxes_error(self, client): + def test_host_has_capsules_error(self, client): respx.delete("https://api.wrenn.dev/v1/hosts/h-1").respond( 409, json={ "error": { - "code": "host_has_sandboxes", - "message": "host has running sandboxes", + "code": "host_has_capsules", + "message": "host has running capsules", }, "sandbox_ids": ["sb-1", "sb-2"], }, ) - with pytest.raises(WrennHostHasSandboxesError) as exc_info: + with pytest.raises(WrennHostHasCapsulesError) as exc_info: client.hosts.delete("h-1") - assert exc_info.value.sandbox_ids == ["sb-1", "sb-2"] + assert exc_info.value.capsule_ids == ["sb-1", "sb-2"] @respx.mock def test_agent_error(self, client): - respx.post("https://api.wrenn.dev/v1/sandboxes").respond( + respx.post("https://api.wrenn.dev/v1/capsules").respond( 502, json={"error": {"code": "agent_error", "message": "host agent failed"}}, ) with pytest.raises(WrennAgentError): - client.sandboxes.create() + client.capsules.create() @respx.mock def test_internal_error(self, client): - respx.get("https://api.wrenn.dev/v1/sandboxes/sb-1").respond( + respx.get("https://api.wrenn.dev/v1/capsules/sb-1").respond( 500, json={"error": {"code": "internal_error", "message": "oops"}}, ) with pytest.raises(WrennInternalError): - client.sandboxes.get("sb-1") + client.capsules.get("sb-1") @respx.mock def test_unknown_error_code_falls_back(self, client): - respx.get("https://api.wrenn.dev/v1/sandboxes/sb-1").respond( + respx.get("https://api.wrenn.dev/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.sandboxes.get("sb-1") + client.capsules.get("sb-1") assert exc_info.value.code == "teapot" @@ -379,22 +379,22 @@ class TestAuthModes: class TestAsyncClient: @pytest.mark.asyncio @respx.mock - async def test_async_sandboxes_create(self, async_client): + async def test_async_capsules_create(self, async_client): async with async_client: - respx.post("https://api.wrenn.dev/v1/sandboxes").respond( + respx.post("https://api.wrenn.dev/v1/capsules").respond( 201, json={"id": "sb-1", "status": "pending"} ) - resp = await async_client.sandboxes.create(template="base-python") + resp = await async_client.capsules.create(template="base-python") assert resp.id == "sb-1" @pytest.mark.asyncio @respx.mock - async def test_async_sandboxes_list(self, async_client): + async def test_async_capsules_list(self, async_client): async with async_client: - respx.get("https://api.wrenn.dev/v1/sandboxes").respond( + respx.get("https://api.wrenn.dev/v1/capsules").respond( 200, json=[{"id": "sb-1"}] ) - boxes = await async_client.sandboxes.list() + boxes = await async_client.capsules.list() assert len(boxes) == 1 @pytest.mark.asyncio @@ -409,9 +409,9 @@ class TestAsyncClient: @respx.mock async def test_async_error_handling(self, async_client): async with async_client: - respx.get("https://api.wrenn.dev/v1/sandboxes/nope").respond( + respx.get("https://api.wrenn.dev/v1/capsules/nope").respond( 404, json={"error": {"code": "not_found", "message": "not found"}}, ) with pytest.raises(WrennNotFoundError): - await async_client.sandboxes.get("nope") + await async_client.capsules.get("nope") diff --git a/tests/test_filesystem_pty.py b/tests/test_filesystem_pty.py index 983daa6..6b494a6 100644 --- a/tests/test_filesystem_pty.py +++ b/tests/test_filesystem_pty.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest import respx +from wrenn.capsule import Capsule from wrenn.client import WrennClient from wrenn.models import FileEntry from wrenn.pty import ( @@ -15,7 +16,6 @@ from wrenn.pty import ( PtySession, _parse_pty_event, ) -from wrenn.sandbox import Sandbox @pytest.fixture @@ -24,18 +24,18 @@ def client(): yield c -def _make_sandbox(client: WrennClient, sb_id: str = "cl-abc") -> Sandbox: - respx.post("https://api.wrenn.dev/v1/sandboxes").respond( - 201, json={"id": sb_id, "status": "running"} +def _make_capsule(client: WrennClient, cap_id: str = "cl-abc") -> Capsule: + respx.post("https://api.wrenn.dev/v1/capsules").respond( + 201, json={"id": cap_id, "status": "running"} ) - return client.sandboxes.create() + return client.capsules.create() class TestListDir: @respx.mock def test_list_dir_returns_entries(self, client): - sb = _make_sandbox(client) - respx.post("https://api.wrenn.dev/v1/sandboxes/cl-abc/files/list").respond( + cap = _make_capsule(client) + respx.post("https://api.wrenn.dev/v1/capsules/cl-abc/files/list").respond( 200, json={ "entries": [ @@ -66,7 +66,7 @@ class TestListDir: ] }, ) - entries = sb.list_dir("/home/user") + entries = cap.list_dir("/home/user") assert len(entries) == 2 assert isinstance(entries[0], FileEntry) assert entries[0].name == "main.py" @@ -76,27 +76,27 @@ class TestListDir: @respx.mock def test_list_dir_with_depth(self, client): - sb = _make_sandbox(client) + cap = _make_capsule(client) route = respx.post( - "https://api.wrenn.dev/v1/sandboxes/cl-abc/files/list" + "https://api.wrenn.dev/v1/capsules/cl-abc/files/list" ).respond(200, json={"entries": []}) - sb.list_dir("/home/user", depth=3) + cap.list_dir("/home/user", depth=3) body = json.loads(route.calls[0].request.content) assert body["depth"] == 3 @respx.mock def test_list_dir_empty(self, client): - sb = _make_sandbox(client) - respx.post("https://api.wrenn.dev/v1/sandboxes/cl-abc/files/list").respond( + cap = _make_capsule(client) + respx.post("https://api.wrenn.dev/v1/capsules/cl-abc/files/list").respond( 200, json={"entries": []} ) - entries = sb.list_dir("/empty") + entries = cap.list_dir("/empty") assert entries == [] @respx.mock def test_list_dir_symlink(self, client): - sb = _make_sandbox(client) - respx.post("https://api.wrenn.dev/v1/sandboxes/cl-abc/files/list").respond( + cap = _make_capsule(client) + respx.post("https://api.wrenn.dev/v1/capsules/cl-abc/files/list").respond( 200, json={ "entries": [ @@ -115,7 +115,7 @@ class TestListDir: ] }, ) - entries = sb.list_dir("/home/user") + entries = cap.list_dir("/home/user") assert len(entries) == 1 assert entries[0].type == "symlink" assert entries[0].symlink_target == "/bin" @@ -124,8 +124,8 @@ class TestListDir: class TestMkdir: @respx.mock def test_mkdir_returns_entry(self, client): - sb = _make_sandbox(client) - respx.post("https://api.wrenn.dev/v1/sandboxes/cl-abc/files/mkdir").respond( + cap = _make_capsule(client) + respx.post("https://api.wrenn.dev/v1/capsules/cl-abc/files/mkdir").respond( 200, json={ "entry": { @@ -142,19 +142,19 @@ class TestMkdir: } }, ) - entry = sb.mkdir("/home/user/data") + entry = cap.mkdir("/home/user/data") assert isinstance(entry, FileEntry) assert entry.name == "data" assert entry.type == "directory" @respx.mock def test_mkdir_existing_returns_gracefully(self, client): - sb = _make_sandbox(client) - respx.post("https://api.wrenn.dev/v1/sandboxes/cl-abc/files/mkdir").respond( + cap = _make_capsule(client) + respx.post("https://api.wrenn.dev/v1/capsules/cl-abc/files/mkdir").respond( 409, json={"error": {"code": "conflict", "message": "already exists"}}, ) - respx.post("https://api.wrenn.dev/v1/sandboxes/cl-abc/files/list").respond( + respx.post("https://api.wrenn.dev/v1/capsules/cl-abc/files/list").respond( 200, json={ "entries": [ @@ -173,27 +173,27 @@ class TestMkdir: ] }, ) - entry = sb.mkdir("/home/user/data") + entry = cap.mkdir("/home/user/data") assert entry.name == "data" class TestRemove: @respx.mock def test_remove_succeeds(self, client): - sb = _make_sandbox(client) + cap = _make_capsule(client) route = respx.post( - "https://api.wrenn.dev/v1/sandboxes/cl-abc/files/remove" + "https://api.wrenn.dev/v1/capsules/cl-abc/files/remove" ).respond(204) - sb.remove("/home/user/old_data") + cap.remove("/home/user/old_data") assert route.called @respx.mock def test_remove_sends_path(self, client): - sb = _make_sandbox(client) + cap = _make_capsule(client) route = respx.post( - "https://api.wrenn.dev/v1/sandboxes/cl-abc/files/remove" + "https://api.wrenn.dev/v1/capsules/cl-abc/files/remove" ).respond(204) - sb.remove("/tmp/test.txt") + cap.remove("/tmp/test.txt") body = json.loads(route.calls[0].request.content) assert body["path"] == "/tmp/test.txt" @@ -201,23 +201,23 @@ class TestRemove: class TestUpload: @respx.mock def test_upload_sends_multipart(self, client): - sb = _make_sandbox(client) + cap = _make_capsule(client) route = respx.post( - "https://api.wrenn.dev/v1/sandboxes/cl-abc/files/write" + "https://api.wrenn.dev/v1/capsules/cl-abc/files/write" ).respond(204) - sb.upload("/app/main.py", b"print('hello')") + cap.upload("/app/main.py", b"print('hello')") assert route.called req = route.calls[0].request assert b"multipart/form-data" in req.headers.get("content-type", "").encode() @respx.mock def test_download_returns_bytes(self, client): - sb = _make_sandbox(client) + cap = _make_capsule(client) content = b"file contents here" - respx.post("https://api.wrenn.dev/v1/sandboxes/cl-abc/files/read").respond( + respx.post("https://api.wrenn.dev/v1/capsules/cl-abc/files/read").respond( 200, content=content ) - data = sb.download("/app/main.py") + data = cap.download("/app/main.py") assert data == content @@ -500,7 +500,8 @@ class TestExports: assert APS is not None def test_pty_event_importable(self): - from wrenn import PtyEvent as PE, PtyEventType as PET + from wrenn import PtyEvent as PE + from wrenn import PtyEventType as PET assert PE is not None assert PET is not None diff --git a/tests/test_integration.py b/tests/test_integration.py index ca99b14..9cba1c8 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -64,74 +64,74 @@ def bearer_client() -> Generator[WrennClient, None, None]: @requires_auth -class TestSandboxLifecycle: +class TestCapsuleLifecycle: def test_create_exec_destroy(self, client): - with client.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - result = sb.exec("echo", args=["hello"]) + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - result = sb.exec("echo", args=["hello", "world"]) + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - result = sb.exec("sh", args=["-c", "exit 42"]) + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - result = sb.exec("sh", args=["-c", "echo err>&2"]) + 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): - sb = client.sandboxes.create(template="minimal", timeout_sec=120) - sb_id = sb.id + cap = client.capsules.create(template="minimal", timeout_sec=120) + cap_id = cap.id - with sb: - sb.wait_ready(timeout=60, interval=1) + with cap: + cap.wait_ready(timeout=60, interval=1) - fetched = client.sandboxes.get(sb_id) + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) + with client.capsules.create(template="minimal", timeout_sec=120) as cap: + cap.wait_ready(timeout=60, interval=1) content = b"Hello from integration test!" - sb.upload("/tmp/test_file.txt", content) - downloaded = sb.download("/tmp/test_file.txt") + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) + with client.capsules.create(template="minimal", timeout_sec=120) as cap: + cap.wait_ready(timeout=60, interval=1) with pytest.raises(Exception): - sb.download("/tmp/no_such_file_12345") + cap.download("/tmp/no_such_file_12345") @requires_auth class TestPauseResume: def test_pause_and_resume(self, client): - with client.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - sb.pause() - assert sb.status == "paused" + with client.capsules.create(template="minimal", timeout_sec=120) as cap: + cap.wait_ready(timeout=60, interval=1) + cap.pause() + assert cap.status == "paused" - sb.resume() - sb.wait_ready(timeout=60, interval=1) + cap.resume() + cap.wait_ready(timeout=60, interval=1) - result = sb.exec("echo", args=["resumed"]) + result = cap.exec("echo", args=["resumed"]) assert result.exit_code == 0 assert "resumed" in result.stdout @@ -139,10 +139,10 @@ class TestPauseResume: @requires_auth class TestPing: def test_ping_resets_timer(self, client): - with client.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - sb.ping() - result = sb.exec("echo", args=["still_alive"]) + 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 @@ -150,32 +150,32 @@ class TestPing: @requires_auth class TestProxy: def test_get_url(self, client): - with client.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - url = sb.get_url(8888) - assert sb.id in url + 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_sandboxes(self, client): - with client.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - boxes = client.sandboxes.list() + 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 sb.id in ids + assert cap.id in ids - def test_get_existing_sandbox(self, client): - with client.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - fetched = client.sandboxes.get(sb.id) - assert fetched.id == sb.id + 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_sandbox(self, client): + def test_get_nonexistent_capsule(self, client): with pytest.raises((WrennNotFoundError, WrennValidationError)): - client.sandboxes.get("cl-nonexistent00000000000000000") + client.capsules.get("cl-nonexistent00000000000000000") @requires_auth @@ -204,117 +204,117 @@ class TestAPIKeys: @requires_auth class TestRunCode: def test_basic_execution(self, client): - with client.sandboxes.create( + with client.capsules.create( template="python-interpreter-v0-beta", timeout_sec=120 - ) as sb: - sb.wait_ready(timeout=60, interval=1) + ) as cap: + cap.wait_ready(timeout=60, interval=1) - r = sb.run_code("x = 42") + r = cap.run_code("x = 42") assert r.error is None - r = sb.run_code("x * 2") + r = cap.run_code("x * 2") assert r.text == "84" def test_state_persists(self, client): - with client.sandboxes.create( + with client.capsules.create( template="python-interpreter-v0-beta", timeout_sec=120 - ) as sb: - sb.wait_ready(timeout=60, interval=1) + ) as cap: + cap.wait_ready(timeout=60, interval=1) - sb.run_code("def greet(name): return f'hello {name}'") - r = sb.run_code("greet('sandbox')") - assert "hello sandbox" in (r.text or "") + 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.sandboxes.create( + with client.capsules.create( template="python-interpreter-v0-beta", timeout_sec=120 - ) as sb: - sb.wait_ready(timeout=60, interval=1) + ) as cap: + cap.wait_ready(timeout=60, interval=1) - r = sb.run_code("1/0") + 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.sandboxes.create( + with client.capsules.create( template="python-interpreter-v0-beta", timeout_sec=120 - ) as sb: - sb.wait_ready(timeout=60, interval=1) + ) as cap: + cap.wait_ready(timeout=60, interval=1) - r = sb.run_code("print('hello from kernel')") + r = cap.run_code("print('hello from kernel')") assert "hello from kernel" in r.stdout @requires_auth -class TestAsyncSandboxLifecycle: +class TestAsyncCapsuleLifecycle: @pytest.mark.asyncio async def test_async_create_exec_destroy(self, async_client): async with async_client: - sb = await async_client.sandboxes.create( + cap = await async_client.capsules.create( template="minimal", timeout_sec=120 ) try: - await sb.async_wait_ready(timeout=60, interval=1) - result = await sb.async_exec("echo", args=["async_hello"]) + 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 sb.async_destroy() + await cap.async_destroy() @pytest.mark.asyncio async def test_async_upload_download(self, async_client): async with async_client: - sb = await async_client.sandboxes.create( + cap = await async_client.capsules.create( template="minimal", timeout_sec=120 ) try: - await sb.async_wait_ready(timeout=60, interval=1) + await cap.async_wait_ready(timeout=60, interval=1) content = b"Async upload test" - await sb.async_upload("/tmp/async_test.txt", content) - downloaded = await sb.async_download("/tmp/async_test.txt") + await cap.async_upload("/tmp/async_test.txt", content) + downloaded = await cap.async_download("/tmp/async_test.txt") assert downloaded == content finally: - await sb.async_destroy() + await cap.async_destroy() @pytest.mark.asyncio async def test_async_run_code(self, async_client): async with async_client: - sb = await async_client.sandboxes.create( + cap = await async_client.capsules.create( template="python-interpreter-v0-beta", timeout_sec=120 ) try: - await sb.async_wait_ready(timeout=60, interval=1) - r = await sb.async_run_code("42 * 2") + await cap.async_wait_ready(timeout=60, interval=1) + r = await cap.async_run_code("42 * 2") assert r.text == "84" finally: - await sb.async_destroy() + await cap.async_destroy() @requires_auth class TestFilesystemListDir: def test_list_dir_root(self, client: WrennClient): - with client.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - sb.mkdir("/tmp/ls_test_root") - sb.upload("/tmp/ls_test_root/hello.txt", b"hello") - entries = sb.list_dir("/tmp/ls_test_root") + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - sb.mkdir("/tmp/fs_test_dir") - entries = sb.list_dir("/tmp") + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - sb.upload("/tmp/meta_test.txt", b"hello world") - entries = sb.list_dir("/tmp") + 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] @@ -326,100 +326,100 @@ class TestFilesystemListDir: assert f.modified_at is not None def test_list_dir_depth(self, client): - with client.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - sb.mkdir("/tmp/depth_a/depth_b") - sb.upload("/tmp/depth_a/depth_b/nested.txt", b"deep") - entries = sb.list_dir("/tmp/depth_a", depth=2) + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - sb.mkdir("/tmp/empty_dir_test") - entries = sb.list_dir("/tmp/empty_dir_test") + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - entry = sb.mkdir("/tmp/mkdir_test") + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - entry = sb.mkdir("/tmp/a/b/c/d") + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - sb.mkdir("/tmp/exist_test") - entry = sb.mkdir("/tmp/exist_test") + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - sb.upload("/tmp/rm_test.txt", b"delete me") - entries_before = sb.list_dir("/tmp") + 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) - sb.remove("/tmp/rm_test.txt") - entries_after = sb.list_dir("/tmp") + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - sb.mkdir("/tmp/rm_dir_test") - sb.upload("/tmp/rm_dir_test/file.txt", b"inside") - sb.remove("/tmp/rm_dir_test") - entries = sb.list_dir("/tmp") + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) + 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 - sb.upload("/tmp/rt.txt", content) - downloaded = sb.download("/tmp/rt.txt") + cap.upload("/tmp/rt.txt", content) + downloaded = cap.download("/tmp/rt.txt") assert downloaded == content - sb.remove("/tmp/rt.txt") + cap.remove("/tmp/rt.txt") with pytest.raises(Exception): - sb.download("/tmp/rt.txt") + cap.download("/tmp/rt.txt") @requires_auth class TestStreamUploadDownload: def test_stream_upload_and_download(self, client: WrennClient): - with client.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) + 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 - sb.stream_upload("/tmp/stream_test.bin", data_gen()) - downloaded = sb.download("/tmp/stream_test.bin") + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) + with client.capsules.create(template="minimal", timeout_sec=120) as cap: + cap.wait_ready(timeout=60, interval=1) content = b"x" * 65536 * 3 - sb.upload("/tmp/large.bin", content) + cap.upload("/tmp/large.bin", content) collected = b"" - for chunk in sb.stream_download("/tmp/large.bin"): + for chunk in cap.stream_download("/tmp/large.bin"): collected += chunk assert collected == content @@ -427,9 +427,9 @@ class TestStreamUploadDownload: @requires_auth class TestPty: def test_pty_basic_output(self, client): - with client.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - with sb.pty(cmd="/bin/sh", cwd="/tmp") as term: + 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: @@ -442,9 +442,9 @@ class TestPty: assert b"pty_hello" in output def test_pty_tag_and_pid(self, client): - with client.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - with sb.pty(cmd="/bin/sh") as term: + 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: @@ -459,18 +459,18 @@ class TestPty: assert started def test_pty_exit_on_command_exit(self, client): - with client.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - with sb.pty(cmd="/bin/echo", args=["immediate"]) as term: + 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.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - with sb.pty(cmd="/bin/sh", cols=80, rows=24) as term: + 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) @@ -479,9 +479,9 @@ class TestPty: break def test_pty_envs(self, client): - with client.sandboxes.create(template="minimal", timeout_sec=120) as sb: - sb.wait_ready(timeout=60, interval=1) - with sb.pty(cmd="/bin/sh", envs={"MY_VAR": "hello_env"}) as term: + 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: @@ -500,69 +500,69 @@ class TestAsyncFilesystem: @pytest.mark.asyncio async def test_async_list_dir(self, async_client): async with async_client: - sb = await async_client.sandboxes.create( + cap = await async_client.capsules.create( template="minimal", timeout_sec=120 ) try: - await sb.async_wait_ready(timeout=60, interval=1) - await sb.async_mkdir("/tmp/async_ls_test") - await sb.async_upload("/tmp/async_ls_test/file.txt", b"data") - entries = await sb.async_list_dir("/tmp/async_ls_test") + 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 sb.async_destroy() + await cap.async_destroy() @pytest.mark.asyncio async def test_async_mkdir(self, async_client): async with async_client: - sb = await async_client.sandboxes.create( + cap = await async_client.capsules.create( template="minimal", timeout_sec=120 ) try: - await sb.async_wait_ready(timeout=60, interval=1) - entry = await sb.async_mkdir("/tmp/async_mkdir_test") + 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 sb.async_destroy() + await cap.async_destroy() @pytest.mark.asyncio async def test_async_remove(self, async_client): async with async_client: - sb = await async_client.sandboxes.create( + cap = await async_client.capsules.create( template="minimal", timeout_sec=120 ) try: - await sb.async_wait_ready(timeout=60, interval=1) - await sb.async_upload("/tmp/async_rm.txt", b"bye") - entries = await sb.async_list_dir("/tmp") + 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 sb.async_remove("/tmp/async_rm.txt") - entries = await sb.async_list_dir("/tmp") + 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 sb.async_destroy() + await cap.async_destroy() @pytest.mark.asyncio async def test_async_full_filesystem_roundtrip(self, async_client): async with async_client: - sb = await async_client.sandboxes.create( + cap = await async_client.capsules.create( template="minimal", timeout_sec=120 ) try: - await sb.async_wait_ready(timeout=60, interval=1) + await cap.async_wait_ready(timeout=60, interval=1) - await sb.async_mkdir("/tmp/async_rt") - await sb.async_upload("/tmp/async_rt/file.txt", b"async content") - entries = await sb.async_list_dir("/tmp/async_rt") + 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 sb.async_download("/tmp/async_rt/file.txt") + data = await cap.async_download("/tmp/async_rt/file.txt") assert data == b"async content" - await sb.async_remove("/tmp/async_rt/file.txt") - entries = await sb.async_list_dir("/tmp/async_rt") + 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 sb.async_destroy() + await cap.async_destroy()