From a863655e38f72196d531dca097af0d9c50d40530 Mon Sep 17 00:00:00 2001
From: Albert Li <albert@runloop.ai>
Date: Wed, 9 Jul 2025 17:45:25 -0700
Subject: [PATCH] Update awaitCompleted async exec command to long poll

---
 src/resources/devboxes/devboxes.ts   |  7 ++++++-
 src/resources/devboxes/executions.ts | 25 +++++++++++++++++++++++--
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/src/resources/devboxes/devboxes.ts b/src/resources/devboxes/devboxes.ts
index a9e49119b..239e9d85f 100644
--- a/src/resources/devboxes/devboxes.ts
+++ b/src/resources/devboxes/devboxes.ts
@@ -158,6 +158,9 @@ export class Devboxes extends APIResource {
   /**
    * Wait for a devbox to reach the running state.
    * Polls the devbox status until it reaches running state or fails with an error.
+   *
+   * @param id - Devbox ID
+   * @param options - request options to specify retries, timeout, polling, etc.
    */
   async awaitRunning(
     id: string,
@@ -168,7 +171,6 @@ export class Devboxes extends APIResource {
       // Otherwise it throws an 408 error when times out.
       return this._client.post(`/v1/devboxes/${id}/wait_for_status`, {
         body: { statuses: ['running', 'failure'] },
-        ...options,
       });
     };
 
@@ -202,6 +204,9 @@ export class Devboxes extends APIResource {
   /**
    * Create a devbox and wait for it to reach the running state.
    * This is a convenience method that combines create() and awaitDevboxRunning().
+   *
+   * @param body - DevboxCreateParams
+   * @param options - request options to specify retries, timeout, polling, etc.
    */
   async createAndAwaitRunning(
     body?: DevboxCreateParams,
diff --git a/src/resources/devboxes/executions.ts b/src/resources/devboxes/executions.ts
index bddef4b45..7d436510d 100755
--- a/src/resources/devboxes/executions.ts
+++ b/src/resources/devboxes/executions.ts
@@ -49,6 +49,10 @@ export class Executions extends APIResource {
   /**
    * Wait for an async execution to complete.
    * Polls the execution status until it reaches completed state.
+   *
+   * @param id - Devbox ID
+   * @param executionId - Execution ID
+   * @param options - request options to specify retries, timeout, polling, etc.
    */
   async awaitCompleted(
     id: string,
@@ -57,14 +61,31 @@ export class Executions extends APIResource {
       polling?: Partial<PollingOptions<DevboxesAPI.DevboxAsyncExecutionDetailView>>;
     },
   ): Promise<DevboxesAPI.DevboxAsyncExecutionDetailView> {
+    const longPoll = (): Promise<DevboxesAPI.DevboxAsyncExecutionDetailView> => {
+      // This either returns a DevboxAsyncExecutionDetailView when execution status is completed;
+      // Otherwise it throws an 408 error when times out.
+      return this._client.post(`/v1/devboxes/${id}/executions/${executionId}/wait_for_status`, {
+        body: { statuses: ['completed'] },
+      });
+    };
+
     const finalResult = await poll(
-      () => this.retrieve(id, executionId, options),
-      () => this.retrieve(id, executionId, options),
+      () => longPoll(),
+      () => longPoll(),
       {
         ...options?.polling,
         shouldStop: (result: DevboxesAPI.DevboxAsyncExecutionDetailView) => {
           return result.status === 'completed';
         },
+        onError: (error) => {
+          if (error.status === 408) {
+            // Return a placeholder result to continue polling
+            return { status: 'running' } as DevboxesAPI.DevboxAsyncExecutionDetailView;
+          }
+
+          // For any other error, rethrow it
+          throw error;
+        },
       },
     );