From 213af4aee7889b8856a643d04759f654faf8ec8e Mon Sep 17 00:00:00 2001 From: pptx704 Date: Sat, 2 May 2026 04:44:26 +0600 Subject: [PATCH] Increased timeout for long running API calls and updated typehints --- .pre-commit-config.yaml | 25 +++++++++++++++++++++ src/wrenn/async_capsule.py | 8 ++++--- src/wrenn/capsule.py | 11 ++++----- src/wrenn/client.py | 15 +++++++++---- src/wrenn/code_interpreter/async_capsule.py | 2 +- src/wrenn/code_interpreter/capsule.py | 2 +- src/wrenn/commands.py | 19 +++++++++++----- 7 files changed, 62 insertions(+), 20 deletions(-) create mode 100644 .pre-commit-config.yaml diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..8dc6f53 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,25 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.15.10 + hooks: + - id: ruff + - id: ruff-format + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.20.0 + hooks: + - id: mypy + additional_dependencies: + - pydantic>=2.12.5 + - httpx>=0.28.1 + - httpx-ws>=0.9.0 + - email-validator>=2.3.0 + + - repo: local + hooks: + - id: unit-tests + name: unit tests + entry: uv run pytest -m "not integration" -x -q + language: system + pass_filenames: false + always_run: true diff --git a/src/wrenn/async_capsule.py b/src/wrenn/async_capsule.py index 3e92de7..509f53a 100644 --- a/src/wrenn/async_capsule.py +++ b/src/wrenn/async_capsule.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import builtins import time from collections.abc import AsyncIterator from contextlib import asynccontextmanager @@ -102,6 +103,7 @@ class AsyncCapsule: memory_mb=memory_mb, timeout_sec=timeout, ) + assert info.id is not None capsule = cls( _capsule_id=info.id, _client=client, @@ -284,7 +286,7 @@ class AsyncCapsule: async def pty( self, cmd: str = "/bin/bash", - args: list[str] | None = None, + args: builtins.list[str] | None = None, cols: int = 80, rows: int = 24, envs: dict[str, str] | None = None, @@ -316,7 +318,7 @@ class AsyncCapsule: """ async with httpx_ws.aconnect_ws( f"/v1/capsules/{self._id}/pty", client=self._client.http - ) as ws: + ) as ws: # type: httpx_ws.AsyncWebSocketSession session = AsyncPtySession(ws, self._id) await session._send_start( cmd=cmd, args=args, cols=cols, rows=rows, envs=envs, cwd=cwd @@ -335,7 +337,7 @@ class AsyncCapsule: """ async with httpx_ws.aconnect_ws( f"/v1/capsules/{self._id}/pty", client=self._client.http - ) as ws: + ) as ws: # type: httpx_ws.AsyncWebSocketSession session = AsyncPtySession(ws, self._id) await session._send_connect(tag) yield session diff --git a/src/wrenn/capsule.py b/src/wrenn/capsule.py index 400409f..a5ed36a 100644 --- a/src/wrenn/capsule.py +++ b/src/wrenn/capsule.py @@ -1,5 +1,6 @@ from __future__ import annotations +import builtins import time from collections.abc import Iterator from contextlib import contextmanager @@ -94,9 +95,8 @@ class Capsule: ``WRENN_BASE_URL`` or the default production endpoint. """ if _capsule_id is not None: - # Internal construction path (from create/connect classmethods) assert _client is not None - self._id = _capsule_id + self._id: str = _capsule_id self._client = _client self._info = _info else: @@ -108,6 +108,7 @@ class Capsule: memory_mb=memory_mb, timeout_sec=timeout, ) + assert self._info.id is not None self._id = self._info.id self.commands = Commands(self._id, self._client.http) @@ -360,7 +361,7 @@ class Capsule: def pty( self, cmd: str = "/bin/bash", - args: list[str] | None = None, + args: builtins.list[str] | None = None, cols: int = 80, rows: int = 24, envs: dict[str, str] | None = None, @@ -391,7 +392,7 @@ class Capsule: """ with httpx_ws.connect_ws( f"/v1/capsules/{self._id}/pty", client=self._client.http - ) as ws: + ) as ws: # type: httpx_ws.WebSocketSession session = PtySession(ws, self._id) session._send_start( cmd=cmd, args=args, cols=cols, rows=rows, envs=envs, cwd=cwd @@ -410,7 +411,7 @@ class Capsule: """ with httpx_ws.connect_ws( f"/v1/capsules/{self._id}/pty", client=self._client.http - ) as ws: + ) as ws: # type: httpx_ws.WebSocketSession session = PtySession(ws, self._id) session._send_connect(tag) yield session diff --git a/src/wrenn/client.py b/src/wrenn/client.py index c927396..c51b190 100644 --- a/src/wrenn/client.py +++ b/src/wrenn/client.py @@ -6,6 +6,7 @@ import httpx from wrenn._config import DEFAULT_BASE_URL, ENV_API_KEY, ENV_BASE_URL from wrenn.exceptions import handle_response + from wrenn.models import ( Template, ) @@ -13,6 +14,8 @@ from wrenn.models import ( Capsule as CapsuleModel, ) +_LONG_TIMEOUT = httpx.Timeout(60.0) + def _resolve_api_key(api_key: str | None) -> str: resolved = api_key or os.environ.get(ENV_API_KEY) @@ -108,7 +111,7 @@ class CapsulesResource: Raises: WrennNotFoundError: If no capsule with the given ID exists. """ - resp = self._http.post(f"/v1/capsules/{id}/pause") + resp = self._http.post(f"/v1/capsules/{id}/pause", timeout=_LONG_TIMEOUT) return CapsuleModel.model_validate(handle_response(resp)) def resume(self, id: str) -> CapsuleModel: @@ -224,7 +227,7 @@ class AsyncCapsulesResource: Raises: WrennNotFoundError: If no capsule with the given ID exists. """ - resp = await self._http.post(f"/v1/capsules/{id}/pause") + resp = await self._http.post(f"/v1/capsules/{id}/pause", timeout=_LONG_TIMEOUT) return CapsuleModel.model_validate(handle_response(resp)) async def resume(self, id: str) -> CapsuleModel: @@ -285,7 +288,9 @@ class SnapshotsResource: params: dict = {} if overwrite: params["overwrite"] = "true" - resp = self._http.post("/v1/snapshots", json=payload, params=params) + resp = self._http.post( + "/v1/snapshots", json=payload, params=params, timeout=_LONG_TIMEOUT + ) return Template.model_validate(handle_response(resp)) def list(self, type: str | None = None) -> list[Template]: @@ -347,7 +352,9 @@ class AsyncSnapshotsResource: params: dict = {} if overwrite: params["overwrite"] = "true" - resp = await self._http.post("/v1/snapshots", json=payload, params=params) + resp = await self._http.post( + "/v1/snapshots", json=payload, params=params, timeout=_LONG_TIMEOUT + ) return Template.model_validate(handle_response(resp)) async def list(self, type: str | None = None) -> list[Template]: diff --git a/src/wrenn/code_interpreter/async_capsule.py b/src/wrenn/code_interpreter/async_capsule.py index fb99752..f61937c 100644 --- a/src/wrenn/code_interpreter/async_capsule.py +++ b/src/wrenn/code_interpreter/async_capsule.py @@ -207,7 +207,7 @@ class AsyncCapsule(BaseAsyncCapsule): deadline = time.monotonic() + timeout headers = {"X-API-Key": self._client._api_key} - async with httpx_ws.aconnect_ws(ws_url, headers=headers) as ws: + async with httpx_ws.aconnect_ws(ws_url, headers=headers) as ws: # type: httpx_ws.AsyncWebSocketSession await ws.send_text(json.dumps(msg)) while time.monotonic() < deadline: time_left = deadline - time.monotonic() diff --git a/src/wrenn/code_interpreter/capsule.py b/src/wrenn/code_interpreter/capsule.py index 1b1f7ea..344c3f3 100644 --- a/src/wrenn/code_interpreter/capsule.py +++ b/src/wrenn/code_interpreter/capsule.py @@ -233,7 +233,7 @@ class Capsule(BaseCapsule): deadline = time.monotonic() + timeout headers = {"X-API-Key": self._client._api_key} - with httpx_ws.connect_ws(ws_url, headers=headers) as ws: + with httpx_ws.connect_ws(ws_url, headers=headers) as ws: # type: httpx_ws.WebSocketSession ws.send_text(json.dumps(msg)) while time.monotonic() < deadline: time_left = deadline - time.monotonic() diff --git a/src/wrenn/commands.py b/src/wrenn/commands.py index 7ca9f44..e24f898 100644 --- a/src/wrenn/commands.py +++ b/src/wrenn/commands.py @@ -1,6 +1,7 @@ from __future__ import annotations import base64 +import builtins import json from collections.abc import AsyncIterator, Iterator from dataclasses import dataclass @@ -199,6 +200,7 @@ class Commands: resp = self._http.post(f"/v1/capsules/{self._capsule_id}/exec", json=payload) data = handle_response(resp) + assert isinstance(data, dict) if background: return CommandHandle( @@ -217,6 +219,7 @@ class Commands: """ resp = self._http.get(f"/v1/capsules/{self._capsule_id}/processes") data = handle_response(resp) + assert isinstance(data, dict) return [ ProcessInfo( pid=p.get("pid", 0), @@ -252,7 +255,7 @@ class Commands: with httpx_ws.connect_ws( f"/v1/capsules/{self._capsule_id}/processes/{pid}/stream", self._http, - ) as ws: + ) as ws: # type: httpx_ws.WebSocketSession while True: try: raw = ws.receive_json() @@ -263,7 +266,9 @@ class Commands: except httpx_ws.WebSocketDisconnect: break - def stream(self, cmd: str, args: list[str] | None = None) -> Iterator[StreamEvent]: + def stream( + self, cmd: str, args: builtins.list[str] | None = None + ) -> Iterator[StreamEvent]: """Execute a command via WebSocket, streaming output as events. Args: @@ -280,7 +285,7 @@ class Commands: with httpx_ws.connect_ws( f"/v1/capsules/{self._capsule_id}/exec/stream", self._http, - ) as ws: + ) as ws: # type: httpx_ws.WebSocketSession if args: start_msg: dict = {"type": "start", "cmd": cmd, "args": args} else: @@ -378,6 +383,7 @@ class AsyncCommands: f"/v1/capsules/{self._capsule_id}/exec", json=payload ) data = handle_response(resp) + assert isinstance(data, dict) if background: return CommandHandle( @@ -396,6 +402,7 @@ class AsyncCommands: """ resp = await self._http.get(f"/v1/capsules/{self._capsule_id}/processes") data = handle_response(resp) + assert isinstance(data, dict) return [ ProcessInfo( pid=p.get("pid", 0), @@ -433,7 +440,7 @@ class AsyncCommands: async with httpx_ws.aconnect_ws( f"/v1/capsules/{self._capsule_id}/processes/{pid}/stream", self._http, - ) as ws: + ) as ws: # type: httpx_ws.AsyncWebSocketSession try: while True: raw = await ws.receive_json() @@ -445,7 +452,7 @@ class AsyncCommands: pass async def stream( - self, cmd: str, args: list[str] | None = None + self, cmd: str, args: builtins.list[str] | None = None ) -> AsyncIterator[StreamEvent]: """Execute a command via WebSocket, streaming output as events. @@ -463,7 +470,7 @@ class AsyncCommands: async with httpx_ws.aconnect_ws( f"/v1/capsules/{self._capsule_id}/exec/stream", self._http, - ) as ws: + ) as ws: # type: httpx_ws.AsyncWebSocketSession if args: start_msg: dict = {"type": "start", "cmd": cmd, "args": args} else: