From 51c698751579e75df33fd0278539499efcfd7617 Mon Sep 17 00:00:00 2001 From: pptx704 Date: Tue, 19 May 2026 15:06:49 +0600 Subject: [PATCH] fix: sync SDK with v0.2 API, add wait kwargs to lifecycle ops - Drop AuthResponse from models __init__ (renamed SessionResponse server-side; SDK auths via API key, doesn't need either) - Regenerate models from updated 0.2 openapi spec - Add wait: bool = False kwarg to Capsule/AsyncCapsule destroy/pause/resume (instance + _static_*); 500ms poll for resume/destroy, 2s for pause - Unify polling into _poll_until / _apoll_until + _wait_for_status helper; remove duplicated _POLL_INTERVALS tables - wait_ready: drop implicit paused->resume side effect; treat missing as fail - Capsule.connect: handle transient pausing (wait for paused first) before resuming, fixes hang when caller pauses then connects immediately - Drop dead "if self._id is None" branch in Capsule.__init__ after assigning from already-truthy _capsule_id - files.make_dir: detect already_exists across 409/wrapped error messages via shared _is_already_exists helper - tests/test_integration.py: assertions on final lifecycle state use wait=True --- api/openapi.yaml | 1273 +++++++++++++++++++++++++++++--- src/wrenn/async_capsule.py | 142 ++-- src/wrenn/capsule.py | 147 ++-- src/wrenn/files.py | 62 +- src/wrenn/models/__init__.py | 2 - src/wrenn/models/_generated.py | 162 +++- tests/test_integration.py | 10 +- 7 files changed, 1551 insertions(+), 247 deletions(-) diff --git a/api/openapi.yaml b/api/openapi.yaml index dfc5c75..f3fb110 100644 --- a/api/openapi.yaml +++ b/api/openapi.yaml @@ -53,7 +53,7 @@ paths: tags: [auth] description: | Consumes the activation token sent via email and activates the user account. - Creates a default team and returns a JWT to log the user in. + Creates a default team and sets a session cookie to log the user in. requestBody: required: true content: @@ -66,11 +66,11 @@ paths: type: string responses: "200": - description: Account activated, JWT issued + description: Account activated, session cookie set content: application/json: schema: - $ref: "#/components/schemas/AuthResponse" + $ref: "#/components/schemas/SessionResponse" "400": description: Invalid or expired token content: @@ -78,17 +78,113 @@ paths: schema: $ref: "#/components/schemas/Error" + /v1/auth/logout: + post: + summary: Revoke the current session + operationId: logout + tags: [auth] + security: + - sessionAuth: [] + responses: + "204": + description: Session revoked; cookies cleared + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + + /v1/auth/logout-all: + post: + summary: Revoke every session for the current user + operationId: logoutAll + tags: [auth] + description: | + Revokes every active session for the calling user across all devices, + including the caller's own. Returns 204 and clears cookies on the + response. Triggered automatically by password change, password add, + and password reset. + security: + - sessionAuth: [] + responses: + "204": + description: All sessions revoked + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + + /v1/me/sessions: + get: + summary: List the caller's active sessions + operationId: listSessions + tags: [me] + security: + - sessionAuth: [] + responses: + "200": + description: Sessions list + content: + application/json: + schema: + type: object + properties: + sessions: + type: array + items: + type: object + properties: + id: + type: string + user_agent: + type: string + ip_address: + type: string + created_at: + type: string + format: date-time + last_seen_at: + type: string + format: date-time + expires_at: + type: string + format: date-time + current: + type: boolean + "401": + $ref: "#/components/responses/Unauthorized" + + /v1/me/sessions/{id}: + delete: + summary: Revoke a single session + operationId: revokeSession + tags: [me] + security: + - sessionAuth: [] + parameters: + - name: id + in: path + required: true + schema: + type: string + responses: + "204": + description: Session revoked + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + /v1/auth/switch-team: post: summary: Switch active team operationId: switchTeam tags: [auth] security: - - bearerAuth: [] + - sessionAuth: [] description: | - Re-issues a JWT scoped to a different team. The user must be a member of - the target team (verified from DB). Use the returned token for subsequent - requests to that team's resources. + Rotates the session SID and updates its team scope. The user must be a + member of the target team (verified from DB). The new wrenn_sid and + wrenn_csrf cookies are set on the response. requestBody: required: true content: @@ -101,11 +197,11 @@ paths: type: string responses: "200": - description: New JWT issued for the target team + description: New session issued for the target team; cookies refreshed content: application/json: schema: - $ref: "#/components/schemas/AuthResponse" + $ref: "#/components/schemas/SessionResponse" "403": description: Not a member of this team content: @@ -136,7 +232,7 @@ paths: content: application/json: schema: - $ref: "#/components/schemas/AuthResponse" + $ref: "#/components/schemas/SessionResponse" "401": description: Invalid credentials content: @@ -144,7 +240,7 @@ paths: schema: $ref: "#/components/schemas/Error" - /v1/auth/oauth/{provider}: + /auth/oauth/{provider}: parameters: - name: provider in: path @@ -171,7 +267,7 @@ paths: schema: $ref: "#/components/schemas/Error" - /v1/auth/oauth/{provider}/callback: + /auth/oauth/{provider}/callback: parameters: - name: provider in: path @@ -188,9 +284,10 @@ paths: description: | Handles the OAuth provider's callback after user authorization. Exchanges the authorization code for a user profile, creates or - logs in the user, and redirects to the frontend with a JWT token. + logs in the user, sets the wrenn_sid + wrenn_csrf cookies, and + redirects to the SPA callback page. - **On success:** redirects to `{OAUTH_REDIRECT_URL}/auth/{provider}/callback?token=...&user_id=...&team_id=...&email=...` + **On success:** redirects to `{OAUTH_REDIRECT_URL}/auth/{provider}/callback` (no tokens in URL). **On error:** redirects to `{OAUTH_REDIRECT_URL}/auth/{provider}/callback?error=...` @@ -217,7 +314,7 @@ paths: operationId: getMe tags: [account] security: - - bearerAuth: [] + - sessionAuth: [] responses: "200": description: User profile @@ -231,7 +328,7 @@ paths: operationId: updateName tags: [account] security: - - bearerAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -245,12 +342,8 @@ paths: minLength: 1 maxLength: 100 responses: - "200": - description: Name updated, new JWT issued - content: - application/json: - schema: - $ref: "#/components/schemas/AuthResponse" + "204": + description: Name updated; session caches refreshed "400": description: Invalid name content: @@ -263,7 +356,7 @@ paths: operationId: deleteAccount tags: [account] security: - - bearerAuth: [] + - sessionAuth: [] description: | Soft-deletes the account (sets status=deleted, deleted_at=now). The account is permanently removed after 15 days. Blocked if the user @@ -301,7 +394,7 @@ paths: operationId: changePassword tags: [account] security: - - bearerAuth: [] + - sessionAuth: [] description: | For users with an existing password: requires `current_password` and `new_password`. For OAuth-only users adding a password: requires `new_password` and `confirm_password`. @@ -398,7 +491,7 @@ paths: operationId: connectProvider tags: [account] security: - - bearerAuth: [] + - sessionAuth: [] description: | Sets OAuth state and link cookies, then returns the provider's authorization URL. The frontend navigates to this URL to start the @@ -437,7 +530,7 @@ paths: operationId: disconnectProvider tags: [account] security: - - bearerAuth: [] + - sessionAuth: [] description: | Unlinks the OAuth provider from the current account. Blocked if this is the user's only login method (no password and no other providers). @@ -463,7 +556,7 @@ paths: operationId: createAPIKey tags: [api-keys] security: - - bearerAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -489,7 +582,7 @@ paths: operationId: listAPIKeys tags: [api-keys] security: - - bearerAuth: [] + - sessionAuth: [] responses: "200": description: List of API keys (plaintext keys are never returned) @@ -513,7 +606,7 @@ paths: operationId: deleteAPIKey tags: [api-keys] security: - - bearerAuth: [] + - sessionAuth: [] responses: "204": description: API key deleted @@ -524,7 +617,7 @@ paths: operationId: searchUsers tags: [users] security: - - bearerAuth: [] + - sessionAuth: [] description: | Returns up to 10 users whose email starts with the given prefix. The prefix must contain "@". Intended for the add-member UI autocomplete. @@ -557,7 +650,7 @@ paths: operationId: listTeams tags: [teams] security: - - bearerAuth: [] + - sessionAuth: [] responses: "200": description: Teams the user belongs to, each with their role @@ -573,7 +666,7 @@ paths: operationId: createTeam tags: [teams] security: - - bearerAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -613,7 +706,7 @@ paths: operationId: getTeam tags: [teams] security: - - bearerAuth: [] + - sessionAuth: [] responses: "200": description: Team details with members @@ -639,7 +732,7 @@ paths: operationId: renameTeam tags: [teams] security: - - bearerAuth: [] + - sessionAuth: [] description: Admin or owner role required (verified from DB). requestBody: required: true @@ -672,10 +765,10 @@ paths: operationId: deleteTeam tags: [teams] security: - - bearerAuth: [] + - sessionAuth: [] description: | Owner only. Soft-deletes the team and destroys all running/paused/starting - capsulees. All DB records are preserved. The team slug is permanently reserved. + capsules. All DB records are preserved. The team slug is permanently reserved. responses: "204": description: Team deleted @@ -699,7 +792,7 @@ paths: operationId: listTeamMembers tags: [teams] security: - - bearerAuth: [] + - sessionAuth: [] responses: "200": description: Members with roles @@ -715,7 +808,7 @@ paths: operationId: addTeamMember tags: [teams] security: - - bearerAuth: [] + - sessionAuth: [] description: Admin or owner role required. User is added instantly as a member. requestBody: required: true @@ -773,7 +866,7 @@ paths: operationId: updateMemberRole tags: [teams] security: - - bearerAuth: [] + - sessionAuth: [] description: | Admin or owner required. Valid target roles: admin, member. The owner's role cannot be changed. @@ -809,7 +902,7 @@ paths: operationId: removeTeamMember tags: [teams] security: - - bearerAuth: [] + - sessionAuth: [] description: Admin or owner required. Owner cannot be removed. responses: "204": @@ -840,7 +933,7 @@ paths: operationId: leaveTeam tags: [teams] security: - - bearerAuth: [] + - sessionAuth: [] description: The owner cannot leave; they must delete the team instead. responses: "204": @@ -859,6 +952,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -880,14 +974,15 @@ paths: $ref: "#/components/schemas/Error" get: - summary: List capsulees for your team + summary: List capsules for your team operationId: listCapsules tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] responses: "200": - description: List of capsulees + description: List of capsules content: application/json: schema: @@ -902,6 +997,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] parameters: - name: range in: query @@ -928,6 +1024,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] parameters: - name: from in: query @@ -967,6 +1064,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] responses: "200": description: Capsule details @@ -987,6 +1085,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] responses: "202": description: Capsule destruction initiated @@ -1005,6 +1104,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -1051,6 +1151,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] description: | Returns all running processes inside the capsule, including background processes and any processes started by templates or init scripts. @@ -1094,6 +1195,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] parameters: - name: signal in: query @@ -1139,6 +1241,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] description: | Opens a WebSocket connection to stream stdout/stderr from a running background process. The selector can be a numeric PID or a string tag. @@ -1167,9 +1270,10 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] description: | Resets the last_active_at timestamp for a running capsule, preventing - the auto-pause TTL from expiring. Use this as a keepalive for capsulees + the auto-pause TTL from expiring. Use this as a keepalive for capsules that are idle but should remain running. responses: "204": @@ -1201,7 +1305,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] - - bearerAuth: [] + - sessionAuth: [] description: | Returns time-series CPU, memory, and disk metrics for a capsule. Three tiers are available with different granularity and retention: @@ -1209,9 +1313,9 @@ paths: - `2h`: 30-second averages, last 2 hours - `24h`: 5-minute averages, last 24 hours - 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. + For running capsules, data comes from the host agent's in-memory + ring buffer. For paused capsules, data is read from persisted + snapshots in the database. Stopped/destroyed capsules return 404. parameters: - name: range in: query @@ -1255,6 +1359,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] description: | Takes a snapshot of the capsule (VM state + memory + rootfs), then destroys all running resources. The capsule exists only as files on @@ -1287,10 +1392,12 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] description: | - 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. + Restores a paused capsule from its snapshot. Cloud Hypervisor is + relaunched in --restore mode with memory_restore_mode=ondemand so + guest pages fault in lazily via userfaultfd. The original network + slot (and host-reachable IP) is preserved across pause/resume. responses: "202": description: Capsule resume initiated (status will be "resuming") @@ -1312,18 +1419,15 @@ paths: tags: [snapshots] security: - apiKeyAuth: [] + - sessionAuth: [] description: | - Pauses a running capsule, takes a full snapshot, copies the snapshot - files to the images directory as a reusable template, then destroys - the capsule. The template can be used to create new capsulees. - parameters: - - name: overwrite - in: query - required: false - schema: - type: string - enum: ["true"] - description: Set to "true" to overwrite an existing snapshot with the same name. + Live snapshot: briefly pauses the capsule, writes its VM state + + memory + flattened rootfs to a new template directory, then resumes + the capsule. The source capsule keeps running after the snapshot; + the resulting template can be used to create new capsules. + + Snapshots are immutable: each call must use a fresh name. Re-using + an existing name returns 409 Conflict. requestBody: required: true content: @@ -1350,6 +1454,7 @@ paths: tags: [snapshots] security: - apiKeyAuth: [] + - sessionAuth: [] parameters: - name: type in: query @@ -1382,6 +1487,7 @@ paths: tags: [snapshots] security: - apiKeyAuth: [] + - sessionAuth: [] description: Removes the snapshot files from disk and deletes the database record. responses: "204": @@ -1407,6 +1513,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -1452,6 +1559,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -1487,6 +1595,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -1527,6 +1636,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -1547,7 +1657,9 @@ paths: schema: $ref: "#/components/schemas/Error" "409": - description: Capsule not running + description: > + Capsule not running, or a directory already exists at the + target path (error code `already_exists`). content: application/json: schema: @@ -1567,6 +1679,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -1603,6 +1716,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] description: | Opens a WebSocket connection for streaming command execution. @@ -1656,6 +1770,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] description: | Opens a WebSocket connection for an interactive PTY (terminal) session. Supports creating new sessions, sending input, resizing, killing, and @@ -1733,6 +1848,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] description: | Streams file content to the capsule without buffering in memory. Suitable for large files. Uses the same multipart/form-data format @@ -1782,6 +1898,7 @@ paths: tags: [capsules] security: - apiKeyAuth: [] + - sessionAuth: [] description: | Streams file content from the capsule without buffering in memory. Suitable for large files. Returns raw bytes with chunked transfer encoding. @@ -1818,7 +1935,7 @@ paths: operationId: createHost tags: [hosts] security: - - bearerAuth: [] + - sessionAuth: [] description: | Creates a new host record and returns a one-time registration token. Regular hosts can only be created by admins. BYOC hosts can be created @@ -1854,7 +1971,7 @@ paths: operationId: listHosts tags: [hosts] security: - - bearerAuth: [] + - sessionAuth: [] description: | Admins see all hosts. Non-admins see only BYOC hosts belonging to their team. responses: @@ -1880,7 +1997,7 @@ paths: operationId: getHost tags: [hosts] security: - - bearerAuth: [] + - sessionAuth: [] responses: "200": description: Host details @@ -1900,18 +2017,18 @@ paths: operationId: deleteHost tags: [hosts] security: - - bearerAuth: [] + - sessionAuth: [] 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 capsulees. With `?force=true`, destroys all capsulees first. + has active capsules. With `?force=true`, destroys all capsules first. parameters: - name: force in: query required: false schema: type: boolean - description: If true, destroy all capsulees on the host before deleting. + description: If true, destroy all capsules on the host before deleting. responses: "204": description: Host deleted @@ -1922,7 +2039,7 @@ paths: schema: $ref: "#/components/schemas/Error" "409": - description: Host has active capsulees (only when force is not set) + description: Host has active capsules (only when force is not set) content: application/json: schema: @@ -1941,7 +2058,7 @@ paths: operationId: regenerateHostToken tags: [hosts] security: - - bearerAuth: [] + - sessionAuth: [] description: | Issues a new registration token for a host still in "pending" status. Use this when a previous registration attempt failed after consuming @@ -2056,7 +2173,13 @@ paths: properties: event: type: string - enum: [sandbox.auto_paused] + description: | + Lifecycle event type. Known values: + * `sandbox.auto_paused` — TTL reaper paused the capsule + * `sandbox.stopped` — autonomous destroy (crash/eviction) + * `sandbox.error` — VMM/crash watcher reported error + Unknown event names are accepted and forwarded to the + stream consumer as-is (future-compatible). sandbox_id: type: string host_id: @@ -2122,7 +2245,7 @@ paths: operationId: getHostDeletePreview tags: [hosts] security: - - bearerAuth: [] + - sessionAuth: [] description: | Returns the list of capsule IDs that would be destroyed if the host were deleted with `?force=true`. No state is modified. @@ -2159,7 +2282,7 @@ paths: operationId: listHostTags tags: [hosts] security: - - bearerAuth: [] + - sessionAuth: [] responses: "200": description: List of tags @@ -2175,7 +2298,7 @@ paths: operationId: addHostTag tags: [hosts] security: - - bearerAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -2210,7 +2333,7 @@ paths: operationId: removeHostTag tags: [hosts] security: - - bearerAuth: [] + - sessionAuth: [] responses: "204": description: Tag removed @@ -2227,7 +2350,7 @@ paths: operationId: createChannel tags: [channels] security: - - bearerAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -2248,7 +2371,7 @@ paths: operationId: listChannels tags: [channels] security: - - bearerAuth: [] + - sessionAuth: [] responses: "200": description: Channels list @@ -2268,7 +2391,7 @@ paths: operationId: testChannel tags: [channels] security: - - bearerAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -2301,7 +2424,7 @@ paths: operationId: getChannel tags: [channels] security: - - bearerAuth: [] + - sessionAuth: [] responses: "200": description: Channel details @@ -2320,7 +2443,7 @@ paths: operationId: updateChannel tags: [channels] security: - - bearerAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -2347,7 +2470,7 @@ paths: operationId: deleteChannel tags: [channels] security: - - bearerAuth: [] + - sessionAuth: [] responses: "204": description: Channel deleted @@ -2368,7 +2491,7 @@ paths: operationId: rotateChannelConfig tags: [channels] security: - - bearerAuth: [] + - sessionAuth: [] requestBody: required: true content: @@ -2398,11 +2521,11 @@ paths: tags: [admin] description: | Sets the platform admin flag on a user. Cannot remove the last admin. - Requires platform admin access (JWT + is_admin). - The target user's JWT is not re-issued — their frontend will reflect the - change on next login or team switch. + Requires platform admin access. Session caches for the target user + are invalidated immediately so the flag flip takes effect on the + user's next request. security: - - bearerAuth: [] + - sessionAuth: [] parameters: - name: id in: path @@ -2439,7 +2562,811 @@ paths: schema: $ref: "#/components/schemas/Error" + /v1/events/stream: + get: + summary: Real-time lifecycle event stream + operationId: streamEvents + tags: [events] + description: | + Server-Sent Events stream of capsule, template, and host lifecycle + events scoped to the caller's active team. Browsers send the + wrenn_sid cookie automatically on EventSource connections; SDKs + authenticate via X-API-Key. + + Frame format follows the standard SSE protocol: + ``` + event: capsule.create + data: {"event":"capsule.create","outcome":"success","resource":{"id":"sb-..."},"sandbox":{...},"timestamp":"2026-05-19T02:00:00Z"} + + : keepalive + ``` + A `: keepalive` comment is emitted every 30s. + security: + - apiKeyAuth: [] + - sessionAuth: [] + responses: + "200": + description: SSE stream opened + content: + text/event-stream: + schema: + $ref: "#/components/schemas/SSEEvent" + "401": + description: Missing or invalid auth + content: + application/json: + schema: + $ref: "#/components/schemas/Error" + + /v1/audit-logs: + get: + summary: List team audit log entries + operationId: listAuditLogs + tags: [audit] + description: Paginated cursor list of audit events for the caller's team. + security: + - sessionAuth: [] + parameters: + - name: before + in: query + required: false + schema: + type: string + format: date-time + - name: before_id + in: query + required: false + schema: + type: string + - name: limit + in: query + required: false + schema: + type: integer + minimum: 1 + maximum: 200 + default: 50 + responses: + "200": + description: Audit log page + content: + application/json: + schema: + type: object + properties: + entries: + type: array + items: + $ref: "#/components/schemas/AuditLogEntry" + next_cursor: + type: object + nullable: true + properties: + before: + type: string + format: date-time + before_id: + type: string + + /v1/admin/events/stream: + get: + summary: Admin SSE event stream (all teams) + operationId: adminStreamEvents + tags: [admin, events] + description: | + Admin variant of /v1/events/stream that emits events across all teams. + Requires an admin session cookie. + security: + - sessionAuth: [] + responses: + "200": + description: SSE stream opened + content: + text/event-stream: + schema: + $ref: "#/components/schemas/SSEEvent" + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + + /v1/admin/audit-logs: + get: + summary: List audit log entries (all teams) + operationId: adminListAuditLogs + tags: [admin, audit] + security: + - sessionAuth: [] + parameters: + - name: before + in: query + schema: {type: string, format: date-time} + - name: before_id + in: query + schema: {type: string} + - name: limit + in: query + schema: {type: integer, minimum: 1, maximum: 200, default: 50} + responses: + "200": + description: Audit log page (all teams) + content: + application/json: + schema: + type: object + properties: + entries: + type: array + items: + $ref: "#/components/schemas/AuditLogEntry" + + /v1/admin/teams: + get: + summary: List all teams (admin) + operationId: adminListTeams + tags: [admin] + security: + - sessionAuth: [] + responses: + "200": + description: Teams list + content: + application/json: + schema: + type: array + items: {type: object} + + /v1/admin/teams/{id}/byoc: + put: + summary: Toggle BYOC for a team (admin) + operationId: adminSetTeamBYOC + tags: [admin] + security: + - sessionAuth: [] + parameters: + - name: id + in: path + required: true + schema: {type: string} + requestBody: + required: true + content: + application/json: + schema: + type: object + required: [byoc] + properties: + byoc: {type: boolean} + responses: + "204": + description: Updated + + /v1/admin/teams/{id}: + delete: + summary: Delete a team (admin) + operationId: adminDeleteTeam + tags: [admin] + security: + - sessionAuth: [] + parameters: + - name: id + in: path + required: true + schema: {type: string} + responses: + "204": + description: Deleted + + /v1/admin/users: + get: + summary: List all users (admin) + operationId: adminListUsers + tags: [admin] + security: + - sessionAuth: [] + responses: + "200": + description: Users list + content: + application/json: + schema: + type: array + items: {type: object} + + /v1/admin/users/{id}/active: + put: + summary: Activate or deactivate a user (admin) + operationId: adminSetUserActive + tags: [admin] + security: + - sessionAuth: [] + parameters: + - name: id + in: path + required: true + schema: {type: string} + requestBody: + required: true + content: + application/json: + schema: + type: object + required: [active] + properties: + active: {type: boolean} + responses: + "204": + description: Updated + + /v1/admin/templates: + get: + summary: List all templates (admin) + operationId: adminListTemplates + tags: [admin] + security: + - sessionAuth: [] + responses: + "200": + description: Templates list + content: + application/json: + schema: + type: array + items: + $ref: "#/components/schemas/Template" + + /v1/admin/templates/{name}: + delete: + summary: Delete a template (admin) + operationId: adminDeleteTemplate + tags: [admin] + security: + - sessionAuth: [] + parameters: + - name: name + in: path + required: true + schema: {type: string} + responses: + "204": + description: Deleted + + /v1/admin/builds: + post: + summary: Submit a template build (admin) + operationId: adminCreateBuild + tags: [admin] + security: + - sessionAuth: [] + requestBody: + required: true + content: + application/json: + schema: {type: object} + responses: + "202": + description: Build queued + content: + application/json: + schema: {type: object} + get: + summary: List builds (admin) + operationId: adminListBuilds + tags: [admin] + security: + - sessionAuth: [] + responses: + "200": + description: Builds list + content: + application/json: + schema: + type: array + items: {type: object} + + /v1/admin/builds/{id}: + get: + summary: Get build detail (admin) + operationId: adminGetBuild + tags: [admin] + security: + - sessionAuth: [] + parameters: + - name: id + in: path + required: true + schema: {type: string} + responses: + "200": + description: Build detail + content: + application/json: + schema: {type: object} + + /v1/admin/builds/{id}/cancel: + post: + summary: Cancel a build (admin) + operationId: adminCancelBuild + tags: [admin] + security: + - sessionAuth: [] + parameters: + - name: id + in: path + required: true + schema: {type: string} + responses: + "204": + description: Cancelled + + /v1/admin/capsules: + post: + summary: Create a capsule on behalf of any team (admin) + operationId: adminCreateCapsule + tags: [admin] + security: + - sessionAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/CreateCapsuleRequest" + responses: + "201": + description: Capsule created + content: + application/json: + schema: + $ref: "#/components/schemas/Capsule" + get: + summary: List capsules across all teams (admin) + operationId: adminListCapsules + tags: [admin] + security: + - sessionAuth: [] + responses: + "200": + description: Capsules list + content: + application/json: + schema: + type: array + items: + $ref: "#/components/schemas/Capsule" + + /v1/admin/capsules/{id}: + parameters: + - name: id + in: path + required: true + schema: {type: string} + get: + summary: Get capsule detail (admin) + operationId: adminGetCapsule + tags: [admin] + security: + - sessionAuth: [] + responses: + "200": + description: Capsule detail + content: + application/json: + schema: + $ref: "#/components/schemas/Capsule" + delete: + summary: Destroy capsule (admin) + operationId: adminDestroyCapsule + tags: [admin] + security: + - sessionAuth: [] + responses: + "204": + description: Destroyed + + /v1/admin/capsules/{id}/snapshot: + post: + summary: Create snapshot from any capsule (admin) + operationId: adminCreateSnapshotFromCapsule + tags: [admin] + security: + - sessionAuth: [] + parameters: + - name: id + in: path + required: true + schema: {type: string} + requestBody: + required: true + content: + application/json: + schema: + type: object + required: [name] + properties: + name: {type: string} + responses: + "201": + description: Snapshot created + content: + application/json: + schema: + $ref: "#/components/schemas/Template" + + /v1/admin/capsules/{id}/exec: + parameters: + - name: id + in: path + required: true + schema: {type: string} + post: + summary: Execute a command on any capsule (admin) + operationId: adminExecCommand + tags: [admin] + security: + - sessionAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/ExecRequest" + responses: + "200": + description: Command output (foreground exec) + content: + application/json: + schema: + $ref: "#/components/schemas/ExecResponse" + "202": + description: Background process started + content: + application/json: + schema: + $ref: "#/components/schemas/BackgroundExecResponse" + "404": + $ref: "#/components/responses/NotFound" + "409": + $ref: "#/components/responses/FailedPrecondition" + + /v1/admin/capsules/{id}/metrics: + parameters: + - name: id + in: path + required: true + schema: {type: string} + get: + summary: Get per-capsule resource metrics (admin) + operationId: adminGetCapsuleMetrics + tags: [admin] + security: + - sessionAuth: [] + parameters: + - name: range + in: query + required: false + schema: + type: string + enum: ["5m", "10m", "1h", "2h", "6h", "12h", "24h"] + default: "10m" + responses: + "200": + description: Metrics retrieved + content: + application/json: + schema: + $ref: "#/components/schemas/CapsuleMetrics" + "404": + $ref: "#/components/responses/NotFound" + + /v1/admin/capsules/{id}/processes: + parameters: + - name: id + in: path + required: true + schema: {type: string} + get: + summary: List running processes on any capsule (admin) + operationId: adminListProcesses + tags: [admin] + security: + - sessionAuth: [] + responses: + "200": + description: Process list + content: + application/json: + schema: + $ref: "#/components/schemas/ProcessListResponse" + "404": + $ref: "#/components/responses/NotFound" + "409": + $ref: "#/components/responses/FailedPrecondition" + + /v1/admin/capsules/{id}/processes/{selector}: + parameters: + - name: id + in: path + required: true + schema: {type: string} + - name: selector + in: path + required: true + schema: {type: string} + description: Process PID (numeric) or tag (string) + delete: + summary: Kill a process on any capsule (admin) + operationId: adminKillProcess + tags: [admin] + security: + - sessionAuth: [] + parameters: + - name: signal + in: query + required: false + schema: + type: string + enum: [SIGKILL, SIGTERM] + default: SIGKILL + responses: + "204": + description: Process killed + "404": + $ref: "#/components/responses/NotFound" + "409": + $ref: "#/components/responses/FailedPrecondition" + + /v1/admin/capsules/{id}/files/write: + parameters: + - name: id + in: path + required: true + schema: {type: string} + post: + summary: Upload a file to any capsule (admin) + operationId: adminUploadFile + tags: [admin] + security: + - sessionAuth: [] + requestBody: + required: true + content: + multipart/form-data: + schema: + type: object + required: [path, file] + properties: + path: {type: string} + file: {type: string, format: binary} + responses: + "204": + description: File uploaded + "409": + $ref: "#/components/responses/FailedPrecondition" + + /v1/admin/capsules/{id}/files/read: + parameters: + - name: id + in: path + required: true + schema: {type: string} + post: + summary: Download a file from any capsule (admin) + operationId: adminDownloadFile + tags: [admin] + security: + - sessionAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/ReadFileRequest" + responses: + "200": + description: File content + content: + application/octet-stream: + schema: + type: string + format: binary + "404": + $ref: "#/components/responses/NotFound" + + /v1/admin/capsules/{id}/files/list: + parameters: + - name: id + in: path + required: true + schema: {type: string} + post: + summary: List directory contents on any capsule (admin) + operationId: adminListDir + tags: [admin] + security: + - sessionAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/ListDirRequest" + responses: + "200": + description: Directory listing + content: + application/json: + schema: + $ref: "#/components/schemas/ListDirResponse" + "404": + $ref: "#/components/responses/NotFound" + "409": + $ref: "#/components/responses/FailedPrecondition" + + /v1/admin/capsules/{id}/files/mkdir: + parameters: + - name: id + in: path + required: true + schema: {type: string} + post: + summary: Create a directory on any capsule (admin) + operationId: adminMakeDir + tags: [admin] + security: + - sessionAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/MakeDirRequest" + responses: + "200": + description: Directory created + content: + application/json: + schema: + $ref: "#/components/schemas/MakeDirResponse" + "404": + $ref: "#/components/responses/NotFound" + "409": + $ref: "#/components/responses/FailedPrecondition" + + /v1/admin/capsules/{id}/files/remove: + parameters: + - name: id + in: path + required: true + schema: {type: string} + post: + summary: Remove a file or directory on any capsule (admin) + operationId: adminRemovePath + tags: [admin] + security: + - sessionAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/RemoveRequest" + responses: + "204": + description: File or directory removed + "404": + $ref: "#/components/responses/NotFound" + "409": + $ref: "#/components/responses/FailedPrecondition" + + /v1/admin/capsules/{id}/exec/stream: + parameters: + - name: id + in: path + required: true + schema: {type: string} + get: + summary: Stream command execution on any capsule via WebSocket (admin) + operationId: adminExecStream + tags: [admin] + security: + - sessionAuth: [] + description: | + Admin variant of /v1/capsules/{id}/exec/stream. Same protocol — WebSocket + upgrade, client sends `{"type":"start", "cmd":..., "args":...}` to start; + server streams stdout/stderr/exit frames. + responses: + "101": + description: WebSocket upgrade + "404": + $ref: "#/components/responses/NotFound" + "409": + $ref: "#/components/responses/FailedPrecondition" + + /v1/admin/capsules/{id}/pty: + parameters: + - name: id + in: path + required: true + schema: {type: string} + get: + summary: Interactive PTY session on any capsule via WebSocket (admin) + operationId: adminPtySession + tags: [admin] + security: + - sessionAuth: [] + description: | + Admin variant of /v1/capsules/{id}/pty. Same protocol — base64-encoded + PTY bytes, start/connect/input/resize/kill control messages, persistent + sessions reconnectable via tag. + responses: + "101": + description: WebSocket upgrade + "404": + $ref: "#/components/responses/NotFound" + "409": + $ref: "#/components/responses/FailedPrecondition" + + /v1/admin/capsules/{id}/processes/{selector}/stream: + parameters: + - name: id + in: path + required: true + schema: {type: string} + - name: selector + in: path + required: true + schema: {type: string} + description: Process PID (numeric) or tag (string) + get: + summary: Stream process output on any capsule via WebSocket (admin) + operationId: adminConnectProcess + tags: [admin] + security: + - sessionAuth: [] + responses: + "101": + description: WebSocket upgrade + "404": + $ref: "#/components/responses/NotFound" + components: + responses: + BadRequest: + description: Invalid request parameters + content: + application/json: + schema: + $ref: "#/components/schemas/Error" + + Unauthorized: + description: Missing or invalid auth + content: + application/json: + schema: + $ref: "#/components/schemas/Error" + + Forbidden: + description: Authenticated but not permitted (e.g. non-admin on /v1/admin/*) + content: + application/json: + schema: + $ref: "#/components/schemas/Error" + + NotFound: + description: Resource not found + content: + application/json: + schema: + $ref: "#/components/schemas/Error" + + FailedPrecondition: + description: Resource state does not allow this operation (e.g. exec on a paused capsule) + content: + application/json: + schema: + $ref: "#/components/schemas/Error" + securitySchemes: apiKeyAuth: type: apiKey @@ -2447,11 +3374,23 @@ components: name: X-API-Key description: API key for capsule lifecycle operations. Create via POST /v1/api-keys. - bearerAuth: - type: http - scheme: bearer - bearerFormat: JWT - description: JWT token from /v1/auth/login or /v1/auth/signup. Valid for 6 hours. + sessionAuth: + type: apiKey + in: cookie + name: wrenn_sid + description: | + Opaque session cookie set by POST /v1/auth/login, /v1/auth/activate, or + the OAuth callback. HttpOnly, Secure, SameSite=Strict. Idle window 6h, + absolute lifetime 24h. State-changing requests also require an + X-CSRF-Token header matching the wrenn_csrf cookie (double-submit). + csrfHeader: + type: apiKey + in: header + name: X-CSRF-Token + description: | + Double-submit CSRF token whose value must match the wrenn_csrf cookie. + Required on all non-GET requests authenticated via session cookie. + Not required for API key auth. hostTokenAuth: type: apiKey @@ -2491,12 +3430,13 @@ components: type: string description: Confirmation message instructing user to check email - AuthResponse: + SessionResponse: type: object + description: | + Returned by login, activate, and switch-team. The actual auth credential + is the wrenn_sid cookie set on the response. The body carries identity + data the SPA needs to bootstrap. properties: - token: - type: string - description: JWT token (valid for 6 hours) user_id: type: string team_id: @@ -2505,6 +3445,10 @@ components: type: string name: type: string + role: + type: string + is_admin: + type: boolean CreateAPIKeyRequest: type: object @@ -2549,13 +3493,22 @@ components: memory_mb: type: integer default: 512 + disk_size_mb: + type: integer + default: 5120 + description: > + Maximum size of the per-capsule copy-on-write disk in MB. Capped + at 5 GB by default; the actual size is max(disk_size_mb, origin + rootfs size). timeout_sec: type: integer + minimum: 0 default: 0 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. + no auto-pause. Positive values below 60 are silently clamped + to 60 (the agent's startup envelope). UsageResponse: type: object @@ -2664,6 +3617,17 @@ components: last_updated: type: string format: date-time + metadata: + type: object + additionalProperties: {type: string} + nullable: true + description: | + Free-form key/value labels attached at create-time. Also carries + agent-side version info (kernel_version, vmm_version, + agent_version, envd_version) when running. + disk_size_mb: + type: integer + nullable: true CreateSnapshotRequest: type: object @@ -2696,6 +3660,16 @@ components: created_at: type: string format: date-time + platform: + type: boolean + description: | + True when the template is platform-managed (visible to all teams, + e.g. the built-in `minimal` rootfs). False for team-owned + snapshot templates. + metadata: + type: object + additionalProperties: {type: string} + nullable: true ExecRequest: type: object @@ -2997,7 +3971,7 @@ components: type: array items: type: string - description: IDs of capsulees that would be destroyed on force-delete. + description: IDs of capsules that would be destroyed on force-delete. HostHasCapsulesError: type: object @@ -3014,7 +3988,7 @@ components: type: array items: type: string - description: IDs of active capsulees blocking deletion. + description: IDs of active capsules blocking deletion. AddTagRequest: type: object @@ -3104,7 +4078,7 @@ components: mem_bytes: type: integer format: int64 - description: "Resident memory in bytes (VmRSS of Firecracker process)" + description: "Resident memory in bytes (VmRSS of Cloud Hypervisor process)" disk_bytes: type: integer format: int64 @@ -3135,12 +4109,12 @@ components: items: type: string enum: - - capsule.created - - capsule.running - - capsule.paused - - capsule.destroyed - - template.snapshot.created - - template.snapshot.deleted + - capsule.create + - capsule.pause + - capsule.resume + - capsule.destroy + - template.snapshot.create + - template.snapshot.delete - host.up - host.down @@ -3180,12 +4154,12 @@ components: items: type: string enum: - - capsule.created - - capsule.running - - capsule.paused - - capsule.destroyed - - template.snapshot.created - - template.snapshot.deleted + - capsule.create + - capsule.pause + - capsule.resume + - capsule.destroy + - template.snapshot.create + - template.snapshot.delete - host.up - host.down @@ -3257,3 +4231,78 @@ components: type: string message: type: string + + AuditLogEntry: + type: object + properties: + id: {type: string} + actor_type: {type: string, enum: [user, api_key, host, system]} + actor_id: {type: string} + actor_name: {type: string} + resource_type: {type: string} + resource_id: {type: string} + action: {type: string} + scope: {type: string} + status: {type: string, enum: [success, failure]} + metadata: + type: object + additionalProperties: true + created_at: + type: string + format: date-time + + SSEEvent: + type: object + description: | + Wire format of one SSE message body. The event name (`event:` line) is + the `kind` and the JSON below is the `data:` line. + properties: + event: + type: string + enum: + - connected + - capsule.create + - capsule.pause + - capsule.resume + - capsule.destroy + - capsule.state.changed + - template.snapshot.create + - template.snapshot.delete + - host.up + - host.down + outcome: + type: string + enum: [success, error] + description: | + Present for action events (capsule.* except state.changed, + template.snapshot.*). Absent for host.up/down, capsule.state.changed, + and the connected sentinel. + resource: + type: object + properties: + id: {type: string} + type: {type: string} + actor: + type: object + properties: + type: {type: string, enum: [user, api_key, system]} + id: {type: string} + name: {type: string} + metadata: + type: object + additionalProperties: {type: string} + description: | + Event-specific context. Examples: `reason` (ttl_expired, + host_failure, cleanup_after_create_error, orphaned), + `host_ip`, `from`/`to` (for capsule.state.changed). + error: + type: string + description: Failure reason; only set when outcome=error. + sandbox: + allOf: + - $ref: "#/components/schemas/Capsule" + nullable: true + description: Populated for capsule.* events; null if DB lookup failed. + timestamp: + type: string + format: date-time diff --git a/src/wrenn/async_capsule.py b/src/wrenn/async_capsule.py index 44a9e84..4cf4c96 100644 --- a/src/wrenn/async_capsule.py +++ b/src/wrenn/async_capsule.py @@ -10,15 +10,54 @@ from contextlib import asynccontextmanager import httpx_ws from wrenn._git import AsyncGit -from wrenn.capsule import _DualMethod, _build_proxy_url +from wrenn.capsule import ( + _DEFAULT_WAIT_TIMEOUT, + _DESTROY_INTERVAL, + _FAIL_STATUSES, + _PAUSE_INTERVAL, + _RESUME_INTERVAL, + _START_INTERVAL, + _DualMethod, + _build_proxy_url, +) from wrenn.client import AsyncWrennClient from wrenn.commands import AsyncCommands +from wrenn.exceptions import WrennNotFoundError from wrenn.files import AsyncFiles from wrenn.models import Capsule as CapsuleModel from wrenn.models import Status, Template from wrenn.pty import AsyncPtySession +async def _apoll_until( + fetch, + targets: set[Status], + interval: float, + timeout: float = _DEFAULT_WAIT_TIMEOUT, + fail_on: set[Status] | None = None, +) -> CapsuleModel: + fail = fail_on if fail_on is not None else _FAIL_STATUSES + treat_missing_as_target = Status.missing in targets + deadline = time.monotonic() + timeout + last: CapsuleModel | None = None + while time.monotonic() < deadline: + try: + last = await fetch() + except WrennNotFoundError: + if treat_missing_as_target: + return CapsuleModel(status=Status.missing) + raise + if last.status in targets: + return last + if last.status is not None and last.status in fail: + raise RuntimeError(f"Capsule entered {last.status} state while waiting") + await asyncio.sleep(interval) + raise TimeoutError( + f"Capsule did not reach {targets} within {timeout}s " + f"(last status: {last.status if last else 'unknown'})" + ) + + class AsyncCapsule: """Async Wrenn capsule with e2b-compatible interface. @@ -139,15 +178,16 @@ class AsyncCapsule: client = AsyncWrennClient(api_key=api_key, base_url=base_url) info = await client.capsules.get(capsule_id) - if info.status == Status.paused: - await client.capsules.resume(capsule_id) - capsule = cls( _capsule_id=capsule_id, _client=client, _info=info, ) + if info.status == Status.pausing: + info = await capsule._wait_for_status({Status.paused}, _PAUSE_INTERVAL) + if info.status == Status.paused: + await client.capsules.resume(capsule_id) if info.status != Status.running: await capsule.wait_ready() @@ -160,22 +200,35 @@ class AsyncCapsule: resume = _DualMethod("_instance_resume", "_static_resume") get_info = _DualMethod("_instance_get_info", "_static_get_info") - async def _instance_destroy(self) -> None: + async def _instance_destroy(self, wait: bool = False) -> None: await self._client.capsules.destroy(self._id) + if wait: + await self._wait_for_status( + {Status.stopped, Status.missing}, _DESTROY_INTERVAL + ) @classmethod async def _static_destroy( cls, capsule_id: str, *, + wait: bool = False, api_key: str | None = None, base_url: str | None = None, ) -> None: async with AsyncWrennClient(api_key=api_key, base_url=base_url) as client: await client.capsules.destroy(capsule_id) + if wait: + await _apoll_until( + lambda: client.capsules.get(capsule_id), + {Status.stopped, Status.missing}, + _DESTROY_INTERVAL, + ) - async def _instance_pause(self) -> CapsuleModel: + async def _instance_pause(self, wait: bool = False) -> CapsuleModel: self._info = await self._client.capsules.pause(self._id) + if wait: + self._info = await self._wait_for_status({Status.paused}, _PAUSE_INTERVAL) return self._info @classmethod @@ -183,14 +236,24 @@ class AsyncCapsule: cls, capsule_id: str, *, + wait: bool = False, api_key: str | None = None, base_url: str | None = None, ) -> CapsuleModel: async with AsyncWrennClient(api_key=api_key, base_url=base_url) as client: - return await client.capsules.pause(capsule_id) + info = await client.capsules.pause(capsule_id) + if wait: + info = await _apoll_until( + lambda: client.capsules.get(capsule_id), + {Status.paused}, + _PAUSE_INTERVAL, + ) + return info - async def _instance_resume(self) -> CapsuleModel: + async def _instance_resume(self, wait: bool = False) -> CapsuleModel: self._info = await self._client.capsules.resume(self._id) + if wait: + self._info = await self._wait_for_status({Status.running}, _RESUME_INTERVAL) return self._info @classmethod @@ -198,11 +261,19 @@ class AsyncCapsule: cls, capsule_id: str, *, + wait: bool = False, api_key: str | None = None, base_url: str | None = None, ) -> CapsuleModel: async with AsyncWrennClient(api_key=api_key, base_url=base_url) as client: - return await client.capsules.resume(capsule_id) + info = await client.capsules.resume(capsule_id) + if wait: + info = await _apoll_until( + lambda: client.capsules.get(capsule_id), + {Status.running}, + _RESUME_INTERVAL, + ) + return info async def _instance_get_info(self) -> CapsuleModel: self._info = await self._client.capsules.get(self._id) @@ -229,43 +300,30 @@ class AsyncCapsule: """ await self._client.capsules.ping(self._id) - _POLL_INTERVALS: dict[Status, float] = { - Status.starting: 0.5, - Status.resuming: 0.5, - Status.pausing: 2.0, - Status.stopping: 1.0, - } + async def _wait_for_status( + self, + targets: set[Status], + interval: float, + timeout: float = _DEFAULT_WAIT_TIMEOUT, + ) -> CapsuleModel: + info = await _apoll_until( + lambda: self._client.capsules.get(self._id), + targets, + interval, + timeout, + fail_on={Status.error, Status.stopped, Status.missing} - targets, + ) + self._info = info + return info - async def wait_ready(self, timeout: float = 30) -> None: - """Await until the capsule status is ``running``. - - Polling interval adapts to the current transient status: - 0.5 s for starting/resuming, 2 s for pausing, 1 s for stopping. - - Args: - timeout (float): Maximum seconds to wait. Defaults to ``30``. + async def wait_ready(self, timeout: float = _DEFAULT_WAIT_TIMEOUT) -> None: + """Await until capsule status is ``running``. Raises: - TimeoutError: If the capsule does not reach ``running`` state - within ``timeout`` seconds. - RuntimeError: If the capsule enters an error, stopped, or paused - state while waiting. + TimeoutError: If capsule does not reach ``running`` within ``timeout``. + RuntimeError: If capsule enters error/stopped/missing while waiting. """ - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - info = await self._client.capsules.get(self._id) - if info.status == Status.running: - self._info = info - return - if info.status in (Status.error, Status.stopped): - raise RuntimeError(f"Capsule entered {info.status} state while waiting") - if info.status == Status.paused: - await self._client.capsules.resume(self._id) - interval = ( - self._POLL_INTERVALS.get(info.status, 0.5) if info.status else 0.5 - ) - await asyncio.sleep(interval) - raise TimeoutError(f"Capsule {self._id} did not become ready within {timeout}s") + await self._wait_for_status({Status.running}, _START_INTERVAL, timeout) async def is_running(self) -> bool: """Check whether the capsule is currently running. diff --git a/src/wrenn/capsule.py b/src/wrenn/capsule.py index 24dd4a5..f533205 100644 --- a/src/wrenn/capsule.py +++ b/src/wrenn/capsule.py @@ -13,6 +13,7 @@ import httpx_ws from wrenn._git import Git from wrenn.client import WrennClient from wrenn.commands import Commands +from wrenn.exceptions import WrennNotFoundError from wrenn.files import Files from wrenn.models import Capsule as CapsuleModel from wrenn.models import Status, Template @@ -28,6 +29,44 @@ def _build_proxy_url(base_url: str, capsule_id: str | None, port: int) -> str: return f"{scheme}://{port}-{capsule_id}.{host}" +_RESUME_INTERVAL = 0.5 +_DESTROY_INTERVAL = 0.5 +_PAUSE_INTERVAL = 2.0 +_START_INTERVAL = 0.5 +_DEFAULT_WAIT_TIMEOUT = 30.0 +_FAIL_STATUSES = {Status.error} + + +def _poll_until( + fetch, + targets: set[Status], + interval: float, + timeout: float = _DEFAULT_WAIT_TIMEOUT, + fail_on: set[Status] | None = None, +) -> CapsuleModel: + """Poll ``fetch()`` until status ∈ ``targets``. Raise on ``fail_on``/timeout.""" + fail = fail_on if fail_on is not None else _FAIL_STATUSES + treat_missing_as_target = Status.missing in targets + deadline = time.monotonic() + timeout + last: CapsuleModel | None = None + while time.monotonic() < deadline: + try: + last = fetch() + except WrennNotFoundError: + if treat_missing_as_target: + return CapsuleModel(status=Status.missing) + raise + if last.status in targets: + return last + if last.status is not None and last.status in fail: + raise RuntimeError(f"Capsule entered {last.status} state while waiting") + time.sleep(interval) + raise TimeoutError( + f"Capsule did not reach {targets} within {timeout}s " + f"(last status: {last.status if last else 'unknown'})" + ) + + class _DualMethod: """Descriptor that dispatches to instance method or classmethod depending on call site.""" @@ -100,9 +139,6 @@ class Capsule: self._id: str = _capsule_id self._client = _client self._info = _info - if self._id is None: - self._client.close() - raise RuntimeError("API returned a capsule without an ID") else: self._client = WrennClient(api_key=api_key, base_url=base_url) try: @@ -213,15 +249,16 @@ class Capsule: client = WrennClient(api_key=api_key, base_url=base_url) info = client.capsules.get(capsule_id) - if info.status == Status.paused: - client.capsules.resume(capsule_id) - capsule = cls( _capsule_id=capsule_id, _client=client, _info=info, ) + if info.status == Status.pausing: + info = capsule._wait_for_status({Status.paused}, _PAUSE_INTERVAL) + if info.status == Status.paused: + client.capsules.resume(capsule_id) if info.status != Status.running: capsule.wait_ready() @@ -234,25 +271,36 @@ class Capsule: resume = _DualMethod("_instance_resume", "_static_resume") get_info = _DualMethod("_instance_get_info", "_static_get_info") - def _instance_destroy(self) -> None: - """Destroy this capsule.""" + def _instance_destroy(self, wait: bool = False) -> None: + """Destroy this capsule. If ``wait``, poll until stopped/missing.""" self._client.capsules.destroy(self._id) + if wait: + self._wait_for_status({Status.stopped, Status.missing}, _DESTROY_INTERVAL) @classmethod def _static_destroy( cls, capsule_id: str, *, + wait: bool = False, api_key: str | None = None, base_url: str | None = None, ) -> None: """Destroy a capsule by ID.""" with WrennClient(api_key=api_key, base_url=base_url) as client: client.capsules.destroy(capsule_id) + if wait: + _poll_until( + lambda: client.capsules.get(capsule_id), + {Status.stopped, Status.missing}, + _DESTROY_INTERVAL, + ) - def _instance_pause(self) -> CapsuleModel: - """Pause this capsule.""" + def _instance_pause(self, wait: bool = False) -> CapsuleModel: + """Pause this capsule. If ``wait``, poll until ``paused``.""" self._info = self._client.capsules.pause(self._id) + if wait: + self._info = self._wait_for_status({Status.paused}, _PAUSE_INTERVAL) return self._info @classmethod @@ -260,16 +308,26 @@ class Capsule: cls, capsule_id: str, *, + wait: bool = False, api_key: str | None = None, base_url: str | None = None, ) -> CapsuleModel: """Pause a capsule by ID.""" with WrennClient(api_key=api_key, base_url=base_url) as client: - return client.capsules.pause(capsule_id) + info = client.capsules.pause(capsule_id) + if wait: + info = _poll_until( + lambda: client.capsules.get(capsule_id), + {Status.paused}, + _PAUSE_INTERVAL, + ) + return info - def _instance_resume(self) -> CapsuleModel: - """Resume this capsule.""" + def _instance_resume(self, wait: bool = False) -> CapsuleModel: + """Resume this capsule. If ``wait``, poll until ``running``.""" self._info = self._client.capsules.resume(self._id) + if wait: + self._info = self._wait_for_status({Status.running}, _RESUME_INTERVAL) return self._info @classmethod @@ -277,12 +335,20 @@ class Capsule: cls, capsule_id: str, *, + wait: bool = False, api_key: str | None = None, base_url: str | None = None, ) -> CapsuleModel: """Resume a capsule by ID.""" with WrennClient(api_key=api_key, base_url=base_url) as client: - return client.capsules.resume(capsule_id) + info = client.capsules.resume(capsule_id) + if wait: + info = _poll_until( + lambda: client.capsules.get(capsule_id), + {Status.running}, + _RESUME_INTERVAL, + ) + return info def _instance_get_info(self) -> CapsuleModel: """Get current info for this capsule.""" @@ -311,43 +377,30 @@ class Capsule: """ self._client.capsules.ping(self._id) - _POLL_INTERVALS: dict[Status, float] = { - Status.starting: 0.5, - Status.resuming: 0.5, - Status.pausing: 2.0, - Status.stopping: 1.0, - } + def _wait_for_status( + self, + targets: set[Status], + interval: float, + timeout: float = _DEFAULT_WAIT_TIMEOUT, + ) -> CapsuleModel: + info = _poll_until( + lambda: self._client.capsules.get(self._id), + targets, + interval, + timeout, + fail_on={Status.error, Status.stopped, Status.missing} - targets, + ) + self._info = info + return info - def wait_ready(self, timeout: float = 30) -> None: - """Block until the capsule status is ``running``. - - Polling interval adapts to the current transient status: - 0.5 s for starting/resuming, 2 s for pausing, 1 s for stopping. - - Args: - timeout (float): Maximum seconds to wait. Defaults to ``30``. + def wait_ready(self, timeout: float = _DEFAULT_WAIT_TIMEOUT) -> None: + """Block until capsule status is ``running``. Raises: - TimeoutError: If the capsule does not reach ``running`` state - within ``timeout`` seconds. - RuntimeError: If the capsule enters an error, stopped, or paused - state while waiting. + TimeoutError: If capsule does not reach ``running`` within ``timeout``. + RuntimeError: If capsule enters error/stopped/missing while waiting. """ - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - info = self._client.capsules.get(self._id) - if info.status == Status.running: - self._info = info - return - if info.status in (Status.error, Status.stopped): - raise RuntimeError(f"Capsule entered {info.status} state while waiting") - if info.status == Status.paused: - self._client.capsules.resume(self._id) - interval = ( - self._POLL_INTERVALS.get(info.status, 0.5) if info.status else 0.5 - ) - time.sleep(interval) - raise TimeoutError(f"Capsule {self._id} did not become ready within {timeout}s") + self._wait_for_status({Status.running}, _START_INTERVAL, timeout) def is_running(self) -> bool: """Check whether the capsule is currently running. diff --git a/src/wrenn/files.py b/src/wrenn/files.py index 477aeca..5a99289 100644 --- a/src/wrenn/files.py +++ b/src/wrenn/files.py @@ -9,6 +9,36 @@ from wrenn.exceptions import WrennNotFoundError, _raise_for_status, handle_respo from wrenn.models import FileEntry, ListDirResponse, MakeDirResponse +def _is_already_exists(resp: httpx.Response) -> bool: + """Detect server's already-exists reply across status codes / code strings. + + Server may return 409 with code "conflict"/"already_exists" or wrap + "already_exists" inside an "internal" 500 message. + """ + if resp.status_code < 400: + return False + try: + body = resp.json() + except Exception: + return False + err = body.get("error", {}) if isinstance(body, dict) else {} + code = err.get("code", "") + msg = err.get("message", "") or "" + return code in {"conflict", "already_exists"} or "already_exists" in msg + + +def _find_entry(list_fn, path: str) -> FileEntry | None: + parent = os.path.dirname(path) + name = os.path.basename(path) + try: + for entry in list_fn(parent, depth=1): + if entry.name == name: + return entry + except WrennNotFoundError: + return None + return None + + class Files: """Sync filesystem interface. Accessed via ``capsule.files``.""" @@ -118,17 +148,10 @@ class Files: f"/v1/capsules/{self._capsule_id}/files/mkdir", json={"path": path}, ) - if resp.status_code == 409: - try: - body = resp.json() - if body.get("error", {}).get("code") == "conflict": - parent = os.path.dirname(path) - name = os.path.basename(path) - for entry in self.list(parent, depth=1): - if entry.name == name: - return entry - except Exception: - pass + if _is_already_exists(resp): + existing = _find_entry(self.list, path) + if existing is not None: + return existing parsed = MakeDirResponse.model_validate(handle_response(resp)) if parsed.entry is None: raise RuntimeError("mkdir response missing entry") @@ -315,17 +338,12 @@ class AsyncFiles: f"/v1/capsules/{self._capsule_id}/files/mkdir", json={"path": path}, ) - if resp.status_code == 409: - try: - body = resp.json() - if body.get("error", {}).get("code") == "conflict": - parent = os.path.dirname(path) - name = os.path.basename(path) - for entry in await self.list(parent, depth=1): - if entry.name == name: - return entry - except Exception: - pass + if _is_already_exists(resp): + parent = os.path.dirname(path) + name = os.path.basename(path) + for entry in await self.list(parent, depth=1): + if entry.name == name: + return entry parsed = MakeDirResponse.model_validate(handle_response(resp)) if parsed.entry is None: raise RuntimeError("mkdir response missing entry") diff --git a/src/wrenn/models/__init__.py b/src/wrenn/models/__init__.py index 5628e11..6fe5eb8 100644 --- a/src/wrenn/models/__init__.py +++ b/src/wrenn/models/__init__.py @@ -1,6 +1,5 @@ from wrenn.models._generated import ( APIKeyResponse, - AuthResponse, Capsule, CreateAPIKeyRequest, CreateCapsuleRequest, @@ -34,7 +33,6 @@ from wrenn.models._generated import ( __all__ = [ "APIKeyResponse", - "AuthResponse", "CreateAPIKeyRequest", "CreateHostRequest", "CreateHostResponse", diff --git a/src/wrenn/models/_generated.py b/src/wrenn/models/_generated.py index 8bcec35..8eb7425 100644 --- a/src/wrenn/models/_generated.py +++ b/src/wrenn/models/_generated.py @@ -1,10 +1,10 @@ # generated by datamodel-codegen: # filename: openapi.yaml -# timestamp: 2026-05-15T07:57:28+00:00 +# timestamp: 2026-05-19T08:54:50+00:00 from __future__ import annotations from pydantic import AwareDatetime, BaseModel, EmailStr, Field -from typing import Annotated +from typing import Annotated, Any from datetime import date as date_aliased from enum import StrEnum @@ -27,14 +27,20 @@ class SignupResponse(BaseModel): ] = None -class AuthResponse(BaseModel): - token: Annotated[str | None, Field(description="JWT token (valid for 6 hours)")] = ( - None - ) +class SessionResponse(BaseModel): + """ + Returned by login, activate, and switch-team. The actual auth credential + is the wrenn_sid cookie set on the response. The body carries identity + data the SPA needs to bootstrap. + + """ + user_id: str | None = None team_id: str | None = None email: str | None = None name: str | None = None + role: str | None = None + is_admin: bool | None = None class CreateAPIKeyRequest(BaseModel): @@ -62,10 +68,17 @@ class CreateCapsuleRequest(BaseModel): template: str | None = "minimal" vcpus: int | None = 1 memory_mb: int | None = 512 + disk_size_mb: Annotated[ + int | None, + Field( + description="Maximum size of the per-capsule copy-on-write disk in MB. Capped at 5 GB by default; the actual size is max(disk_size_mb, origin rootfs size).\n" + ), + ] = 5120 timeout_sec: Annotated[ int | None, Field( - description="Auto-pause TTL in seconds. The capsule is automatically paused after this duration of inactivity (no exec or ping). 0 means no auto-pause.\n" + 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. Positive values below 60 are silently clamped to 60 (the agent's startup envelope).\n", + ge=0, ), ] = 0 @@ -156,6 +169,13 @@ class Capsule(BaseModel): started_at: AwareDatetime | None = None last_active_at: AwareDatetime | None = None last_updated: AwareDatetime | None = None + metadata: Annotated[ + dict[str, str] | None, + Field( + description="Free-form key/value labels attached at create-time. Also carries\nagent-side version info (kernel_version, vmm_version,\nagent_version, envd_version) when running.\n" + ), + ] = None + disk_size_mb: int | None = None class CreateSnapshotRequest(BaseModel): @@ -180,6 +200,13 @@ class Template(BaseModel): memory_mb: int | None = None size_bytes: int | None = None created_at: AwareDatetime | None = None + platform: Annotated[ + bool | None, + Field( + description="True when the template is platform-managed (visible to all teams,\ne.g. the built-in `minimal` rootfs). False for team-owned\nsnapshot templates.\n" + ), + ] = None + metadata: dict[str, str] | None = None class ExecRequest(BaseModel): @@ -402,7 +429,7 @@ class HostDeletePreview(BaseModel): host: Host | None = None sandbox_ids: Annotated[ list[str] | None, - Field(description="IDs of capsulees that would be destroyed on force-delete."), + Field(description="IDs of capsules that would be destroyed on force-delete."), ] = None @@ -410,8 +437,7 @@ class Error(BaseModel): code: Annotated[str | None, Field(examples=["host_has_sandboxes"])] = None message: str | None = None sandbox_ids: Annotated[ - list[str] | None, - Field(description="IDs of active capsulees blocking deletion."), + list[str] | None, Field(description="IDs of active capsules blocking deletion.") ] = None @@ -479,7 +505,9 @@ class MetricPoint(BaseModel): ] = None mem_bytes: Annotated[ int | None, - Field(description="Resident memory in bytes (VmRSS of Firecracker process)"), + Field( + description="Resident memory in bytes (VmRSS of Cloud Hypervisor process)" + ), ] = None disk_bytes: Annotated[ int | None, Field(description="Allocated disk bytes for the CoW sparse file") @@ -497,12 +525,12 @@ class Provider(StrEnum): 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" + capsule_create = "capsule.create" + capsule_pause = "capsule.pause" + capsule_resume = "capsule.resume" + capsule_destroy = "capsule.destroy" + template_snapshot_create = "template.snapshot.create" + template_snapshot_delete = "template.snapshot.delete" host_up = "host.up" host_down = "host.down" @@ -594,6 +622,106 @@ class Error1(BaseModel): error: Error2 | None = None +class ActorType(StrEnum): + user = "user" + api_key = "api_key" + host = "host" + system = "system" + + +class Status2(StrEnum): + success = "success" + failure = "failure" + + +class AuditLogEntry(BaseModel): + id: str | None = None + actor_type: ActorType | None = None + actor_id: str | None = None + actor_name: str | None = None + resource_type: str | None = None + resource_id: str | None = None + action: str | None = None + scope: str | None = None + status: Status2 | None = None + metadata: dict[str, Any] | None = None + created_at: AwareDatetime | None = None + + +class Event2(StrEnum): + connected = "connected" + capsule_create = "capsule.create" + capsule_pause = "capsule.pause" + capsule_resume = "capsule.resume" + capsule_destroy = "capsule.destroy" + capsule_state_changed = "capsule.state.changed" + template_snapshot_create = "template.snapshot.create" + template_snapshot_delete = "template.snapshot.delete" + host_up = "host.up" + host_down = "host.down" + + +class Outcome(StrEnum): + """ + Present for action events (capsule.* except state.changed, + template.snapshot.*). Absent for host.up/down, capsule.state.changed, + and the connected sentinel. + + """ + + success = "success" + error = "error" + + +class Resource(BaseModel): + id: str | None = None + type: str | None = None + + +class Type4(StrEnum): + user = "user" + api_key = "api_key" + system = "system" + + +class Actor(BaseModel): + type: Type4 | None = None + id: str | None = None + name: str | None = None + + +class SSEEvent(BaseModel): + """ + Wire format of one SSE message body. The event name (`event:` line) is + the `kind` and the JSON below is the `data:` line. + + """ + + event: Event2 | None = None + outcome: Annotated[ + Outcome | None, + Field( + description="Present for action events (capsule.* except state.changed,\ntemplate.snapshot.*). Absent for host.up/down, capsule.state.changed,\nand the connected sentinel.\n" + ), + ] = None + resource: Resource | None = None + actor: Actor | None = None + metadata: Annotated[ + dict[str, str] | None, + Field( + description="Event-specific context. Examples: `reason` (ttl_expired,\nhost_failure, cleanup_after_create_error, orphaned),\n`host_ip`, `from`/`to` (for capsule.state.changed).\n" + ), + ] = None + error: Annotated[ + str | None, Field(description="Failure reason; only set when outcome=error.") + ] = None + sandbox: Annotated[ + Capsule | None, + Field(description="Populated for capsule.* events; null if DB lookup failed."), + ] = None + timestamp: AwareDatetime | None = None + + class ListDirResponse(BaseModel): entries: list[FileEntry] | None = None diff --git a/tests/test_integration.py b/tests/test_integration.py index ff66983..d280d2c 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -46,7 +46,7 @@ class TestCapsuleLifecycle: assert capsule_id assert capsule.info is not None finally: - capsule.destroy() + capsule.destroy(wait=True) info = Capsule.get_info(capsule_id) assert info.status in (Status.stopped, Status.missing) @@ -65,7 +65,7 @@ class TestCapsuleLifecycle: assert capsule.is_running() info = Capsule.get_info(capsule_id) - assert info.status in (Status.stopped, Status.missing) + assert info.status in (Status.stopping, Status.stopped, Status.missing) def test_get_info(self): capsule = Capsule(wait=True) @@ -80,11 +80,11 @@ class TestCapsuleLifecycle: def test_pause_and_resume(self): capsule = Capsule(wait=True) try: - paused = capsule.pause() + paused = capsule.pause(wait=True) assert paused.status == Status.paused assert not capsule.is_running() - resumed = capsule.resume() + resumed = capsule.resume(wait=True) assert resumed.status == Status.running finally: capsule.destroy() @@ -93,7 +93,7 @@ class TestCapsuleLifecycle: capsule = Capsule(wait=True) capsule_id = capsule.capsule_id try: - Capsule.destroy(capsule_id) + Capsule.destroy(capsule_id, wait=True) except Exception: capsule.destroy() raise