Merge pull request #613 from runloopai/alb/add-long-poll-to-await-completed

Add long poll endpoint for querying async exec status

Author: alb-rl

src/runloop_api_client/resources/devboxes/executions.py

@@ -17,6 +17,7 @@
     async_to_streamed_response_wrapper,
 )
 from ..._constants import DEFAULT_TIMEOUT
+from ..._exceptions import APIStatusError
 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
@@ -105,7 +106,7 @@ def await_completed(
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
     ) -> DevboxAsyncExecutionDetailView:
         """Wait for an execution to complete.
-        
+
         Args:
             execution_id: The ID of the execution to wait for
             id: The ID of the devbox
@@ -121,20 +122,37 @@ def await_completed(
         Raises:
             PollingTimeout: If polling times out before execution completes
         """
-        def retrieve_execution() -> DevboxAsyncExecutionDetailView:
-            return self.retrieve(
-                execution_id,
-                devbox_id=devbox_id,
-                extra_headers=extra_headers,
-                extra_query=extra_query,
-                extra_body=extra_body,
-                timeout=timeout
+
+        def wait_for_execution_status() -> DevboxAsyncExecutionDetailView:
+            # This wait_for_status endpoint polls the execution status for 60 seconds until it reaches either completed.
+            return self._post(
+                f"/v1/devboxes/{devbox_id}/executions/{execution_id}/wait_for_status",
+                body={"statuses": ["completed"]},
+                options=make_request_options(
+                    extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+                ),
+                cast_to=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:
+                # Return a placeholder result 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 error
+
         def is_done(execution: DevboxAsyncExecutionDetailView) -> bool:
-            return execution.status == 'completed'
+            return execution.status == "completed"
 
-        return poll_until(retrieve_execution, is_done, config)    
+        return poll_until(wait_for_execution_status, is_done, config, handle_timeout_error)
 
     def execute_async(
         self,
@@ -369,7 +387,7 @@ async def retrieve(
 
     async def await_completed(
         self,
-        execution_id: str,        
+        execution_id: str,
         *,
         devbox_id: str,
         polling_config: PollingConfig | None = None,
@@ -381,7 +399,7 @@ async def await_completed(
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
     ) -> DevboxAsyncExecutionDetailView:
         """Wait for an execution to complete.
-        
+
         Args:
             execution_id: The ID of the execution to wait for
             id: The ID of the devbox
@@ -397,20 +415,35 @@ async def await_completed(
         Raises:
             PollingTimeout: If polling times out before execution completes
         """
-        async def retrieve_execution() -> DevboxAsyncExecutionDetailView:
-            return await self.retrieve(
-                execution_id,
-                devbox_id=devbox_id,
-                extra_headers=extra_headers,
-                extra_query=extra_query,
-                extra_body=extra_body,
-                timeout=timeout,
-            )
+
+        async def wait_for_execution_status() -> DevboxAsyncExecutionDetailView:
+            try:
+                return await self._post(
+                    f"/v1/devboxes/{devbox_id}/executions/{execution_id}/wait_for_status",
+                    body={"statuses": ["completed"]},
+                    options=make_request_options(
+                        extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+                    ),
+                    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
 
         def is_done(execution: DevboxAsyncExecutionDetailView) -> bool:
-            return execution.status == 'completed'
+            return execution.status == "completed"
 
-        return await async_poll_until(retrieve_execution, is_done, polling_config)
+        return await async_poll_until(wait_for_execution_status, is_done, polling_config)
 
     async def execute_async(
         self,