make ws recv timeout configurable (#517)

Signed-off-by: Jessie Frazelle <github@jessfraz.com>

Author: jessfraz

CHANGELOG.md

@@ -5,6 +5,16 @@ All notable changes to the KittyCAD Python SDK will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## v1.2.5
+
+### Added
+
+- WebSocket wrappers now honor configurable receive timeouts. Set a client-wide default with `KittyCAD(..., websocket_recv_timeout=120)` or override per connection via `client.modeling.modeling_commands_ws(recv_timeout=300)`. Existing code keeps the previous 60 second timeout without changes.
+
+### Migration
+
+- No action required unless you wish to change the timeout. If you relied on the old hardcoded 60 second timeout, explicitly pass `recv_timeout=60` to preserve that behavior when using a client configured with a different default.
+
 ## v1.1.2
 
 ### Changed - WebSocket message typing
@@ -720,6 +730,7 @@ Example exception message:
 - **Easier Debugging**: No need to dig into Error objects
 - **Cleaner Code**: No error checking boilerplate needed
 - **Better IDE Support**: Improved autocomplete and type checking
+
 ## v1.1.3
 
 ### Fixed – Typed responses instead of raw dicts

generate/templates/__init__.py.jinja2

@@ -1,7 +1,7 @@
 """The KittyCAD Python SDK"""
 
 import os
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 import httpx
 
 {% if has_websockets %}
@@ -51,12 +51,12 @@ class {{ tag_pascal }}API:
 
     {% for func_name, func_info in functions.items() %}
     {% if func_info.is_websocket and func_info.has_websocket_class %}
-    def {{ func_name }}(self{{ func_info.websocket_params if func_info.websocket_params else '' }}) -> "WebSocket{{ func_name|to_pascal_case }}":
+    def {{ func_name }}(self{{ func_info.websocket_params if func_info.websocket_params else '' }}, recv_timeout: Optional[float] = None, ws_factory: Optional[Callable[..., ClientConnectionSync]] = None) -> "WebSocket{{ func_name|to_pascal_case }}":
         """{{ func_info.description }}
 
         Returns a WebSocket wrapper with methods for sending/receiving data.
         """
-        return WebSocket{{ func_name|to_pascal_case }}({{ func_info.websocket_call_args if func_info.websocket_call_args else '' }}{{ ', ' if func_info.websocket_call_args else '' }}client=self.client)
+        return WebSocket{{ func_name|to_pascal_case }}({{ func_info.websocket_call_args if func_info.websocket_call_args else '' }}{{ ', ' if func_info.websocket_call_args else '' }}recv_timeout=recv_timeout, ws_factory=ws_factory, client=self.client)
     {% else %}
 {{ func_info.sync_implementation }}
     {% endif %}
@@ -93,7 +93,7 @@ class WebSocket{{ func_name|to_pascal_case }}:
 
     ws: ClientConnectionSync
 
-    def __init__(self{{ func_info.websocket_params if func_info.websocket_params else '' }}, *, client: Client):
+    def __init__(self{{ func_info.websocket_params if func_info.websocket_params else '' }}, recv_timeout: Optional[float] = None, ws_factory: Optional[Callable[..., ClientConnectionSync]] = None, *, client: Client):
         # Inline WebSocket connection logic
         {% set path_args = func_info.ws_args|selectattr("in_url")|list %}
         url = ("{}" + "{{ func_info.path }}").format(client.base_url{% for arg in path_args %}, {{ arg.name }}={{ arg.name }}{% endfor %})
@@ -114,7 +114,11 @@ class WebSocket{{ func_name|to_pascal_case }}:
         {% endif %}
         {% endfor %}
         headers = client.get_headers()
-        self.ws = ws_connect(url.replace("http", "ws"), additional_headers=headers, close_timeout=120, max_size=None)
+        factory = ws_factory or ws_connect
+        self.ws = factory(url.replace("http", "ws"), additional_headers=headers, close_timeout=120, max_size=None)
+        self._recv_timeout = (
+            client.get_websocket_recv_timeout() if recv_timeout is None else recv_timeout
+        )
 
     def __enter__(self):
         return self
@@ -156,7 +160,7 @@ class WebSocket{{ func_name|to_pascal_case }}:
 
     def recv(self) -> {{ func_info.ws_response_type }}:
         """Receive data from the websocket."""
-        message = self.ws.recv(timeout=60)
+        message = self.ws.recv(timeout=self._recv_timeout)
         {% if func_info.ws_response_is_dict %}
         return json.loads(message)
         {% else %}

kittycad/__init__.py

@@ -2,7 +2,7 @@
 
 import json
 import os
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import bson
 import httpx
@@ -1274,22 +1274,37 @@ def create_text_to_cad_part_feedback(
         return response.json() if response.content else None
 
     def ml_copilot_ws(
-        self, conversation_id: Optional[str] = None, replay: Optional[bool] = None
+        self,
+        conversation_id: Optional[str] = None,
+        replay: Optional[bool] = None,
+        recv_timeout: Optional[float] = None,
+        ws_factory: Optional[Callable[..., ClientConnectionSync]] = None,
     ) -> "WebSocketMlCopilotWs":
         """Open a websocket to prompt the ML copilot.
 
         Returns a WebSocket wrapper with methods for sending/receiving data.
         """
         return WebSocketMlCopilotWs(
-            conversation_id=conversation_id, replay=replay, client=self.client
+            conversation_id=conversation_id,
+            replay=replay,
+            recv_timeout=recv_timeout,
+            ws_factory=ws_factory,
+            client=self.client,
         )
 
-    def ml_reasoning_ws(self, id: str) -> "WebSocketMlReasoningWs":
+    def ml_reasoning_ws(
+        self,
+        id: str,
+        recv_timeout: Optional[float] = None,
+        ws_factory: Optional[Callable[..., ClientConnectionSync]] = None,
+    ) -> "WebSocketMlReasoningWs":
         """Open a websocket to prompt the ML copilot.
 
         Returns a WebSocket wrapper with methods for sending/receiving data.
         """
-        return WebSocketMlReasoningWs(id=id, client=self.client)
+        return WebSocketMlReasoningWs(
+            id=id, recv_timeout=recv_timeout, ws_factory=ws_factory, client=self.client
+        )
 
 
 class AsyncMlAPI:
@@ -4768,12 +4783,18 @@ def create_file_execution(
         # Validate into a Pydantic model (works for BaseModel and RootModel)
         return CodeOutput.model_validate(json_data)
 
-    def create_executor_term(self) -> "WebSocketCreateExecutorTerm":
+    def create_executor_term(
+        self,
+        recv_timeout: Optional[float] = None,
+        ws_factory: Optional[Callable[..., ClientConnectionSync]] = None,
+    ) -> "WebSocketCreateExecutorTerm":
         """Create a terminal.
 
         Returns a WebSocket wrapper with methods for sending/receiving data.
         """
-        return WebSocketCreateExecutorTerm(client=self.client)
+        return WebSocketCreateExecutorTerm(
+            recv_timeout=recv_timeout, ws_factory=ws_factory, client=self.client
+        )
 
 
 class AsyncExecutorAPI:
@@ -11973,6 +11994,8 @@ def modeling_commands_ws(
         video_res_height: Optional[int] = None,
         video_res_width: Optional[int] = None,
         webrtc: Optional[bool] = None,
+        recv_timeout: Optional[float] = None,
+        ws_factory: Optional[Callable[..., ClientConnectionSync]] = None,
     ) -> "WebSocketModelingCommandsWs":
         """Open a websocket which accepts modeling commands.
 
@@ -11989,6 +12012,8 @@ def modeling_commands_ws(
             video_res_height=video_res_height,
             video_res_width=video_res_width,
             webrtc=webrtc,
+            recv_timeout=recv_timeout,
+            ws_factory=ws_factory,
             client=self.client,
         )
 
@@ -12114,6 +12139,8 @@ def __init__(
         self,
         conversation_id: Optional[str] = None,
         replay: Optional[bool] = None,
+        recv_timeout: Optional[float] = None,
+        ws_factory: Optional[Callable[..., ClientConnectionSync]] = None,
         *,
         client: Client,
     ):
@@ -12134,12 +12161,18 @@ def __init__(
                 url = url + "?replay=" + str(replay).lower()
 
         headers = client.get_headers()
-        self.ws = ws_connect(
+        factory = ws_factory or ws_connect
+        self.ws = factory(
             url.replace("http", "ws"),
             additional_headers=headers,
             close_timeout=120,
             max_size=None,
         )
+        self._recv_timeout = (
+            client.get_websocket_recv_timeout()
+            if recv_timeout is None
+            else recv_timeout
+        )
 
     def __enter__(self):
         return self
@@ -12171,7 +12204,7 @@ def send_binary(self, data: MlCopilotClientMessage):
 
     def recv(self) -> MlCopilotServerMessage:
         """Receive data from the websocket."""
-        message = self.ws.recv(timeout=60)
+        message = self.ws.recv(timeout=self._recv_timeout)
 
         return MlCopilotServerMessage.model_validate_json(message)
 
@@ -12185,18 +12218,31 @@ class WebSocketMlReasoningWs:
 
     ws: ClientConnectionSync
 
-    def __init__(self, id: str, *, client: Client):
+    def __init__(
+        self,
+        id: str,
+        recv_timeout: Optional[float] = None,
+        ws_factory: Optional[Callable[..., ClientConnectionSync]] = None,
+        *,
+        client: Client,
+    ):
         # Inline WebSocket connection logic
 
         url = ("{}" + "/ws/ml/reasoning/{id}").format(client.base_url, id=id)
 
         headers = client.get_headers()
-        self.ws = ws_connect(
+        factory = ws_factory or ws_connect
+        self.ws = factory(
             url.replace("http", "ws"),
             additional_headers=headers,
             close_timeout=120,
             max_size=None,
         )
+        self._recv_timeout = (
+            client.get_websocket_recv_timeout()
+            if recv_timeout is None
+            else recv_timeout
+        )
 
     def __enter__(self):
         return self
@@ -12228,7 +12274,7 @@ def send_binary(self, data: MlCopilotClientMessage):
 
     def recv(self) -> MlCopilotServerMessage:
         """Receive data from the websocket."""
-        message = self.ws.recv(timeout=60)
+        message = self.ws.recv(timeout=self._recv_timeout)
 
         return MlCopilotServerMessage.model_validate_json(message)
 
@@ -12242,18 +12288,30 @@ class WebSocketCreateExecutorTerm:
 
     ws: ClientConnectionSync
 
-    def __init__(self, *, client: Client):
+    def __init__(
+        self,
+        recv_timeout: Optional[float] = None,
+        ws_factory: Optional[Callable[..., ClientConnectionSync]] = None,
+        *,
+        client: Client,
+    ):
         # Inline WebSocket connection logic
 
         url = ("{}" + "/ws/executor/term").format(client.base_url)
 
         headers = client.get_headers()
-        self.ws = ws_connect(
+        factory = ws_factory or ws_connect
+        self.ws = factory(
             url.replace("http", "ws"),
             additional_headers=headers,
             close_timeout=120,
             max_size=None,
         )
+        self._recv_timeout = (
+            client.get_websocket_recv_timeout()
+            if recv_timeout is None
+            else recv_timeout
+        )
 
     def __enter__(self):
         return self
@@ -12285,7 +12343,7 @@ def send_binary(self, data: Dict[str, Any]):
 
     def recv(self) -> Dict[str, Any]:
         """Receive data from the websocket."""
-        message = self.ws.recv(timeout=60)
+        message = self.ws.recv(timeout=self._recv_timeout)
 
         return json.loads(message)
 
@@ -12311,6 +12369,8 @@ def __init__(
         video_res_height: Optional[int] = None,
         video_res_width: Optional[int] = None,
         webrtc: Optional[bool] = None,
+        recv_timeout: Optional[float] = None,
+        ws_factory: Optional[Callable[..., ClientConnectionSync]] = None,
         *,
         client: Client,
     ):
@@ -12379,12 +12439,18 @@ def __init__(
                 url = url + "?webrtc=" + str(webrtc).lower()
 
         headers = client.get_headers()
-        self.ws = ws_connect(
+        factory = ws_factory or ws_connect
+        self.ws = factory(
             url.replace("http", "ws"),
             additional_headers=headers,
             close_timeout=120,
             max_size=None,
         )
+        self._recv_timeout = (
+            client.get_websocket_recv_timeout()
+            if recv_timeout is None
+            else recv_timeout
+        )
 
     def __enter__(self):
         return self
@@ -12416,7 +12482,7 @@ def send_binary(self, data: WebSocketRequest):
 
     def recv(self) -> WebSocketResponse:
         """Receive data from the websocket."""
-        message = self.ws.recv(timeout=60)
+        message = self.ws.recv(timeout=self._recv_timeout)
 
         return WebSocketResponse.model_validate_json(message)
 

kittycad/client.py

@@ -17,6 +17,7 @@ class Client:
     cookies: Dict[str, str] = attr.ib(factory=dict, kw_only=True)
     headers: Dict[str, str] = attr.ib(factory=dict, kw_only=True)
     timeout: float = attr.ib(120.0, kw_only=True)
+    websocket_recv_timeout: Optional[float] = attr.ib(60.0, kw_only=True)
     verify_ssl: Union[str, bool, ssl.SSLContext, truststore.SSLContext] = attr.ib(
         True, kw_only=True
     )
@@ -47,6 +48,13 @@ def with_timeout(self, timeout: float) -> "Client":
         """Get a new client matching this one with a new timeout (in seconds)"""
         return attr.evolve(self, timeout=timeout)
 
+    def get_websocket_recv_timeout(self) -> Optional[float]:
+        return self.websocket_recv_timeout
+
+    def with_websocket_recv_timeout(self, timeout: Optional[float]) -> "Client":
+        """Get a new client matching this one with a new websocket recv timeout"""
+        return attr.evolve(self, websocket_recv_timeout=timeout)
+
     def with_base_url(self, url: str) -> "Client":
         """Get a new client matching this one with a new base url"""
         return attr.evolve(self, base_url=url)
@@ -85,6 +93,7 @@ class AsyncClient:
     cookies: Dict[str, str] = attr.ib(factory=dict, kw_only=True)
     headers: Dict[str, str] = attr.ib(factory=dict, kw_only=True)
     timeout: float = attr.ib(120.0, kw_only=True)
+    websocket_recv_timeout: Optional[float] = attr.ib(60.0, kw_only=True)
     verify_ssl: Union[str, bool, ssl.SSLContext, truststore.SSLContext] = attr.ib(
         True, kw_only=True
     )
@@ -115,6 +124,13 @@ def with_timeout(self, timeout: float) -> "AsyncClient":
         """Get a new client matching this one with a new timeout (in seconds)"""
         return attr.evolve(self, timeout=timeout)
 
+    def get_websocket_recv_timeout(self) -> Optional[float]:
+        return self.websocket_recv_timeout
+
+    def with_websocket_recv_timeout(self, timeout: Optional[float]) -> "AsyncClient":
+        """Get a new client matching this one with a new websocket recv timeout"""
+        return attr.evolve(self, websocket_recv_timeout=timeout)
+
     def with_base_url(self, url: str) -> "AsyncClient":
         """Get a new client matching this one with a new base url"""
         return attr.evolve(self, base_url=url)