Use new wait for status endpoint in await_running

alb-rl · alb-rl · commit bc470ba51c51 · 2025-07-08T16:34:49.000-07:00
diff --git a/src/runloop_api_client/lib/polling.py b/src/runloop_api_client/lib/polling.py
@@ -21,6 +21,7 @@ def poll_until(
     retriever: Callable[[], T],
     is_terminal: Callable[[T], bool],
     config: Optional[PollingConfig] = None,
+    on_error: Optional[Callable[[Exception], T]] = None,
 ) -> T:
     """
     Poll until a condition is met or timeout/max attempts are reached.
@@ -29,6 +30,8 @@ def poll_until(
         retriever: Callable that returns the object to check
         is_terminal: Callable that returns True when polling should stop
         config: Optional polling configuration
+        on_error: Optional error handler that can return a value to continue polling
+                 or re-raise the exception to stop polling
         
     Returns:
         The final state of the polled object
@@ -44,7 +47,13 @@ def poll_until(
     last_result = None
     
     while True:
-        last_result = retriever()
+        try:
+            last_result = retriever()
+        except Exception as e:
+            if on_error is not None:
+                last_result = on_error(e)
+            else:
+                raise
         
         if is_terminal(last_result):
             return last_result
diff --git a/src/runloop_api_client/lib/polling_async.py b/src/runloop_api_client/lib/polling_async.py
@@ -10,6 +10,7 @@ async def async_poll_until(
     retriever: Callable[[], Awaitable[T]],
     is_terminal: Callable[[T], bool],
     config: Optional[PollingConfig] = None,
+    on_error: Optional[Callable[[Exception], T]] = None,
 ) -> T:
     """
     Poll until a condition is met or timeout/max attempts are reached.
@@ -18,6 +19,8 @@ async def async_poll_until(
         retriever: Async or sync callable that returns the object to check
         is_terminal: Callable that returns True when polling should stop
         config: Optional polling configuration
+        on_error: Optional error handler that can return a value to continue polling
+                 or re-raise the exception to stop polling
         
     Returns:
         The final state of the polled object
@@ -33,7 +36,13 @@ async def async_poll_until(
     last_result: Union[T, None] = None
     
     while True:
-        last_result = await retriever()
+        try:
+            last_result = await retriever()
+        except Exception as e:
+            if on_error is not None:
+                last_result = on_error(e)
+            else:
+                raise
         
         if is_terminal(last_result):
             return last_result
diff --git a/src/runloop_api_client/resources/devboxes/devboxes.py b/src/runloop_api_client/resources/devboxes/devboxes.py
@@ -108,10 +108,12 @@
 from ...types.shared_params.launch_parameters import LaunchParameters
 from ...types.devbox_async_execution_detail_view import DevboxAsyncExecutionDetailView
 from ...types.shared_params.code_mount_parameters import CodeMountParameters
+from ...types.shared.launch_parameters import LaunchParameters as SharedLaunchParameters
 
 __all__ = ["DevboxesResource", "AsyncDevboxesResource"]
 
-DEVBOX_BOOTING_STATES = frozenset(('provisioning', 'initializing'))
+DEVBOX_BOOTING_STATES = frozenset(("provisioning", "initializing"))
+
 
 class DevboxesResource(SyncAPIResource):
     @cached_property
@@ -357,7 +359,7 @@ def update(
     def await_running(
         self,
         id: str,
-        *,        
+        *,
         polling_config: PollingConfig | None = None,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
@@ -367,7 +369,7 @@ def await_running(
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
     ) -> DevboxView:
         """Wait for a devbox to be in running state.
-        
+
         Args:
             id: The ID of the devbox to wait for
             config: Optional polling configuration
@@ -383,24 +385,43 @@ def await_running(
             PollingTimeout: If polling times out before devbox is running
             RunloopError: If devbox enters a non-running terminal state
         """
-        def retrieve_devbox() -> DevboxView:
-            return self.retrieve(
-                id,
-                extra_headers=extra_headers,
-                extra_query=extra_query,
-                extra_body=extra_body,
-                timeout=timeout
+
+        def wait_for_devbox_status() -> DevboxView:
+            # This wait_for_status endpoint polls the devbox status for 10 seconds until it reaches either running or failure.
+            # If it's neither, it will throw an error.
+            return self._post(
+                f"/v1/devboxes/{id}/wait_for_status",
+                body={"statuses": ["running", "failure"]},
+                options=make_request_options(
+                    extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+                ),
+                cast_to=DevboxView,
             )
 
+        def handle_timeout_error(error: Exception) -> DevboxView:
+            # Handle 408 timeout errors by returning current devbox state to continue polling
+            if isinstance(error, httpx.HTTPStatusError) 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
+
         def is_done_booting(devbox: DevboxView) -> bool:
             return devbox.status not in DEVBOX_BOOTING_STATES
 
-        devbox = poll_until(retrieve_devbox, is_done_booting, polling_config)
+        devbox = poll_until(wait_for_devbox_status, is_done_booting, polling_config, handle_timeout_error)
 
         if devbox.status != "running":
-            raise RunloopError(
-                f"Devbox entered non-running terminal state: {devbox.status}"
-            )
+            raise RunloopError(f"Devbox entered non-running terminal state: {devbox.status}")
 
         return devbox
 
@@ -427,7 +448,7 @@ def create_and_await_running(
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
     ) -> DevboxView:
         """Create a new devbox and wait for it to be in running state.
-        
+
         Args:
             blueprint_id: The ID of the blueprint to use
             blueprint_name: The name of the blueprint to use
@@ -1660,7 +1681,7 @@ async def await_running(
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
     ) -> DevboxView:
         """Wait for a devbox to be in running state.
-        
+
         Args:
             id: The ID of the devbox to wait for
             config: Optional polling configuration
@@ -1676,27 +1697,45 @@ async def await_running(
             PollingTimeout: If polling times out before devbox is running
             RunloopError: If devbox enters a non-running terminal state
         """
-        async def retrieve_devbox() -> DevboxView:
-            return await self.retrieve(
-                id,
-                extra_headers=extra_headers,
-                extra_query=extra_query,
-                extra_body=extra_body,
-                timeout=timeout
-            )
-        
+
+        async def wait_for_devbox_status() -> DevboxView:
+            # This wait_for_status endpoint polls the devbox status for 10 seconds until it reaches either running or failure.
+            # If it's neither, it will throw an error.
+            try:
+                return await self._post(
+                    f"/v1/devboxes/{id}/wait_for_status",
+                    body={"statuses": ["running", "failure"]},
+                    options=make_request_options(
+                        extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+                    ),
+                    cast_to=DevboxView,
+                )
+            except httpx.HTTPStatusError 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
+
         def is_done_booting(devbox: DevboxView) -> bool:
             return devbox.status not in DEVBOX_BOOTING_STATES
 
-        devbox = await async_poll_until(retrieve_devbox, is_done_booting, polling_config)
+        devbox = await async_poll_until(wait_for_devbox_status, is_done_booting, polling_config)
 
         if devbox.status != "running":
-            raise RunloopError(
-                f"Devbox entered non-running terminal state: {devbox.status}"
-            )        
+            raise RunloopError(f"Devbox entered non-running terminal state: {devbox.status}")
 
         return devbox
-    
+
     async def update(
         self,
         id: str,
