Use new wait for status endpoint in await_running

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

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

src/runloop_api_client/resources/devboxes/devboxes.py

@@ -103,6 +103,7 @@
 from ...types.devbox_view import DevboxView
 from ...types.devbox_tunnel_view import DevboxTunnelView
 from ...types.devbox_snapshot_view import DevboxSnapshotView
+from ...types.shared.launch_parameters import LaunchParameters as SharedLaunchParameters
 from ...types.devbox_execution_detail_view import DevboxExecutionDetailView
 from ...types.devbox_create_ssh_key_response import DevboxCreateSSHKeyResponse
 from ...types.shared_params.launch_parameters import LaunchParameters
@@ -111,7 +112,8 @@
 
 __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,

src/runloop_api_client/resources/devboxes/executions.py

@@ -16,8 +16,8 @@
     async_to_raw_response_wrapper,
     async_to_streamed_response_wrapper,
 )
-from ...lib.polling import PollingConfig, poll_until
 from ..._constants import DEFAULT_TIMEOUT
+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
 from ...lib.polling_async import async_poll_until
@@ -103,7 +103,6 @@ def await_completed(
         extra_query: Query | None = None,
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
-        idempotency_key: str | None = None,
     ) -> DevboxAsyncExecutionDetailView:
         """Wait for an execution to complete.
 
@@ -380,7 +379,6 @@ async def await_completed(
         extra_query: Query | None = None,
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
-        idempotency_key: str | None = None,
     ) -> DevboxAsyncExecutionDetailView:
         """Wait for an execution to complete.
 
@@ -406,7 +404,7 @@ async def retrieve_execution() -> DevboxAsyncExecutionDetailView:
                 extra_headers=extra_headers,
                 extra_query=extra_query,
                 extra_body=extra_body,
-                timeout=timeout
+                timeout=timeout,
             )
 
         def is_done(execution: DevboxAsyncExecutionDetailView) -> bool: