Clean up error handling in long polling

Author: alb-rl

src/runloop_api_client/resources/devboxes/devboxes.py

@@ -88,7 +88,7 @@
     SyncDiskSnapshotsCursorIDPage,
     AsyncDiskSnapshotsCursorIDPage,
 )
-from ..._exceptions import RunloopError, APIStatusError
+from ..._exceptions import RunloopError, APIStatusError, APITimeoutError
 from ...lib.polling import PollingConfig, poll_until
 from ..._base_client import AsyncPaginator, make_request_options
 from .disk_snapshots import (
@@ -115,6 +115,18 @@
 DEVBOX_BOOTING_STATES = frozenset(("provisioning", "initializing"))
 
 
+def placeholder_devbox_view(id: str) -> DevboxView:
+    return DevboxView(
+        id=id,
+        status="provisioning",
+        capabilities=[],
+        create_time_ms=0,
+        launch_parameters=SharedLaunchParameters(),
+        metadata={},
+        state_transitions=[],
+    )
+
+
 class DevboxesResource(SyncAPIResource):
     @cached_property
     def disk_snapshots(self) -> DiskSnapshotsResource:
@@ -391,21 +403,15 @@ def wait_for_devbox_status() -> DevboxView:
             )
 
         def handle_timeout_error(error: Exception) -> DevboxView:
-            # Handle 408 timeout errors by returning current devbox state to continue polling
-            if isinstance(error, APIStatusError) and error.response.status_code == 408:
+            # Handle timeout errors by returning current devbox state to continue polling
+            if isinstance(error, APITimeoutError) or (
+                isinstance(error, APIStatusError) and error.response.status_code == 408
+            ):
                 # Return a placeholder result to continue polling
-                return DevboxView(
-                    id=id,
-                    status="provisioning",
-                    capabilities=[],
-                    create_time_ms=0,
-                    launch_parameters=SharedLaunchParameters(),
-                    metadata={},
-                    state_transitions=[],
-                )
-            else:
-                # Re-raise other errors to stop polling
-                raise error
+                return placeholder_devbox_view(id)
+
+            # Re-raise other errors to stop polling
+            raise error
 
         def is_done_booting(devbox: DevboxView) -> bool:
             return devbox.status not in DEVBOX_BOOTING_STATES
@@ -1686,21 +1692,14 @@ async def wait_for_devbox_status() -> DevboxView:
                     body={"statuses": ["running", "failure"]},
                     cast_to=DevboxView,
                 )
-            except APIStatusError as error:
-                if error.response.status_code == 408:
-                    # Handle 408 timeout errors by returning a placeholder result to continue polling
-                    return DevboxView(
-                        id=id,
-                        status="provisioning",
-                        capabilities=[],
-                        create_time_ms=0,
-                        launch_parameters=SharedLaunchParameters(),
-                        metadata={},
-                        state_transitions=[],
-                    )
-                else:
-                    # Re-raise other errors to stop polling
-                    raise
+            except (APITimeoutError, APIStatusError) as error:
+                # Handle timeout errors by returning current devbox state to continue polling
+                if isinstance(error, APITimeoutError) or error.response.status_code == 408:
+                    # Return a placeholder result to continue polling
+                    return placeholder_devbox_view(id)
+
+                # Re-raise other errors to stop polling
+                raise
 
         def is_done_booting(devbox: DevboxView) -> bool:
             return devbox.status not in DEVBOX_BOOTING_STATES

src/runloop_api_client/resources/devboxes/executions.py

@@ -17,7 +17,7 @@
     async_to_streamed_response_wrapper,
 )
 from ..._constants import DEFAULT_TIMEOUT
-from ..._exceptions import APIStatusError
+from ..._exceptions import APIStatusError, APITimeoutError
 from ...lib.polling import PollingConfig, poll_until
 from ..._base_client import make_request_options
 from ...types.devboxes import execution_retrieve_params, execution_execute_sync_params, execution_execute_async_params
@@ -28,6 +28,16 @@
 __all__ = ["ExecutionsResource", "AsyncExecutionsResource"]
 
 
+def placeholder_execution_detail_view(devbox_id: str, execution_id: str) -> DevboxAsyncExecutionDetailView:
+    return DevboxAsyncExecutionDetailView(
+        devbox_id=devbox_id,
+        execution_id=execution_id,
+        status="queued",
+        stdout="",
+        stderr="",
+    )
+
+
 class ExecutionsResource(SyncAPIResource):
     @cached_property
     def with_raw_response(self) -> ExecutionsResourceWithRawResponse:
@@ -127,16 +137,12 @@ def wait_for_execution_status() -> DevboxAsyncExecutionDetailView:
             )
 
         def handle_timeout_error(error: Exception) -> DevboxAsyncExecutionDetailView:
-            # Handle 408 timeout errors by returning current execution state to continue polling
-            if isinstance(error, APIStatusError) and error.response.status_code == 408:
+            # Handle timeout errors by returning current execution state to continue polling
+            if isinstance(error, APITimeoutError) or (
+                isinstance(error, APIStatusError) and error.response.status_code == 408
+            ):
                 # Return a placeholder result to continue polling
-                return DevboxAsyncExecutionDetailView(
-                    devbox_id=devbox_id,
-                    execution_id=execution_id,
-                    status="queued",
-                    stdout="",
-                    stderr="",
-                )
+                return placeholder_execution_detail_view(devbox_id, execution_id)
             else:
                 # Re-raise other errors to stop polling
                 raise error
@@ -410,19 +416,13 @@ async def wait_for_execution_status() -> DevboxAsyncExecutionDetailView:
                     body={"statuses": ["completed"]},
                     cast_to=DevboxAsyncExecutionDetailView,
                 )
-            except APIStatusError as error:
-                if error.response.status_code == 408:
-                    # Handle 408 timeout errors by returning current execution state to continue polling
-                    return DevboxAsyncExecutionDetailView(
-                        devbox_id=devbox_id,
-                        execution_id=execution_id,
-                        status="queued",
-                        stdout="",
-                        stderr="",
-                    )
-                else:
-                    # Re-raise other errors to stop polling
-                    raise
+            except (APITimeoutError, APIStatusError) as error:
+                # Handle timeout errors by returning placeholder to continue polling
+                if isinstance(error, APITimeoutError) or error.response.status_code == 408:
+                    return placeholder_execution_detail_view(devbox_id, execution_id)
+
+                # Re-raise other errors to stop polling
+                raise
 
         def is_done(execution: DevboxAsyncExecutionDetailView) -> bool:
             return execution.status == "completed"