PR Review

Author: tode-rl

README.md

@@ -69,8 +69,7 @@ The URL must be of the form `https://api.<domain>`. The CLI derives other servic
 | API      | `https://api.<domain>` (the value of `RUNLOOP_BASE_URL`) |
 | Platform | `https://platform.<domain>` |
 | SSH      | `ssh.<domain>:443` |
-| PTY      | `https://pty.<domain>` |
-| Tunnels  | `tunnel.<domain>` |
+| Tunnels  | `tunnel.<domain>` (PTY sessions reach the devbox over a tunnel created via the API) |
 
 ## Usage
 

src/commands/devbox/pty.ts

@@ -116,41 +116,21 @@ async function execCommand(
     authToken,
   });
   const ws = await openPtyWebSocket(wsUrl, authToken);
-  await refreshPtySessionAfterAttach(
-    ws,
-    baseUrl,
-    sessionName,
-    80,
-    24,
-    authToken,
-  );
-  ws.send(command + "\n");
-
-  return new Promise<void>((resolve, reject) => {
-    const releaseOnce = createPtySessionReleaser(
-      baseUrl,
-      sessionName,
-      authToken,
-    );
-    const disposeInterruptSignals = registerPtyInterruptHandlers(
-      ws,
-      releaseOnce,
-    );
-
-    let timeoutId: ReturnType<typeof setTimeout> | undefined;
-    if (PTY_EXEC_TIMEOUT_MS > 0) {
-      timeoutId = setTimeout(() => {
-        releaseOnce();
-        ws.close();
-      }, PTY_EXEC_TIMEOUT_MS);
-    }
 
+  const releaseOnce = createPtySessionReleaser(baseUrl, sessionName, authToken);
+  const disposeInterruptSignals = registerPtyInterruptHandlers(ws, releaseOnce);
+
+  let timeoutId: ReturnType<typeof setTimeout> | undefined;
+
+  const completion = new Promise<void>((resolve, reject) => {
     const finish = () => {
       if (timeoutId !== undefined) clearTimeout(timeoutId);
       releaseOnce();
       disposeInterruptSignals();
     };
 
+    // Attach listeners synchronously *before* any await — otherwise messages
+    // emitted between ws open and listener registration are dropped.
     ws.on("message", (data: WebSocket.RawData) => {
       writePtyStreamToStdout(data);
     });
@@ -164,7 +144,29 @@ async function execCommand(
       finish();
       reject(err);
     });
+
+    if (PTY_EXEC_TIMEOUT_MS > 0) {
+      timeoutId = setTimeout(() => {
+        releaseOnce();
+        ws.close();
+      }, PTY_EXEC_TIMEOUT_MS);
+    }
   });
+
+  await refreshPtySessionAfterAttach(
+    ws,
+    baseUrl,
+    sessionName,
+    80,
+    24,
+    authToken,
+  );
+  // Send the command followed by `exit` so the shell terminates and the
+  // server closes the WebSocket. Without this, the session would stay open
+  // after the command finished and the CLI would hang.
+  ws.send(command + "\nexit\n");
+
+  return completion;
 }
 
 async function interactiveSession(
@@ -181,23 +183,26 @@ async function interactiveSession(
     authToken,
   });
   const ws = await openPtyWebSocket(wsUrl, authToken);
-  await refreshPtySessionAfterAttach(
+
+  // Attach IO listeners before the refresh round-trip so server output
+  // emitted during the ptyControl HTTP call is not dropped.
+  const releaseOnce = createPtySessionReleaser(baseUrl, sessionName, authToken);
+  const { dispose, done } = startPtyIoSession(
     ws,
     baseUrl,
     sessionName,
-    cols,
-    rows,
     authToken,
   );
+  const disposeSignals = registerPtyInterruptHandlers(ws, releaseOnce);
 
-  const releaseOnce = createPtySessionReleaser(baseUrl, sessionName, authToken);
-  const { dispose, done } = startPtyIoSession(
+  await refreshPtySessionAfterAttach(
     ws,
     baseUrl,
     sessionName,
+    cols,
+    rows,
     authToken,
   );
-  const disposeSignals = registerPtyInterruptHandlers(ws, releaseOnce);
 
   try {
     await done;

src/components/InteractivePty.tsx

@@ -5,6 +5,7 @@ import {
   resolvePtyWebSocketUrl,
   refreshPtySessionAfterAttach,
   startPtyIoSession,
+  PTY_NORMAL_CLOSE_CODE,
 } from "../lib/pty-client.js";
 import { openPtyWebSocket } from "../lib/pty-ws.js";
 import { clearScreen } from "../utils/screen.js";
@@ -92,15 +93,8 @@ export const InteractivePty: React.FC<InteractivePtyProps> = ({
 
         wsRef.current = ws;
 
-        await refreshPtySessionAfterAttach(
-          ws,
-          baseUrl,
-          sessionName,
-          cols,
-          rows,
-          authToken,
-        );
-
+        // Attach IO listeners before the refresh round-trip so server output
+        // emitted during the ptyControl HTTP call is not dropped.
         const releaseServerSession = createPtySessionReleaser(
           baseUrl,
           sessionName,
@@ -124,13 +118,22 @@ export const InteractivePty: React.FC<InteractivePtyProps> = ({
           return;
         }
 
+        await refreshPtySessionAfterAttach(
+          ws,
+          baseUrl,
+          sessionName,
+          cols,
+          rows,
+          authToken,
+        );
+
         done
           .then((code) => {
             wsRef.current = null;
             ioCleanup();
             restoreTerminal();
             hasStartedRef.current = false;
-            onExitRef.current?.(code === 4000 ? 0 : code);
+            onExitRef.current?.(code === PTY_NORMAL_CLOSE_CODE ? 0 : code);
           })
           .catch((err: Error) => {
             wsRef.current = null;

src/lib/pty-client.ts

@@ -240,6 +240,13 @@ export function createPtySessionReleaser(
 /** Same numeric value as `WebSocket.OPEN` from the `ws` package. */
 const WS_READY_STATE_OPEN = 1;
 
+/**
+ * Application-defined WebSocket close code (range 4000-4999 is reserved for
+ * applications) the PTY server uses to signal a clean session end — e.g. the
+ * user typed `exit` and the shell terminated normally. Treated as exit 0.
+ */
+export const PTY_NORMAL_CLOSE_CODE = 4000;
+
 /**
  * After the attach WebSocket is open: re-send terminal size (refreshes session geometry)
  * and send CR so the shell redraws the prompt (avoids a blank display until the user hits Enter).

src/lib/pty-ws.ts

@@ -6,6 +6,25 @@ const PTY_WS_MAX_ATTEMPTS = Math.max(
   parseInt(process.env.RUNLOOP_PTY_WS_RETRIES || "3", 10) || 3,
 );
 
+/**
+ * Per-attempt connect timeout. Kept short so the worst-case wall-clock spent
+ * stacking ptyConnect (up to 3 × 10s back-off) and WS attach retries stays
+ * bounded; override via env if a slow tunnel ever needs it.
+ */
+const PTY_WS_CONNECT_TIMEOUT_MS = (() => {
+  const raw = parseInt(
+    process.env.RUNLOOP_PTY_WS_CONNECT_TIMEOUT_MS || "15000",
+    10,
+  );
+  return Number.isFinite(raw) && raw > 0 ? raw : 15_000;
+})();
+
+/**
+ * Tunnel-edge HTTP statuses worth retrying on WebSocket upgrade. 502/503 are
+ * emitted by the mux while the upstream Rage REST listener is still warming up.
+ */
+const RETRYABLE_UPGRADE_STATUS = /HTTP\s+(502|503)\b/;
+
 function delay(ms: number): Promise<void> {
   return new Promise((r) => setTimeout(r, ms));
 }
@@ -23,7 +42,7 @@ function connectWebSocketOnce(
       settled = true;
       ws.terminate();
       reject(new Error("WebSocket connection timed out"));
-    }, 45_000);
+    }, PTY_WS_CONNECT_TIMEOUT_MS);
 
     function finish(ok: boolean, result: WebSocket | Error) {
       if (settled) return;
@@ -69,11 +88,7 @@ export async function openPtyWebSocket(
       return await connectWebSocketOnce(wsUrl, protocols);
     } catch (err) {
       lastErr = err instanceof Error ? err : new Error(String(err));
-      const msg = lastErr.message;
-      const retryable =
-        /HTTP\s+(502|503)\b/.test(msg) ||
-        msg.includes("502") ||
-        msg.includes("503");
+      const retryable = RETRYABLE_UPGRADE_STATUS.test(lastErr.message);
 
       if (!retryable || attempt === PTY_WS_MAX_ATTEMPTS) {
         throw lastErr;

src/screens/PtySessionScreen.tsx

@@ -17,17 +17,25 @@ export function PtySessionScreen() {
   const sessionName = params.ptySessionName || params.devboxId;
   const authToken = params.ptyAuthToken;
   const devboxName = params.devboxName || params.devboxId || "devbox";
-  const returnScreen = (params.returnScreen as ScreenName) || "devbox-list";
-  const returnParams = (params.returnParams as RouteParams) || {};
+  const returnScreen: ScreenName =
+    (params.returnScreen as ScreenName) || "devbox-list";
+
+  // Stabilize returnParams across renders — the inline `|| {}` would otherwise
+  // produce a fresh object every render and re-fire any effect depending on it.
+  const returnParamsRaw = params.returnParams as RouteParams | undefined;
+  const returnParams = React.useMemo(
+    () => returnParamsRaw ?? {},
+    [returnParamsRaw],
+  );
+
+  const goBack = React.useCallback(() => {
+    replace(returnScreen, returnParams);
+  }, [replace, returnScreen, returnParams]);
 
   const configOk = !!(baseUrl && sessionName);
   React.useEffect(() => {
-    if (configOk) return;
-    const id = setTimeout(() => {
-      replace(returnScreen, returnParams || {});
-    }, 100);
-    return () => clearTimeout(id);
-  }, [configOk, replace, returnScreen, returnParams]);
+    if (!configOk) goBack();
+  }, [configOk, goBack]);
 
   if (!baseUrl || !sessionName) {
     return (
@@ -57,16 +65,8 @@ export function PtySessionScreen() {
         baseUrl={baseUrl}
         sessionName={sessionName}
         authToken={authToken}
-        onExit={(_code) => {
-          setTimeout(() => {
-            replace(returnScreen, returnParams || {});
-          }, 100);
-        }}
-        onError={(_error) => {
-          setTimeout(() => {
-            replace(returnScreen, returnParams || {});
-          }, 100);
-        }}
+        onExit={goBack}
+        onError={goBack}
       />
     </>
   );