[interactive]: Ensure that command server is healthy before connecting

.changeset/wide-ants-buy.md

@@ -0,0 +1,5 @@
+---
+"sandbox": patch
+---
+
+Ensure that the interactive server is available before creating websocket connection requests.

packages/sandbox/package.json

@@ -4,9 +4,9 @@
   "version": "4.1.0",
   "scripts": {
     "clean": "rm -rf node_modules dist",
-    "sandbox": "ts-node ./src/sandbox.ts",
+    "sandbox": "tsx ./src/sandbox.ts",
     "build": "tsdown",
-    "vercel-build": "ts-node ./scripts/create-preview-deployment-tgz.ts",
+    "vercel-build": "tsx ./scripts/create-preview-deployment-tgz.ts",
     "test": "vitest",
     "example-whoami": "ts-node ./src/example-whoami.ts",
     "example-next": "ts-node ./src/example-next.ts",
@@ -61,6 +61,7 @@
     "open": "^10.2.0",
     "ora": "^8.2.0",
     "tsdown": "catalog:",
+    "tsx": "4.20.3",
     "vitest": "catalog:",
     "xdg-app-paths": "5.1.0"
   }

packages/sandbox/src/interactive-shell/interactive-shell.ts

@@ -1,9 +1,14 @@
 import { Sandbox } from "@vercel/sandbox";
 import createDebugger from "debug";
+import retry from "async-retry";
 import { WebSocket } from "ws";
 import { printCommand } from "../util/print-command";
 import ora from "ora";
-import { acquireRelease, createAbortController, defer } from "../util/disposables";
+import {
+  acquireRelease,
+  createAbortController,
+  defer,
+} from "../util/disposables";
 import chalk from "chalk";
 import { extendSandboxTimeoutPeriodically } from "./extend-sandbox-timeout";
 
@@ -68,6 +73,7 @@ export async function startInteractiveShell(options: {
     (s) => s.stop(),
   );
 
+  debug("Opening interactive session...");
   progress.text = "Opening interactive session...";
   const { url, token } = await options.sandbox.openInteractive();
 
@@ -76,8 +82,47 @@ export async function startInteractiveShell(options: {
     ? ["sudo", command, ...args]
     : [command, ...args];
 
+  debug("Connecting...");
   progress.text = "Connecting...";
-  const client = new WebSocket(`${url}?token=${encodeURIComponent(token)}`);
+
+  let client: WebSocket;
+  try {
+    // Wait for the websocket server to be ready
+    // to accept traffic. This is important for
+    // cross region connections because the sevrer
+    // will have been instructed to start, but may not
+    // be fully ready to accept connections yet.
+    const { hostname } = new URL(url);
+    await retry(
+      async (_, attempt) => {
+        debug("Validating server health... (attempt %d)", attempt);
+        const response = await fetch(`https://${hostname}/health`);
+        if (!response.ok) {
+          throw new Error(`Server health check failed: ${response.statusText}`);
+        }
+      },
+      {
+        onRetry: (error: unknown, attempt: number) => {
+          debug(
+            "Server health check failed: %s (attempt %d)",
+            (error as Error).message,
+            attempt,
+          );
+        },
+        retries: 5,
+        minTimeout: 100,
+        randomize: true,
+      },
+    );
+
+    client = new WebSocket(`${url}?token=${encodeURIComponent(token)}`);
+  } catch (e) {
+    console.error(
+      chalk.dim(`\n╰▶ connection to ▲ ${options.sandbox.name} failed.`),
+    );
+    process.exitCode = 1;
+    return;
+  }
   using _client = defer(() => {
     try {
       client.close();
@@ -91,6 +136,7 @@ export async function startInteractiveShell(options: {
     client.once("error", (err) => reject(err));
   });
   debug("connected to %s", url);
+  progress.text = "Connected successfully.";
 
   client.send(
     JSON.stringify({
@@ -165,7 +211,9 @@ export async function startInteractiveShell(options: {
   process.removeListener("SIGWINCH", onResize);
   process.stdin.removeListener("data", onStdin);
 
-  console.error(chalk.dim(`\n╰▶ connection to ▲ ${options.sandbox.name} closed.`));
+  console.error(
+    chalk.dim(`\n╰▶ connection to ▲ ${options.sandbox.name} closed.`),
+  );
 }
 
 function toEnvArray(env: Record<string, string>): string[] {