diff --git a/src/runloop_api_client/lib/polling.py b/src/runloop_api_client/lib/polling.py index f9a8953f2..8031fd92b 100644 --- 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 index f3acdc93d..1f52dae07 100644 --- 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 index fbff382d6..3ed478e1f 100644 --- a/src/runloop_api_client/resources/devboxes/devboxes.py +++ b/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, diff --git a/src/runloop_api_client/resources/devboxes/executions.py b/src/runloop_api_client/resources/devboxes/executions.py index ad7df57e5..dbb262659 100755 --- a/src/runloop_api_client/resources/devboxes/executions.py +++ b/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: