fix: include response handling in retries and dispose clients in build to avoid reconnecst (#162)

Author: christianalfoni

src/API.ts

@@ -1,6 +1,5 @@
 import type { Client, Config } from "@hey-api/client-fetch";
-import { createClient, createConfig } from "@hey-api/client-fetch";
-import { handleResponse, retryWithDelay } from "./utils/api";
+import { handleResponse, retryWithDelay, createApiClient } from "./utils/api";
 import { getInferredBaseUrl } from "./utils/constants";
 import {
   metaInfo,
@@ -56,54 +55,6 @@ import type {
 } from "./api-clients/client";
 import { PitcherManagerResponse } from "./types";
 
-async function enhanceFetch(
-  request: Request,
-  instrumentation?: (request: Request) => Promise<Response>
-) {
-  // Clone the request to modify headers
-  const headers = new Headers(request.headers);
-  const existingUserAgent = headers.get("User-Agent") || "";
-
-  // Extend User-Agent with SDK version
-  headers.set(
-    "User-Agent",
-    `${existingUserAgent ? `${existingUserAgent} ` : ""}codesandbox-sdk/${
-      // @ts-expect-error - Replaced at build time
-      CSB_SDK_VERSION
-    }`.trim()
-  );
-
-  // Create new request with updated headers and optionally add instrumentation
-  return instrumentation
-    ? instrumentation(
-        new Request(request, {
-          headers,
-        })
-      )
-    : fetch(
-        new Request(request, {
-          headers,
-        })
-      );
-}
-
-function createApiClient(
-  apiKey: string,
-  config: Config = {},
-  instrumentation?: (request: Request) => Promise<Response>
-) {
-  return createClient(
-    createConfig({
-      baseUrl: config.baseUrl || getInferredBaseUrl(apiKey),
-      fetch: (request) => enhanceFetch(request, instrumentation),
-      ...config,
-      headers: {
-        Authorization: `Bearer ${apiKey}`,
-        ...config.headers,
-      },
-    })
-  );
-}
 
 export interface APIOptions {
   apiKey: string;
@@ -297,17 +248,18 @@ export class API {
   }
 
   async hibernate(id: string, data?: VmHibernateData["body"]) {
-    const response = await retryWithDelay(
-      () =>
-        vmHibernate({
+    return retryWithDelay(
+      async () => {
+        const response = await vmHibernate({
           client: this.client,
           path: { id },
           body: data,
-        }),
+        });
+        return handleResponse(response, `Failed to hibernate VM ${id}`);
+      },
       3,
       200
     );
-    return handleResponse(response, `Failed to hibernate VM ${id}`);
   }
 
   async updateHibernationTimeout(
@@ -335,17 +287,18 @@ export class API {
   }
 
   async shutdown(id: string, data?: VmShutdownData["body"]) {
-    const response = await retryWithDelay(
-      () =>
-        vmShutdown({
+    return retryWithDelay(
+      async () => {
+        const response = await vmShutdown({
           client: this.client,
           path: { id },
           body: data,
-        }),
+        });
+        return handleResponse(response, `Failed to shutdown VM ${id}`);
+      },
       3,
       200
     );
-    return handleResponse(response, `Failed to shutdown VM ${id}`);
   }
 
   async updateSpecs(id: string, data: VmUpdateSpecsData["body"]) {
@@ -359,20 +312,18 @@ export class API {
 
   async startVm(id: string, options?: StartVmOptions) {
     const { retryDelay = 200, ...data } = options || {};
-    const response = await retryWithDelay(
-      () =>
-        vmStart({
+    const handledResponse = await retryWithDelay(
+      async () => {
+        const response = await vmStart({
           client: this.client,
           path: { id },
           body: data,
-        }),
+        });
+        return handleResponse(response, `Failed to start VM ${id}`);
+      },
       3,
       retryDelay
     );
-    const handledResponse = handleResponse(
-      response,
-      `Failed to start VM ${id}`
-    );
 
     return {
       bootupType:

src/AgentClient/AgentConnection.ts

@@ -258,11 +258,11 @@ export class AgentConnection {
   }
 
   dispose(): void {
+    this.errorEmitter.dispose();
+    this.messageEmitter.dispose();
     this.connection.dispose();
     this.disposePendingMessages();
     this.pendingMessages.clear();
     this.notificationListeners = {};
-    this.errorEmitter.dispose();
-    this.messageEmitter.dispose();
   }
 }

src/bin/commands/build.ts

@@ -215,6 +215,7 @@ export const buildCommand: yargs.CommandModule<
       };
 
       const tasks = templateData.sandboxes.map(async ({ id }, index) => {
+        let currentSession: SandboxClient | null = null;
         try {
           spinner.start(updateSpinnerMessage(index, "Starting sandbox..."));
 
@@ -225,6 +226,7 @@ export const buildCommand: yargs.CommandModule<
           let sandboxVM = new Sandbox(id, api, startResponse);
 
           let session = await sandboxVM.connect();
+          currentSession = session;
 
           spinner.start(
             updateSpinnerMessage(index, "Writing files to sandbox...")
@@ -251,6 +253,10 @@ export const buildCommand: yargs.CommandModule<
             throw new Error(`Failed to write files to sandbox: ${error}`);
           });
 
+          // Dispose of the session after writing files to prevent reconnection
+          session.dispose();
+          currentSession = null;
+
           spinner.start(updateSpinnerMessage(index, "Building sandbox..."));
 
           sandboxVM = await withCustomError(
@@ -264,9 +270,14 @@ export const buildCommand: yargs.CommandModule<
             sandboxVM.connect(),
             "Failed to connect to sandbox after building"
           );
+          currentSession = session;
 
           await waitForSetup(session, index);
 
+          // Dispose of the session after setup to prevent reconnection
+          session.dispose();
+          currentSession = null;
+
           spinner.start(
             updateSpinnerMessage(index, "Optimizing initial state...")
           );
@@ -281,6 +292,7 @@ export const buildCommand: yargs.CommandModule<
             sandboxVM.connect(),
             "Failed to connect to sandbox after optimizing initial state"
           );
+          currentSession = session;
 
           await waitForSetup(session, index);
 
@@ -329,6 +341,10 @@ export const buildCommand: yargs.CommandModule<
             await new Promise((resolve) => setTimeout(resolve, 5000));
           }
 
+          // Dispose of the session after port operations to prevent reconnection
+          session.dispose();
+          currentSession = null;
+
           spinner.start(updateSpinnerMessage(index, "Creating snapshot..."));
           await withCustomError(
             sdk.sandboxes.hibernate(id),
@@ -338,6 +354,12 @@ export const buildCommand: yargs.CommandModule<
 
           return id;
         } catch (error) {
+          // Dispose of any active session to prevent reconnection attempts
+          if (currentSession) {
+            currentSession.dispose();
+            currentSession = null;
+          }
+
           spinner.start(
             updateSpinnerMessage(
               index,