Improve devbox polling: eliminate client-side sleep, enable HTTP/2 (#793)

Author: jrvb-rl

pyproject.toml

@@ -9,7 +9,7 @@ authors = [
 ]
 
 dependencies = [
-  "httpx>=0.23.0, <1",
+  "httpx[http2]>=0.23.0, <1",
   "pydantic>=2.0, <3",
   "typing-extensions>=4.14, <5",
   "anyio>=3.5.0, <5",

src/runloop_api_client/_base_client.py

@@ -1375,6 +1375,7 @@ def __init__(self, **kwargs: Any) -> None:
         kwargs.setdefault("timeout", DEFAULT_TIMEOUT)
         kwargs.setdefault("limits", DEFAULT_CONNECTION_LIMITS)
         kwargs.setdefault("follow_redirects", True)
+        kwargs.setdefault("http2", True)
         super().__init__(**kwargs)
 
 

src/runloop_api_client/lib/polling_async.py

@@ -58,3 +58,46 @@ async def async_poll_until(
                 raise PollingTimeout(f"Exceeded timeout of {config.timeout_seconds} seconds", last_result)
 
         await asyncio.sleep(config.interval_seconds)
+
+
+async def retry_server_poll_until(
+    retriever: Callable[[float], Awaitable[T]],
+    is_terminal: Callable[[T], bool],
+    timeout_seconds: float = 30.0,
+    on_error: Optional[Callable[[Exception], T]] = None,
+) -> T:
+    """
+    Retry a server-side long-poll until a condition is met or max timeout is reached.
+
+    Args:
+        retriever: Async callable that takes the remaining timeout (seconds) and
+            returns the object to check.
+        is_terminal: Callable that returns True when polling should stop
+        timeout_seconds: Total time to wait.  Must be > 0
+        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
+
+    Raises:
+        PollingTimeout: When max attempts or timeout is reached
+    """
+    last_result: Union[T, None] = None
+    start_time = time.time()
+
+    while True:
+        remaining_time = timeout_seconds - (time.time() - start_time)
+        if remaining_time <= 0:
+            raise PollingTimeout(f"Exceeded timeout of {timeout_seconds} seconds", last_result)
+
+        try:
+            last_result = await retriever(remaining_time)
+        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

@@ -82,7 +82,7 @@
     DiskSnapshotsResourceWithStreamingResponse,
     AsyncDiskSnapshotsResourceWithStreamingResponse,
 )
-from ...lib.polling_async import async_poll_until
+from ...lib.polling_async import async_poll_until, retry_server_poll_until
 from ...types.devbox_view import DevboxView
 from ...types.tunnel_view import TunnelView
 from ...types.shared_params.mount import Mount
@@ -2042,11 +2042,10 @@ async def await_running(
 
         Args:
             id: The ID of the devbox to wait for
-            config: Optional polling configuration
+            polling_config: Optional polling configuration
             extra_headers: Send extra headers
             extra_query: Add additional query parameters to the request
             extra_body: Add additional JSON properties to the request
-            timeout: Override the client-level default timeout for this request, in seconds
 
         Returns:
             The devbox in running state
@@ -2056,13 +2055,13 @@ async def await_running(
             RunloopError: If devbox enters a non-running terminal state
         """
 
-        async def wait_for_devbox_status() -> DevboxView:
+        async def wait_for_devbox_status(remaining_timeout_seconds: float) -> 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", "shutdown"]},
+                    body={"statuses": ["running", "failure", "shutdown"], "timeout_seconds": remaining_timeout_seconds},
                     cast_to=DevboxView,
                 )
             except (APITimeoutError, APIStatusError) as error:
@@ -2077,7 +2076,19 @@ async def wait_for_devbox_status() -> DevboxView:
         def is_done_booting(devbox: DevboxView) -> bool:
             return devbox.status not in DEVBOX_BOOTING_STATES
 
-        devbox = await async_poll_until(wait_for_devbox_status, is_done_booting, polling_config)
+        # calculate the timeout to use.  The PollingConfig doesn't
+        # match the semantics for server-side polling well, so we
+        # instead convert interval*attempts to a total time, and take
+        # the minimum total.
+        config = polling_config
+        if not config:
+            config = PollingConfig()  # use defaults
+
+        timeout = config.interval_seconds * config.max_attempts
+        if config.timeout_seconds is not None and config.timeout_seconds > 0:
+            timeout = min(config.timeout_seconds, timeout)
+
+        devbox = await retry_server_poll_until(wait_for_devbox_status, is_done_booting, timeout)
 
         if devbox.status != "running":
             raise RunloopError(f"Devbox entered non-running terminal state: {devbox.status}")