Merge pull request #616 from runloopai/alb/long-poll-updates

Remove unnecessary arguments for await_running, await_completed, create_and_await_running

Author: alb-rl

src/runloop_api_client/resources/devboxes/devboxes.py

@@ -88,7 +88,7 @@
     SyncDiskSnapshotsCursorIDPage,
     AsyncDiskSnapshotsCursorIDPage,
 )
-from ..._exceptions import RunloopError, APIStatusError
+from ..._exceptions import RunloopError, APIStatusError, APITimeoutError
 from ...lib.polling import PollingConfig, poll_until
 from ..._base_client import AsyncPaginator, make_request_options
 from .disk_snapshots import (
@@ -115,6 +115,18 @@
 DEVBOX_BOOTING_STATES = frozenset(("provisioning", "initializing"))
 
 
+def placeholder_devbox_view(id: str) -> DevboxView:
+    return DevboxView(
+        id=id,
+        status="provisioning",
+        capabilities=[],
+        create_time_ms=0,
+        launch_parameters=SharedLaunchParameters(),
+        metadata={},
+        state_transitions=[],
+    )
+
+
 class DevboxesResource(SyncAPIResource):
     @cached_property
     def disk_snapshots(self) -> DiskSnapshotsResource:
@@ -360,13 +372,8 @@ def await_running(
         self,
         id: str,
         *,
+        # Use polling_config to configure the "long" polling behavior.
         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.
-        extra_headers: Headers | None = None,
-        extra_query: Query | None = None,
-        extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
     ) -> DevboxView:
         """Wait for a devbox to be in running state.
 
@@ -392,28 +399,19 @@ def wait_for_devbox_status() -> DevboxView:
             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, APIStatusError) and error.response.status_code == 408:
+            # Handle timeout errors by returning current devbox state to continue polling
+            if isinstance(error, APITimeoutError) or (
+                isinstance(error, APIStatusError) 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
+                return placeholder_devbox_view(id)
+
+            # Re-raise other errors to stop polling
+            raise error
 
         def is_done_booting(devbox: DevboxView) -> bool:
             return devbox.status not in DEVBOX_BOOTING_STATES
@@ -495,10 +493,6 @@ def create_and_await_running(
         return self.await_running(
             devbox.id,
             polling_config=polling_config,
-            extra_headers=extra_headers,
-            extra_query=extra_query,
-            extra_body=extra_body,
-            timeout=timeout,
         )
 
     def list(
@@ -1662,23 +1656,14 @@ async def create_and_await_running(
         return await self.await_running(
             devbox.id,
             polling_config=polling_config,
-            extra_headers=extra_headers,
-            extra_query=extra_query,
-            extra_body=extra_body,
-            timeout=timeout,
         )
 
     async def await_running(
         self,
         id: str,
         *,
+        # Use polling_config to configure the "long" polling behavior.
         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.
-        extra_headers: Headers | None = None,
-        extra_query: Query | None = None,
-        extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
     ) -> DevboxView:
         """Wait for a devbox to be in running state.
 
@@ -1705,26 +1690,16 @@ async def wait_for_devbox_status() -> DevboxView:
                 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 APIStatusError 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
+            except (APITimeoutError, APIStatusError) as error:
+                # Handle timeout errors by returning current devbox state to continue polling
+                if isinstance(error, APITimeoutError) or error.response.status_code == 408:
+                    # Return a placeholder result to continue polling
+                    return placeholder_devbox_view(id)
+
+                # Re-raise other errors to stop polling
+                raise
 
         def is_done_booting(devbox: DevboxView) -> bool:
             return devbox.status not in DEVBOX_BOOTING_STATES

src/runloop_api_client/resources/devboxes/executions.py

@@ -17,7 +17,7 @@
     async_to_streamed_response_wrapper,
 )
 from ..._constants import DEFAULT_TIMEOUT
-from ..._exceptions import APIStatusError
+from ..._exceptions import APIStatusError, APITimeoutError
 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
@@ -28,6 +28,16 @@
 __all__ = ["ExecutionsResource", "AsyncExecutionsResource"]
 
 
+def placeholder_execution_detail_view(devbox_id: str, execution_id: str) -> DevboxAsyncExecutionDetailView:
+    return DevboxAsyncExecutionDetailView(
+        devbox_id=devbox_id,
+        execution_id=execution_id,
+        status="queued",
+        stdout="",
+        stderr="",
+    )
+
+
 class ExecutionsResource(SyncAPIResource):
     @cached_property
     def with_raw_response(self) -> ExecutionsResourceWithRawResponse:
@@ -97,13 +107,8 @@ def await_completed(
         execution_id: str,
         devbox_id: str,
         *,
-        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.
-        extra_headers: Headers | None = None,
-        extra_query: Query | None = None,
-        extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        # Use polling_config to configure the "long" polling behavior.
+        polling_config: PollingConfig | None = None,
     ) -> DevboxAsyncExecutionDetailView:
         """Wait for an execution to complete.
 
@@ -128,31 +133,24 @@ def wait_for_execution_status() -> DevboxAsyncExecutionDetailView:
             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:
+            # Handle timeout errors by returning current execution state to continue polling
+            if isinstance(error, APITimeoutError) or (
+                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="",
-                )
+                return placeholder_execution_detail_view(devbox_id, execution_id)
             else:
                 # Re-raise other errors to stop polling
                 raise error
 
         def is_done(execution: DevboxAsyncExecutionDetailView) -> bool:
             return execution.status == "completed"
 
-        return poll_until(wait_for_execution_status, is_done, config, handle_timeout_error)
+        return poll_until(wait_for_execution_status, is_done, polling_config, handle_timeout_error)
 
     def execute_async(
         self,
@@ -390,13 +388,8 @@ async def await_completed(
         execution_id: str,
         *,
         devbox_id: str,
+        # Use polling_config to configure the "long" polling behavior.
         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.
-        extra_headers: Headers | None = None,
-        extra_query: Query | None = None,
-        extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
     ) -> DevboxAsyncExecutionDetailView:
         """Wait for an execution to complete.
 
@@ -421,24 +414,15 @@ async def wait_for_execution_status() -> DevboxAsyncExecutionDetailView:
                 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
+            except (APITimeoutError, APIStatusError) as error:
+                # Handle timeout errors by returning placeholder to continue polling
+                if isinstance(error, APITimeoutError) or error.response.status_code == 408:
+                    return placeholder_execution_detail_view(devbox_id, execution_id)
+
+                # Re-raise other errors to stop polling
+                raise
 
         def is_done(execution: DevboxAsyncExecutionDetailView) -> bool:
             return execution.status == "completed"

src/runloop_api_client/resources/scenarios/scenarios.py

@@ -471,10 +471,6 @@ def start_run_and_await_env_ready(
         self._client.devboxes.await_running(
             run.devbox_id,
             polling_config=polling_config,
-            extra_headers=extra_headers,
-            extra_query=extra_query,
-            extra_body=extra_body,
-            timeout=timeout,
         )
 
         return run

tests/api_resources/devboxes/test_executions.py

@@ -11,7 +11,7 @@
 from tests.utils import assert_matches_type
 from runloop_api_client import Runloop, AsyncRunloop
 from runloop_api_client.types import DevboxExecutionDetailView, DevboxAsyncExecutionDetailView
-from runloop_api_client._exceptions import APIStatusError
+from runloop_api_client._exceptions import APIStatusError, APITimeoutError
 from runloop_api_client.lib.polling import PollingConfig, PollingTimeout
 
 base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010")
@@ -338,7 +338,7 @@ def test_method_await_completed_with_config(self, client: Runloop) -> None:
         with patch.object(client.devboxes.executions, "_post") as mock_post:
             mock_post.return_value = mock_execution_completed
 
-            result = client.devboxes.executions.await_completed("execution_id", "devbox_id", config=config)
+            result = client.devboxes.executions.await_completed("execution_id", "devbox_id", polling_config=config)
 
             assert result.execution_id == "execution_id"
             assert result.status == "completed"
@@ -361,7 +361,7 @@ def test_method_await_completed_polling_timeout(self, client: Runloop) -> None:
             mock_post.return_value = mock_execution_running
 
             with pytest.raises(PollingTimeout):
-                client.devboxes.executions.await_completed("execution_id", "devbox_id", config=config)
+                client.devboxes.executions.await_completed("execution_id", "devbox_id", polling_config=config)
 
     @parametrize
     def test_method_await_completed_various_statuses(self, client: Runloop) -> None:
@@ -665,7 +665,7 @@ async def test_method_await_completed_timeout_handling(self, async_client: Async
         # Create a mock 408 response
         mock_response = Mock()
         mock_response.status_code = 408
-        mock_408_error = APIStatusError("Request timeout", response=mock_response, body=None)
+        mock_408_error = APITimeoutError(request=mock_response.request)
 
         mock_execution_completed = DevboxAsyncExecutionDetailView(
             devbox_id="devbox_id",

tests/api_resources/test_devboxes.py

@@ -4,7 +4,7 @@
 
 import os
 from typing import Any, cast
-from unittest.mock import ANY, Mock, patch
+from unittest.mock import Mock, patch
 
 import httpx
 import pytest
@@ -1171,9 +1171,7 @@ def test_method_create_and_await_running_success(self, client: Runloop) -> None:
                 assert result.id == "test_id"
                 assert result.status == "running"
                 mock_create.assert_called_once()
-                mock_await.assert_called_once_with(
-                    "test_id", polling_config=None, extra_headers=None, extra_query=None, extra_body=None, timeout=ANY
-                )
+                mock_await.assert_called_once_with("test_id", polling_config=None)
 
     @parametrize
     def test_method_create_and_await_running_with_config(self, client: Runloop) -> None:
@@ -1210,9 +1208,7 @@ def test_method_create_and_await_running_with_config(self, client: Runloop) -> N
 
                 assert result.id == "test_id"
                 assert result.status == "running"
-                mock_await.assert_called_once_with(
-                    "test_id", polling_config=config, extra_headers=None, extra_query=None, extra_body=None, timeout=ANY
-                )
+                mock_await.assert_called_once_with("test_id", polling_config=config)
 
     @parametrize
     def test_method_create_and_await_running_create_failure(self, client: Runloop) -> None: