fmt

Author: jrvb-rl

src/runloop_api_client/lib/polling_async.py

@@ -61,7 +61,7 @@ async def async_poll_until(
 
 
 async def retry_server_poll_until(
-    retriever: Callable[[], Awaitable[T]],
+    retriever: Callable[[float], Awaitable[T]],
     is_terminal: Callable[[T], bool],
     timeout_seconds: float = 30.0,
     on_error: Optional[Callable[[Exception], T]] = None,
@@ -70,8 +70,8 @@ async def retry_server_poll_until(
     Retry a server-side long-poll until a condition is met or max timeout is reached.
 
     Args:
-        retriever: Async or sync callable that returns the object to check.  This takes should
-          take one argument, which is the remaing time to poll.q
+        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
@@ -83,12 +83,11 @@ async def retry_server_poll_until(
     Raises:
         PollingTimeout: When max attempts or timeout is reached
     """
-    start_time = time.time()
     last_result: Union[T, None] = None
+    start_time = time.time()
 
     while True:
-        elapsed = time.time() - start_time
-        remaining_time = timeout_seconds - elapsed
+        remaining_time = timeout_seconds - (time.time() - start_time)
         if remaining_time <= 0:
             raise PollingTimeout(f"Exceeded timeout of {timeout_seconds} seconds", last_result)
 

src/runloop_api_client/resources/devboxes/devboxes.py

@@ -2055,7 +2055,7 @@ async def await_running(
             RunloopError: If devbox enters a non-running terminal state
         """
 
-        async def wait_for_devbox_status(remaining_timeout_seconds) -> 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: