diff --git a/src/runloop_api_client/resources/devboxes/executions.py b/src/runloop_api_client/resources/devboxes/executions.py index dbb262659..0481039ed 100755 --- a/src/runloop_api_client/resources/devboxes/executions.py +++ b/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,