From ac15e72dec30a1a2feb5c4419889f9eab149d303 Mon Sep 17 00:00:00 2001
From: Albert Li <albert@runloop.ai>
Date: Tue, 8 Jul 2025 14:15:43 -0700
Subject: [PATCH] Update awaitRunning to use long polling endpoint

---
 src/lib/polling.ts                 | 38 +++++++++++++++++++++++++-----
 src/resources/devboxes/devboxes.ts | 26 ++++++++++++++++----
 2 files changed, 54 insertions(+), 10 deletions(-)

diff --git a/src/lib/polling.ts b/src/lib/polling.ts
index 794758047..3e148b16b 100644
--- a/src/lib/polling.ts
+++ b/src/lib/polling.ts
@@ -1,3 +1,5 @@
+import { APIError } from '../error';
+
 export interface PollingOptions<T> {
   /** Initial delay before starting polling (in milliseconds) */
   initialDelayMs?: number;
@@ -14,6 +16,11 @@ export interface PollingOptions<T> {
   shouldStop?: (result: T) => boolean;
   /** Optional callback for each polling attempt */
   onPollingAttempt?: (attempt: number, result: T) => void;
+  /**
+   * Optional error handler for polling requests
+   * Return a result to continue polling with that result, or throw to stop polling
+   */
+  onError?: (error: APIError) => T;
 }
 
 const DEFAULT_OPTIONS: Partial<PollingOptions<any>> = {
@@ -59,10 +66,11 @@ export async function poll<T>(
   pollingRequest: () => Promise<T>,
   options: PollingOptions<T> = {},
 ): Promise<T> {
-  const { initialDelayMs, pollingIntervalMs, maxAttempts, timeoutMs, shouldStop, onPollingAttempt } = {
-    ...DEFAULT_OPTIONS,
-    ...options,
-  };
+  const { initialDelayMs, pollingIntervalMs, maxAttempts, timeoutMs, shouldStop, onPollingAttempt, onError } =
+    {
+      ...DEFAULT_OPTIONS,
+      ...options,
+    };
 
   // Start timeout timer if specified
   const timeoutPromise =
@@ -75,7 +83,16 @@ export async function poll<T>(
     : null;
 
   // Initial request
-  let result = await initialRequest();
+  let result: T;
+  try {
+    result = await initialRequest();
+  } catch (error) {
+    if (onError && error instanceof APIError) {
+      result = onError(error);
+    } else {
+      throw error;
+    }
+  }
 
   // Check if we should stop after initial request
   if (shouldStop?.(result)) {
@@ -92,7 +109,16 @@ export async function poll<T>(
 
     // Create polling promise
     const pollingPromise = async () => {
-      result = await pollingRequest();
+      try {
+        result = await pollingRequest();
+      } catch (error) {
+        if (onError && error instanceof APIError) {
+          result = onError(error);
+        } else {
+          throw error;
+        }
+      }
+
       onPollingAttempt?.(attempts, result);
 
       if (shouldStop?.(result)) {
diff --git a/src/resources/devboxes/devboxes.ts b/src/resources/devboxes/devboxes.ts
index cf5667bb1..a9e49119b 100644
--- a/src/resources/devboxes/devboxes.ts
+++ b/src/resources/devboxes/devboxes.ts
@@ -114,9 +114,9 @@ import {
   type DiskSnapshotsCursorIDPageParams,
 } from '../../pagination';
 import { type Response } from '../../_shims/index';
-import { PollingOptions, poll } from '@runloop/api-client/lib/polling';
-import { RunloopError } from '../..';
+import { poll, PollingOptions } from '@runloop/api-client/lib/polling';
 import { DevboxTools } from './tools';
+import { RunloopError } from '../..';
 
 type DevboxStatus = DevboxView['status'];
 const DEVBOX_BOOTING_STATES: DevboxStatus[] = ['provisioning', 'initializing'];
@@ -163,14 +163,32 @@ export class Devboxes extends APIResource {
     id: string,
     options?: Core.RequestOptions & { polling?: Partial<PollingOptions<DevboxView>> },
   ): Promise<DevboxView> {
+    const longPoll = (): Promise<DevboxView> => {
+      // This either returns a DevboxView when status is running or failure;
+      // Otherwise it throws an 408 error when times out.
+      return this._client.post(`/v1/devboxes/${id}/wait_for_status`, {
+        body: { statuses: ['running', 'failure'] },
+        ...options,
+      });
+    };
+
     const finalResult = await poll(
-      () => this.retrieve(id, options),
-      () => this.retrieve(id, options),
+      () => longPoll(),
+      () => longPoll(),
       {
         ...options?.polling,
         shouldStop: (result) => {
           return !DEVBOX_BOOTING_STATES.includes(result.status);
         },
+        onError: (error) => {
+          if (error.status === 408) {
+            // Return a placeholder result to continue polling
+            return { status: 'provisioning' } as DevboxView;
+          }
+
+          // For any other error, rethrow it
+          throw error;
+        },
       },
     );