diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 50ca29c3..66586e81 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -152,12 +152,13 @@ PR titles should follow the conventional commit format as they are used for rele src/ ├── cli.ts # Main CLI entry point ├── commands/ # Command implementations -│ ├── devbox/ # Devbox commands +│ ├── devbox/ # Devbox commands (ssh, pty, exec, etc.) │ ├── snapshot/ # Snapshot commands │ ├── blueprint/ # Blueprint commands │ └── object/ # Object storage commands ├── components/ # React/Ink UI components ├── hooks/ # Custom React hooks +├── lib/ # Reusable protocol clients (PTY, etc.) ├── mcp/ # MCP server implementation ├── router/ # Navigation router ├── screens/ # Full-screen views diff --git a/README.md b/README.md index 66b3e5b4..3414efc7 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,7 @@ rli devbox delete - 🎯 **CLI mode** — Traditional commands with text, JSON, and YAML output for scripting - ⚡ Fast and responsive with pagination - 📦 Manage devboxes, snapshots, and blueprints -- 🚀 Execute commands, SSH, view logs in devboxes +- 🚀 Execute commands, SSH, PTY shell, view logs in devboxes - 🤖 **Model Context Protocol (MCP) server for AI integration** ## Installation @@ -69,7 +69,7 @@ The URL must be of the form `https://api.`. The CLI derives other servic | API | `https://api.` (the value of `RUNLOOP_BASE_URL`) | | Platform | `https://platform.` | | SSH | `ssh.:443` | -| Tunnels | `tunnel.` | +| Tunnels | `tunnel.` (PTY sessions reach the devbox over a tunnel created via the API) | ## Usage @@ -109,6 +109,7 @@ rli devbox suspend # Suspend a devbox rli devbox resume # Resume a suspended devbox rli devbox shutdown # Shutdown a devbox rli devbox ssh # SSH into a devbox +rli devbox pty # Connect to a devbox PTY session via W... rli devbox scp # Copy files to/from a devbox using scp... rli devbox rsync # Sync files to/from a devbox using rsy... rli devbox tunnel # Create a port-forwarding tunnel to a ... diff --git a/package.json b/package.json index 81e2e8af..5b547473 100644 --- a/package.json +++ b/package.json @@ -89,6 +89,7 @@ "ink-spinner": "5.0.0", "ink-text-input": "6.0.0", "react": "19.2.0", + "ws": "^8.18.0", "tar-stream": "3.1.7", "yaml": "2.8.3", "zustand": "5.0.10" @@ -142,6 +143,7 @@ "prettier": "3.8.1", "ts-jest": "29.4.6", "ts-node": "10.9.2", + "@types/ws": "^8.5.0", "typescript": "5.9.3" } } diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 16789115..dbe7712a 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -92,6 +92,9 @@ importers: tar-stream: specifier: 3.1.7 version: 3.1.7 + ws: + specifier: ^8.18.0 + version: 8.19.0 yaml: specifier: 2.8.3 version: 2.8.3 @@ -117,6 +120,9 @@ importers: '@types/tar-stream': specifier: 3.1.4 version: 3.1.4 + '@types/ws': + specifier: ^8.5.0 + version: 8.18.1 '@typescript-eslint/eslint-plugin': specifier: 8.54.0 version: 8.54.0(@typescript-eslint/parser@8.54.0(eslint@9.39.2)(typescript@5.9.3))(eslint@9.39.2)(typescript@5.9.3) @@ -839,6 +845,9 @@ packages: '@types/wrap-ansi@3.0.0': resolution: {integrity: sha512-ltIpx+kM7g/MLRZfkbL7EsCEjfzCcScLpkg37eXEtx5kmrAKBkTJwd1GIAjDSL8wTpM6Hzn5YO4pSb91BEwu1g==} + '@types/ws@8.18.1': + resolution: {integrity: sha512-ThVF6DCVhA8kUGy+aazFQ4kXQ7E1Ty7A3ypFOe0IcJV8O/M511G99AW24irKrW56Wt44yG9+ij8FaqoBGkuBXg==} + '@types/yargs-parser@21.0.3': resolution: {integrity: sha512-I4q9QU9MQv4oEOz4tAHJtNz1cwuLxn2F3xcc2iV5WdqLPpUnj30aUuxt1mAxYTG+oe8CZMV/+6rU4S4gRDzqtQ==} @@ -4048,6 +4057,10 @@ snapshots: '@types/wrap-ansi@3.0.0': {} + '@types/ws@8.18.1': + dependencies: + '@types/node': 22.19.7 + '@types/yargs-parser@21.0.3': {} '@types/yargs@17.0.35': diff --git a/src/commands/devbox/pty.ts b/src/commands/devbox/pty.ts new file mode 100644 index 00000000..13e493b1 --- /dev/null +++ b/src/commands/devbox/pty.ts @@ -0,0 +1,217 @@ +import WebSocket from "ws"; +import { cliStatus } from "../../utils/cliStatus.js"; +import { outputError } from "../../utils/output.js"; +import { waitForReady } from "../../utils/ssh.js"; +import { + getPtyBaseUrl, + createPtySessionReleaser, + resolvePtyWebSocketUrl, + createPtyTunnel, + getPtyTunnelBaseUrl, + isLocalPtyOverride, + settleAfterPtyTunnel, + refreshPtySessionAfterAttach, + startPtyIoSession, +} from "../../lib/pty-client.js"; +import { openPtyWebSocket } from "../../lib/pty-ws.js"; + +/** SIGINT/SIGTERM → optional notify then close socket; returns disposer for ws close/error cleanup. */ +function registerPtyInterruptHandlers( + ws: WebSocket, + beforeClose: () => void, +): () => void { + function dispose() { + process.removeListener("SIGINT", onSigint); + process.removeListener("SIGTERM", onSigterm); + } + function onSigint() { + beforeClose(); + dispose(); + ws.close(); + } + function onSigterm() { + beforeClose(); + dispose(); + ws.close(); + } + process.on("SIGINT", onSigint); + process.on("SIGTERM", onSigterm); + return dispose; +} + +function writePtyStreamToStdout(data: WebSocket.RawData): void { + if (Buffer.isBuffer(data)) { + process.stdout.write(data); + } else if (data instanceof ArrayBuffer) { + process.stdout.write(Buffer.from(data)); + } else if (Array.isArray(data)) { + process.stdout.write(Buffer.concat(data)); + } else { + process.stdout.write(String(data)); + } +} + +interface PtyOptions { + session?: string; + command?: string; + wait?: boolean; + timeout?: string; + pollInterval?: string; + output?: string; +} + +export async function ptyDevbox(devboxId: string, options: PtyOptions = {}) { + try { + if (options.wait !== false) { + cliStatus(`Waiting for devbox ${devboxId} to be ready...`); + const isReady = await waitForReady( + devboxId, + parseInt(options.timeout || "180"), + parseInt(options.pollInterval || "3"), + ); + if (!isReady) { + outputError(`Devbox ${devboxId} is not ready. Please try again later.`); + } + } + + let baseUrl: string; + let authToken: string | undefined; + + if (isLocalPtyOverride()) { + baseUrl = getPtyBaseUrl(); + } else { + cliStatus(`Creating PTY tunnel for ${devboxId}...`); + const tunnel = await createPtyTunnel(devboxId); + await settleAfterPtyTunnel(); + baseUrl = getPtyTunnelBaseUrl(tunnel.tunnel_key); + authToken = tunnel.auth_token; + } + + const sessionName = options.session?.trim() || devboxId; + + if (options.command) { + await execCommand(baseUrl, sessionName, options.command, authToken); + } else { + await interactiveSession(baseUrl, sessionName, authToken); + } + } catch (error) { + outputError("Failed to start PTY session", error); + } +} + +const PTY_EXEC_TIMEOUT_MS = (() => { + const v = parseInt(process.env.RUNLOOP_PTY_EXEC_TIMEOUT_MS || "0", 10); + return isNaN(v) || v < 0 ? 0 : v; +})(); + +async function execCommand( + baseUrl: string, + sessionName: string, + command: string, + authToken?: string, +): Promise { + const wsUrl = await resolvePtyWebSocketUrl(baseUrl, sessionName, { + cols: 80, + rows: 24, + authToken, + }); + const ws = await openPtyWebSocket(wsUrl, authToken); + + const releaseOnce = createPtySessionReleaser(baseUrl, sessionName, authToken); + const disposeInterruptSignals = registerPtyInterruptHandlers(ws, releaseOnce); + + let timeoutId: ReturnType | undefined; + + const completion = new Promise((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); + }); + + ws.on("close", () => { + finish(); + resolve(); + }); + + ws.on("error", (err: Error) => { + 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( + baseUrl: string, + sessionName: string, + authToken?: string, +): Promise { + const cols = process.stdout.columns || 80; + const rows = process.stdout.rows || 24; + + const wsUrl = await resolvePtyWebSocketUrl(baseUrl, sessionName, { + cols, + rows, + authToken, + }); + const ws = await openPtyWebSocket(wsUrl, authToken); + + // 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, + authToken, + ); + const disposeSignals = registerPtyInterruptHandlers(ws, releaseOnce); + + await refreshPtySessionAfterAttach( + ws, + baseUrl, + sessionName, + cols, + rows, + authToken, + ); + + try { + await done; + releaseOnce(); + } catch (err) { + releaseOnce(); + throw err; + } finally { + dispose(); + disposeSignals(); + } +} diff --git a/src/commands/object/upload.ts b/src/commands/object/upload.ts index f87afa57..47d4fcff 100644 --- a/src/commands/object/upload.ts +++ b/src/commands/object/upload.ts @@ -214,7 +214,13 @@ async function readStdinBuffer(): Promise { export async function uploadObject(options: UploadObjectOptions) { try { const client = getClient(); - const { paths, name, contentType, output: outputFormat } = options; + const { + paths: rawPaths, + name, + contentType, + output: outputFormat, + } = options; + const paths = [...rawPaths]; if (paths.length === 0) { if (!processUtils.stdin.isTTY) { diff --git a/src/components/DevboxActionsMenu.tsx b/src/components/DevboxActionsMenu.tsx index e135d669..d359dfd8 100644 --- a/src/components/DevboxActionsMenu.tsx +++ b/src/components/DevboxActionsMenu.tsx @@ -27,12 +27,21 @@ import { } from "../services/devboxService.js"; import { StreamingLogsViewer } from "./StreamingLogsViewer.js"; import { DevboxView } from "@runloop/api-client/resources/devboxes.mjs"; +import { + getPtyBaseUrl, + isLocalPtyOverride, + createPtyTunnel, + getPtyTunnelBaseUrl, + settleAfterPtyTunnel, +} from "../lib/pty-client.js"; +import { waitForReady } from "../utils/ssh.js"; type Operation = | "exec" | "upload" | "snapshot" | "ssh" + | "pty" | "logs" | "tunnel" | "suspend" @@ -167,6 +176,13 @@ export const DevboxActionsMenu = ({ icon: figures.arrowRight, shortcut: "s", }, + { + key: "pty", + label: "PTY Shell", + color: colors.primary, + icon: figures.arrowRight, + shortcut: "y", + }, { key: "tunnel", label: "Open Tunnel", @@ -244,7 +260,7 @@ export const DevboxActionsMenu = ({ // Auto-execute operations that don't need input (except delete which needs confirmation) React.useEffect(() => { - const autoExecuteOps = ["ssh", "logs", "suspend", "resume"]; + const autoExecuteOps = ["ssh", "pty", "logs", "suspend", "resume"]; if ( executingOperation && autoExecuteOps.includes(executingOperation) && @@ -681,6 +697,33 @@ export const DevboxActionsMenu = ({ }); break; + case "pty": { + await waitForReady(devbox.id, 180, 3); + + let ptyBaseUrl: string; + let ptyAuthToken: string | undefined; + + if (isLocalPtyOverride()) { + ptyBaseUrl = getPtyBaseUrl(); + } else { + const tunnel = await createPtyTunnel(devbox.id); + await settleAfterPtyTunnel(); + ptyBaseUrl = getPtyTunnelBaseUrl(tunnel.tunnel_key); + ptyAuthToken = tunnel.auth_token; + } + + navigate("pty-session", { + ptyBaseUrl, + ptySessionName: devbox.id, + ptyAuthToken, + devboxId: devbox.id, + devboxName: devbox.name || devbox.id, + returnScreen: currentScreen, + returnParams: params, + }); + break; + } + case "logs": // Set flag to show streaming logs viewer const logsResult: any = { diff --git a/src/components/InteractivePty.tsx b/src/components/InteractivePty.tsx new file mode 100644 index 00000000..9a2f9236 --- /dev/null +++ b/src/components/InteractivePty.tsx @@ -0,0 +1,169 @@ +import React from "react"; +import WebSocket from "ws"; +import { + createPtySessionReleaser, + 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"; +import { processUtils } from "../utils/processUtils.js"; + +interface InteractivePtyProps { + baseUrl: string; + sessionName: string; + authToken?: string; + onExit?: (code: number | null) => void; + onError?: (error: Error) => void; +} + +/** + * Hand control of the terminal back to the PTY by pausing Ink's raw-mode input. + * Called before opening the WebSocket so Ink doesn't race on stdin. + */ +function releaseTerminal(): void { + process.stdin.pause(); + if (processUtils.stdin.isTTY && processUtils.stdin.setRawMode) { + processUtils.stdin.setRawMode(false); + } + if (processUtils.stdout.isTTY) { + processUtils.stdout.write("\x1b[0m"); + } +} + +/** + * Return the terminal to Ink after a PTY session ends. + * Clears any residual PTY output and re-enables raw mode so Ink can + * receive keyboard events again. + */ +function restoreTerminal(): void { + clearScreen(); + if (processUtils.stdin.isTTY && processUtils.stdin.setRawMode) { + processUtils.stdin.setRawMode(true); + } + process.stdin.resume(); +} + +export const InteractivePty: React.FC = ({ + baseUrl, + sessionName, + authToken, + onExit, + onError, +}) => { + const wsRef = React.useRef(null); + const hasStartedRef = React.useRef(false); + const onExitRef = React.useRef(onExit); + const onErrorRef = React.useRef(onError); + React.useEffect(() => { + onExitRef.current = onExit; + onErrorRef.current = onError; + }); + + React.useEffect(() => { + if (hasStartedRef.current) return; + hasStartedRef.current = true; + + releaseTerminal(); + + let cancelled = false; + let ioDispose: (() => void) | null = null; + + setImmediate(async () => { + try { + const cols = process.stdout.columns || 80; + const rows = process.stdout.rows || 24; + + const wsUrl = await resolvePtyWebSocketUrl(baseUrl, sessionName, { + cols, + rows, + authToken, + }); + + if (cancelled) return; + + const ws = await openPtyWebSocket(wsUrl, authToken); + + if (cancelled) { + ws.close(); + return; + } + + wsRef.current = ws; + + // 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, + authToken, + ); + const { dispose, done } = startPtyIoSession( + ws, + baseUrl, + sessionName, + authToken, + ); + + const ioCleanup = () => { + dispose(); + releaseServerSession(); + }; + ioDispose = ioCleanup; + + if (cancelled) { + ioCleanup(); + return; + } + + await refreshPtySessionAfterAttach( + ws, + baseUrl, + sessionName, + cols, + rows, + authToken, + ); + + done + .then((code) => { + wsRef.current = null; + ioCleanup(); + restoreTerminal(); + hasStartedRef.current = false; + onExitRef.current?.(code === PTY_NORMAL_CLOSE_CODE ? 0 : code); + }) + .catch((err: Error) => { + wsRef.current = null; + ioCleanup(); + restoreTerminal(); + hasStartedRef.current = false; + onErrorRef.current?.(err); + }); + } catch (err) { + if (!cancelled) { + restoreTerminal(); + hasStartedRef.current = false; + onErrorRef.current?.( + err instanceof Error ? err : new Error(String(err)), + ); + } + } + }); + + return () => { + cancelled = true; + ioDispose?.(); + if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) { + wsRef.current.close(); + } + wsRef.current = null; + restoreTerminal(); + hasStartedRef.current = false; + }; + }, [baseUrl, sessionName, authToken]); + + return null; +}; diff --git a/src/lib/pty-client.ts b/src/lib/pty-client.ts new file mode 100644 index 00000000..dd0380ad --- /dev/null +++ b/src/lib/pty-client.ts @@ -0,0 +1,424 @@ +import WebSocket from "ws"; +import { baseUrl, getConfig } from "../utils/config.js"; +import { getTunnelUrl } from "../utils/url.js"; +import { processUtils } from "../utils/processUtils.js"; + +const PTY_CONNECT_MAX_ATTEMPTS = Math.max( + 1, + parseInt(process.env.RUNLOOP_PTY_CONNECT_RETRIES || "3", 10) || 3, +); + +/** Optional pause after `create_pty_tunnel` so the mux can route (ms). `0` disables. */ +export async function settleAfterPtyTunnel(): Promise { + const ms = Math.max( + 0, + parseInt(process.env.RUNLOOP_PTY_POST_TUNNEL_MS || "1500", 10) || 0, + ); + if (ms > 0) await new Promise((r) => setTimeout(r, ms)); +} + +function ptyBootstrapTokenInQuery(): boolean { + const v = + process.env.RUNLOOP_PTY_BOOTSTRAP_TOKEN_IN_QUERY?.toLowerCase().trim(); + return v === "1" || v === "true" || v === "yes"; +} + +function ptyBootstrapConnectionClose(): boolean { + const v = + process.env.RUNLOOP_PTY_BOOTSTRAP_CONNECTION_CLOSE?.toLowerCase().trim(); + return v === "1" || v === "true" || v === "yes"; +} + +async function readErrorSnippet(res: Response): Promise { + try { + const t = (await res.text()).trim(); + if (!t) return ""; + return t.length > 500 ? `${t.slice(0, 500)}…` : t; + } catch { + return ""; + } +} + +function formatHttpError( + prefix: string, + res: Response, + detail: string, +): string { + return detail + ? `${prefix}: ${res.status} ${res.statusText} — ${detail}` + : `${prefix}: ${res.status} ${res.statusText}`; +} + +/** + * Rage REST port on real devboxes (PTY HTTP + WebSocket). Matches platform + * `RAGE_REST_PORT` in Java/K8s; not configurable in normal flows. + * + * MUX PTY origins use {@link getPtyTunnelBaseUrl}: `https://13-{tunnel_key}.tunnel.` + * (via {@link getTunnelUrl}). Must include the `tunnel.` label — not `13-{key}.`. + * Local PTY tests use an ephemeral port on 127.0.0.1 instead (`RUNLOOP_PTY_URL`). + */ +export const RAGE_REST_PORT = 13 as const; + +/** Default local PTY dev server when `RUNLOOP_PTY_URL` is unset (bazel `//src/test/pty:test_bin`). */ +export const PTY_BASE_URL_LOCAL = "http://localhost:5000"; + +export interface PtyConnectResponse { + session_name: string; + status: string; + protocol_version: string; + connect_url: string; + created: boolean; + attached: boolean; + cols: number; + rows: number; + idle_ttl_seconds: number; +} + +export interface PtyControlResult { + session_name: string; + status: string; +} + +export interface PtyTunnelView { + tunnel_key: string; + auth_token: string; +} + +export type ControlAction = + | { action: "resize"; cols: number; rows: number } + | { action: "signal"; signal: string } + | { action: "close" }; + +export async function createPtyTunnel( + devboxId: string, +): Promise { + const apiKey = getConfig().apiKey; + if (!apiKey) throw new Error("API key not configured"); + + const url = `${baseUrl()}/v1/devboxes/${encodeURIComponent(devboxId)}/create_pty_tunnel`; + const res = await fetch(url, { + method: "POST", + headers: { + Authorization: `Bearer ${apiKey}`, + "Content-Type": "application/json", + }, + body: "{}", + }); + if (!res.ok) { + const detail = await readErrorSnippet(res); + throw new Error(formatHttpError("Create PTY tunnel failed", res, detail)); + } + return res.json() as Promise; +} + +/** + * Public origin for PTY after `create_pty_tunnel` (MUX): `https://13-{tunnel_key}.tunnel.` + * (same `{port}-{key}.tunnel…` shape as port-forward tunnels, with port fixed at + * {@link RAGE_REST_PORT}). Paths are still `GET /pty/{session_name}`, WebSocket `…/attach`, etc. + */ +export function getPtyTunnelBaseUrl(tunnelKey: string): string { + return getTunnelUrl(RAGE_REST_PORT, tunnelKey); +} + +export function isLocalPtyOverride(): boolean { + return !!process.env.RUNLOOP_PTY_URL?.trim(); +} + +export function getPtyBaseUrl(): string { + const override = process.env.RUNLOOP_PTY_URL?.trim(); + if (override) return override; + + // Production uses `create_pty_tunnel` + `getPtyTunnelBaseUrl`; this matches the local test server default. + return PTY_BASE_URL_LOCAL; +} + +/** `GET /pty/{session}` bootstrap (JSON + `connect_url`). Required before WebSocket attach unless `RUNLOOP_PTY_ATTACH_ONLY=1`. */ +export async function ptyConnect( + baseUrl: string, + sessionName: string, + opts?: { cols?: number; rows?: number; authToken?: string }, +): Promise { + const params = new URLSearchParams(); + if (opts?.cols) params.set("cols", String(opts.cols)); + if (opts?.rows) params.set("rows", String(opts.rows)); + if (opts?.authToken && ptyBootstrapTokenInQuery()) { + params.set("token", opts.authToken); + } + + const qs = params.toString(); + const pathSession = encodeURIComponent(sessionName); + const url = `${baseUrl}/pty/${pathSession}${qs ? `?${qs}` : ""}`; + + const headers: Record = { + Accept: "application/json", + "User-Agent": "Runloop-rli/pty-bootstrap", + }; + if (opts?.authToken) { + headers["Authorization"] = `Bearer ${opts.authToken}`; + } + if (opts?.authToken && ptyBootstrapConnectionClose()) { + headers["Connection"] = "close"; + } + + for (let attempt = 1; attempt <= PTY_CONNECT_MAX_ATTEMPTS; attempt++) { + const res = await fetch(url, { headers }); + if (res.ok) { + return res.json() as Promise; + } + + let detail = await readErrorSnippet(res); + if (res.status === 502) { + const hint = + "Tunnel 502 usually means Rage REST is unreachable from the mux (same URL with curl typically matches)."; + detail = detail ? `${detail} — ${hint}` : hint; + } + const msg = formatHttpError("PTY connect failed", res, detail); + const retryable = res.status === 502 || res.status === 503; + + if (retryable && attempt < PTY_CONNECT_MAX_ATTEMPTS) { + const delayMs = Math.min(10_000, 400 * 2 ** (attempt - 1)); + await new Promise((r) => setTimeout(r, delayMs)); + continue; + } + + throw new Error(msg); + } + + throw new Error("PTY connect failed: exhausted retries"); +} + +export async function ptyControl( + baseUrl: string, + sessionName: string, + action: ControlAction, + authToken?: string, +): Promise { + const url = `${baseUrl}/pty/${encodeURIComponent(sessionName)}/control`; + const headers: Record = { + "Content-Type": "application/json", + }; + if (authToken) { + headers["Authorization"] = `Bearer ${authToken}`; + } + const res = await fetch(url, { + method: "POST", + headers, + body: JSON.stringify(action), + }); + if (!res.ok) { + const detail = await readErrorSnippet(res); + throw new Error(formatHttpError("PTY control failed", res, detail)); + } + return res.json() as Promise; +} + +/** Best-effort: release the PTY session on the server so the same session name can attach again. */ +export function ptyNotifyClosed( + baseUrl: string, + sessionName: string, + authToken?: string, +): void { + void ptyControl(baseUrl, sessionName, { action: "close" }, authToken).catch( + () => {}, + ); +} + +/** Calls {@link ptyNotifyClosed} at most once (ws close, signals, Ink unmount, etc.). */ +export function createPtySessionReleaser( + baseUrl: string, + sessionName: string, + authToken?: string, +): () => void { + let released = false; + return () => { + if (released) return; + released = true; + ptyNotifyClosed(baseUrl, sessionName, authToken); + }; +} + +/** 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). + */ +export async function refreshPtySessionAfterAttach( + ws: { readyState: number; send(data: string | Buffer): void }, + baseUrl: string, + sessionName: string, + cols: number, + rows: number, + authToken?: string, +): Promise { + await ptyControl( + baseUrl, + sessionName, + { action: "resize", cols, rows }, + authToken, + ).catch(() => {}); + if (ws.readyState === WS_READY_STATE_OPEN) { + ws.send("\r"); + } +} + +/** WebSocket attach URL built from `baseUrl` + `connectUrl` path. No auth token in the URL. */ +export function buildWsUrl(baseUrl: string, connectUrl: string): string { + const parsed = new URL(baseUrl); + const protocol = parsed.protocol === "https:" ? "wss:" : "ws:"; + return `${protocol}//${parsed.host}${connectUrl}`; +} + +/** + * Path for WebSocket attach only (no prior `GET /pty/{session}` bootstrap). + * Optional `cols` / `rows` match query params supported on upgrade. + */ +export function buildPtyAttachPath( + sessionName: string, + cols?: number, + rows?: number, +): string { + const params = new URLSearchParams(); + if (cols) params.set("cols", String(cols)); + if (rows) params.set("rows", String(rows)); + const qs = params.toString(); + return `/pty/${encodeURIComponent(sessionName)}/attach${qs ? `?${qs}` : ""}`; +} + +/** Full `wss:` URL for PTY streaming with no HTTP bootstrap (experimental). */ +export function buildPtyAttachWsUrl( + baseUrl: string, + sessionName: string, + opts?: { cols?: number; rows?: number }, +): string { + const path = buildPtyAttachPath(sessionName, opts?.cols, opts?.rows); + return buildWsUrl(baseUrl, path); +} + +export function isPtyAttachOnlyMode(): boolean { + const v = process.env.RUNLOOP_PTY_ATTACH_ONLY?.toLowerCase().trim(); + return v === "1" || v === "true" || v === "yes"; +} + +/** + * Resolves the WebSocket URL: by default `ptyConnect` bootstrap then `connect_url`; + * if `RUNLOOP_PTY_ATTACH_ONLY=1`, opens `/pty/.../attach` directly (often 502 if server requires bootstrap). + * Auth token is passed via HTTP Authorization header (bootstrap) and as WebSocket subprotocol + * (Sec-WebSocket-Protocol), never in the URL. + */ +export async function resolvePtyWebSocketUrl( + baseUrl: string, + sessionName: string, + opts: { cols: number; rows: number; authToken?: string }, +): Promise { + if (isPtyAttachOnlyMode()) { + return buildPtyAttachWsUrl(baseUrl, sessionName, opts); + } + const connectResponse = await ptyConnect(baseUrl, sessionName, { + cols: opts.cols, + rows: opts.rows, + authToken: opts.authToken, + }); + return buildWsUrl(baseUrl, connectResponse.connect_url); +} + +export interface PtyIoHandle { + /** + * Stop the session: removes all stdin/SIGWINCH/WS listeners, disables raw mode, + * pauses stdin. Does not close the WebSocket or release the server session. + */ + dispose: () => void; + /** + * Resolves with the WS close code when the server closes the session. + * Rejects if a WS error fires before close. Never settles after dispose(). + */ + done: Promise; +} + +/** + * Wire an open WebSocket into an interactive PTY I/O session: + * - Sets stdin to raw mode and pumps keystrokes to the WS + * - Forwards WS messages (binary or text) to stdout + * - Sends resize commands on SIGWINCH + * + * Returns a handle with `dispose()` to tear down listeners and `done` that + * resolves with the WS close code when the session ends naturally. + */ +export function startPtyIoSession( + ws: WebSocket, + baseUrl: string, + sessionName: string, + authToken?: string, +): PtyIoHandle { + let doneResolve!: (code: number) => void; + let doneReject!: (err: Error) => void; + const done = new Promise((resolve, reject) => { + doneResolve = resolve; + doneReject = reject; + }); + + if (processUtils.stdin.isTTY && processUtils.stdin.setRawMode) { + processUtils.stdin.setRawMode(true); + } + process.stdin.resume(); + + const onStdinData = (data: Buffer) => { + if (ws.readyState === WS_READY_STATE_OPEN) { + ws.send(data); + } + }; + + const onResize = () => { + const cols = process.stdout.columns || 80; + const rows = process.stdout.rows || 24; + ptyControl( + baseUrl, + sessionName, + { action: "resize", cols, rows }, + authToken, + ).catch(() => {}); + }; + + const onMessage = (data: WebSocket.RawData) => { + if (data instanceof ArrayBuffer) { + process.stdout.write(Buffer.from(data)); + } else if (Buffer.isBuffer(data)) { + process.stdout.write(data); + } else if (Array.isArray(data)) { + process.stdout.write(Buffer.concat(data)); + } else { + process.stdout.write(String(data)); + } + }; + + const onClose = (code: number) => doneResolve(code); + const onError = (err: Error) => doneReject(err); + + process.stdin.on("data", onStdinData); + process.on("SIGWINCH", onResize); + ws.on("message", onMessage); + ws.on("close", onClose); + ws.on("error", onError); + + const dispose = () => { + process.stdin.removeListener("data", onStdinData); + process.removeListener("SIGWINCH", onResize); + ws.removeListener("message", onMessage); + ws.removeListener("close", onClose); + ws.removeListener("error", onError); + if (processUtils.stdin.isTTY && processUtils.stdin.setRawMode) { + processUtils.stdin.setRawMode(false); + } + process.stdin.pause(); + }; + + return { dispose, done }; +} diff --git a/src/lib/pty-ws.ts b/src/lib/pty-ws.ts new file mode 100644 index 00000000..24d8f0db --- /dev/null +++ b/src/lib/pty-ws.ts @@ -0,0 +1,103 @@ +import type { ClientRequest, IncomingMessage } from "http"; +import WebSocket from "ws"; + +const PTY_WS_MAX_ATTEMPTS = Math.max( + 1, + 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 { + return new Promise((r) => setTimeout(r, ms)); +} + +function connectWebSocketOnce( + wsUrl: string, + protocols: string[] | undefined, +): Promise { + return new Promise((resolve, reject) => { + const ws = new WebSocket(wsUrl, protocols ?? []); + let settled = false; + + const connectTimer = setTimeout(() => { + if (settled) return; + settled = true; + ws.terminate(); + reject(new Error("WebSocket connection timed out")); + }, PTY_WS_CONNECT_TIMEOUT_MS); + + function finish(ok: boolean, result: WebSocket | Error) { + if (settled) return; + settled = true; + clearTimeout(connectTimer); + if (ok) resolve(result as WebSocket); + else reject(result as Error); + } + + ws.once("open", () => finish(true, ws)); + + ws.once( + "unexpected-response", + (_req: ClientRequest, res: IncomingMessage) => { + const code = res.statusCode ?? 0; + finish( + false, + new Error( + `WebSocket upgrade failed: HTTP ${code} ${res.statusMessage ?? ""}`.trim(), + ), + ); + }, + ); + + ws.once("error", (err: Error) => finish(false, err)); + }); +} + +/** + * Connect to PTY attach URL; retries HTTP 502/503 from tunnel edge during warm-up. + * authToken is passed as a WebSocket subprotocol (Sec-WebSocket-Protocol) so it + * is not visible in server access logs as a URL query parameter. + */ +export async function openPtyWebSocket( + wsUrl: string, + authToken: string | undefined, +): Promise { + const protocols = authToken ? [authToken] : undefined; + let lastErr: Error | undefined; + + for (let attempt = 1; attempt <= PTY_WS_MAX_ATTEMPTS; attempt++) { + try { + return await connectWebSocketOnce(wsUrl, protocols); + } catch (err) { + lastErr = err instanceof Error ? err : new Error(String(err)); + const retryable = RETRYABLE_UPGRADE_STATUS.test(lastErr.message); + + if (!retryable || attempt === PTY_WS_MAX_ATTEMPTS) { + throw lastErr; + } + + const delayMs = Math.min(10_000, 400 * 2 ** (attempt - 1)); + await delay(delayMs); + } + } + + throw lastErr ?? new Error("WebSocket connect failed"); +} diff --git a/src/router/Router.tsx b/src/router/Router.tsx index 5c5248dd..61586b2e 100644 --- a/src/router/Router.tsx +++ b/src/router/Router.tsx @@ -55,6 +55,7 @@ const KNOWN_SCREENS: Set = new Set([ "object-detail", "object-create", "ssh-session", + "pty-session", "benchmark-menu", "benchmark-list", "benchmark-detail", @@ -138,6 +139,7 @@ import { ObjectListScreen } from "../screens/ObjectListScreen.js"; import { ObjectDetailScreen } from "../screens/ObjectDetailScreen.js"; import { ObjectCreateScreen } from "../screens/ObjectCreateScreen.js"; import { SSHSessionScreen } from "../screens/SSHSessionScreen.js"; +import { PtySessionScreen } from "../screens/PtySessionScreen.js"; import { BenchmarkMenuScreen } from "../screens/BenchmarkMenuScreen.js"; import { BenchmarkListScreen } from "../screens/BenchmarkListScreen.js"; import { BenchmarkDetailScreen } from "../screens/BenchmarkDetailScreen.js"; @@ -361,6 +363,9 @@ export function Router() { {currentScreen === "ssh-session" && ( )} + {currentScreen === "pty-session" && ( + + )} {currentScreen === "benchmark-menu" && ( )} diff --git a/src/screens/PtySessionScreen.tsx b/src/screens/PtySessionScreen.tsx new file mode 100644 index 00000000..d3c14136 --- /dev/null +++ b/src/screens/PtySessionScreen.tsx @@ -0,0 +1,73 @@ +import React from "react"; +import { Box, Text } from "ink"; +import { InteractivePty } from "../components/InteractivePty.js"; +import { + useNavigation, + type ScreenName, + type RouteParams, +} from "../store/navigationStore.js"; +import { Breadcrumb } from "../components/Breadcrumb.js"; +import { colors } from "../utils/theme.js"; +import figures from "figures"; + +export function PtySessionScreen() { + const { params, replace } = useNavigation(); + + const baseUrl = params.ptyBaseUrl; + const sessionName = params.ptySessionName || params.devboxId; + const authToken = params.ptyAuthToken; + const devboxName = params.devboxName || params.devboxId || "devbox"; + 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) goBack(); + }, [configOk, goBack]); + + if (!baseUrl || !sessionName) { + return ( + <> + + + + {figures.cross} Missing PTY configuration. Returning... + + + + ); + } + + return ( + <> + + + + {figures.play} Connecting to {devboxName}... + + + Press Ctrl+D or type exit to disconnect + + + + + ); +} diff --git a/src/store/navigationStore.tsx b/src/store/navigationStore.tsx index 0b465131..0c7afef5 100644 --- a/src/store/navigationStore.tsx +++ b/src/store/navigationStore.tsx @@ -46,6 +46,7 @@ export type ScreenName = | "object-detail" | "object-create" | "ssh-session" + | "pty-session" | "benchmark-list" | "benchmark-detail" | "benchmark-run-list" @@ -79,6 +80,10 @@ export interface RouteParams { devboxName?: string; returnScreen?: ScreenName; returnParams?: RouteParams; + // PTY session params + ptyBaseUrl?: string; + ptySessionName?: string; + ptyAuthToken?: string; // Exec session params executionId?: string; execCommand?: string; diff --git a/src/utils/commands.ts b/src/utils/commands.ts index 19c9732c..35663f2d 100644 --- a/src/utils/commands.ts +++ b/src/utils/commands.ts @@ -224,6 +224,28 @@ export function createProgram(): Command { await sshDevbox(id, options); }); + devbox + .command("pty ") + .alias("shell") + .description("Connect to a devbox PTY session via WebSocket") + .option("--session ", "Session name (defaults to devbox ID)") + .option("--command ", "Execute command non-interactively") + .option("--no-wait", "Do not wait for devbox to be ready") + .option( + "--timeout ", + "Timeout in seconds to wait for readiness", + "180", + ) + .option( + "--poll-interval ", + "Polling interval in seconds while waiting", + "3", + ) + .action(async (id, options) => { + const { ptyDevbox } = await import("../commands/devbox/pty.js"); + await ptyDevbox(id, options); + }); + devbox .command("scp ") .description( diff --git a/src/utils/config.ts b/src/utils/config.ts index e9e8801c..3c29af1a 100644 --- a/src/utils/config.ts +++ b/src/utils/config.ts @@ -224,6 +224,6 @@ To make it permanent, add this line to your shell config: • For zsh: echo 'export RUNLOOP_API_KEY=your_api_key_here' >> ~/.zshrc • For bash: echo 'export RUNLOOP_API_KEY=your_api_key_here' >> ~/.bashrc -Then restart your terminal or run: source ~/.zshrc (or ~/.bashrc) +Restart your terminal or run \`source ~/.zshrc\` / \`source ~/.bashrc\` so the variable is picked up. `; } diff --git a/src/utils/url.ts b/src/utils/url.ts index 6ba5b211..309511a2 100644 --- a/src/utils/url.ts +++ b/src/utils/url.ts @@ -72,7 +72,11 @@ export function getSettingsUrl(): string { } /** - * Generate a tunnel URL for the given port and tunnel key. + * Public URL for a devbox tunnel: `https://{remotePort}-{tunnelKey}.tunnel.`. + * User port-forwards pass the chosen remote port; PTY (Rage REST) uses port 13 via + * `getPtyTunnelBaseUrl` → `getTunnelUrl(13, tunnelKey)`. The `tunnel.` hostname + * segment is required (see `tunnelBaseHostname()`), e.g. not `13-{key}.runloop.pro`. + * * Pass a number for a real URL, or a string like "{port}" for a display pattern. */ export function getTunnelUrl(port: number | string, tunnelKey: string): string { diff --git a/tests/__tests__/lib/pty-client.test.ts b/tests/__tests__/lib/pty-client.test.ts new file mode 100644 index 00000000..90c2b4b3 --- /dev/null +++ b/tests/__tests__/lib/pty-client.test.ts @@ -0,0 +1,265 @@ +/** + * Tests for PTY client helpers. + * + * Focused on the pure helpers (URL builders, attach-mode toggle, releaser + * once-only semantics) and the bootstrap retry classifier — i.e. the things + * a future regression is most likely to silently break. + */ + +import { + jest, + describe, + it, + expect, + beforeEach, + afterEach, +} from "@jest/globals"; +import { + buildWsUrl, + buildPtyAttachPath, + isPtyAttachOnlyMode, + resolvePtyWebSocketUrl, + createPtySessionReleaser, + ptyConnect, +} from "@/lib/pty-client.js"; + +describe("buildWsUrl", () => { + it("rewrites https:// to wss:// and keeps the host", () => { + expect(buildWsUrl("https://13-abc.tunnel.example.com", "/pty/foo")).toBe( + "wss://13-abc.tunnel.example.com/pty/foo", + ); + }); + + it("rewrites http:// to ws:// for local dev", () => { + expect(buildWsUrl("http://localhost:5000", "/pty/foo/attach")).toBe( + "ws://localhost:5000/pty/foo/attach", + ); + }); + + it("preserves port and ignores baseUrl path", () => { + expect(buildWsUrl("https://api.example.com:8443/v1", "/pty/x")).toBe( + "wss://api.example.com:8443/pty/x", + ); + }); +}); + +describe("buildPtyAttachPath", () => { + it("URL-encodes session names with special characters", () => { + expect(buildPtyAttachPath("session/with spaces & symbols")).toBe( + "/pty/session%2Fwith%20spaces%20%26%20symbols/attach", + ); + }); + + it("omits the query string when no cols/rows provided", () => { + expect(buildPtyAttachPath("s")).toBe("/pty/s/attach"); + }); + + it("includes cols/rows when set", () => { + expect(buildPtyAttachPath("s", 120, 40)).toBe( + "/pty/s/attach?cols=120&rows=40", + ); + }); + + it("drops zero cols/rows (treated as 'unset')", () => { + // Mirrors the URLSearchParams `if (cols)` filtering — zero cols isn't a + // real terminal size, so emitting `cols=0` would be misleading. + expect(buildPtyAttachPath("s", 0, 0)).toBe("/pty/s/attach"); + }); +}); + +describe("isPtyAttachOnlyMode", () => { + const originalEnv = process.env.RUNLOOP_PTY_ATTACH_ONLY; + + afterEach(() => { + if (originalEnv === undefined) { + delete process.env.RUNLOOP_PTY_ATTACH_ONLY; + } else { + process.env.RUNLOOP_PTY_ATTACH_ONLY = originalEnv; + } + }); + + it.each([ + ["1", true], + ["true", true], + ["TRUE", true], + ["yes", true], + ["0", false], + ["false", false], + ["", false], + ])("env=%s → %s", (value, expected) => { + process.env.RUNLOOP_PTY_ATTACH_ONLY = value; + expect(isPtyAttachOnlyMode()).toBe(expected); + }); +}); + +describe("resolvePtyWebSocketUrl", () => { + const originalFetch = globalThis.fetch; + const originalAttachOnly = process.env.RUNLOOP_PTY_ATTACH_ONLY; + const mockFetch = jest.fn(); + + beforeEach(() => { + globalThis.fetch = mockFetch; + mockFetch.mockReset(); + delete process.env.RUNLOOP_PTY_ATTACH_ONLY; + }); + + afterEach(() => { + globalThis.fetch = originalFetch; + if (originalAttachOnly === undefined) { + delete process.env.RUNLOOP_PTY_ATTACH_ONLY; + } else { + process.env.RUNLOOP_PTY_ATTACH_ONLY = originalAttachOnly; + } + }); + + it("returns the bootstrap connect_url converted to wss when attach-only is off", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => ({ + session_name: "s", + status: "ready", + protocol_version: "1", + connect_url: "/pty/s/attach?token=ephemeral", + created: true, + attached: false, + cols: 80, + rows: 24, + idle_ttl_seconds: 300, + }), + } as Response); + + const url = await resolvePtyWebSocketUrl( + "https://13-abc.tunnel.example.com", + "s", + { cols: 80, rows: 24 }, + ); + expect(url).toBe( + "wss://13-abc.tunnel.example.com/pty/s/attach?token=ephemeral", + ); + expect(mockFetch).toHaveBeenCalledTimes(1); + }); + + it("skips bootstrap and builds attach URL directly when attach-only is set", async () => { + process.env.RUNLOOP_PTY_ATTACH_ONLY = "1"; + const url = await resolvePtyWebSocketUrl( + "https://13-abc.tunnel.example.com", + "s", + { cols: 100, rows: 30 }, + ); + expect(url).toBe( + "wss://13-abc.tunnel.example.com/pty/s/attach?cols=100&rows=30", + ); + expect(mockFetch).not.toHaveBeenCalled(); + }); +}); + +describe("createPtySessionReleaser", () => { + const originalFetch = globalThis.fetch; + const mockFetch = jest.fn(); + + beforeEach(() => { + globalThis.fetch = mockFetch; + mockFetch.mockReset(); + mockFetch.mockResolvedValue({ + ok: true, + json: async () => ({ session_name: "s", status: "closed" }), + } as Response); + }); + + afterEach(() => { + globalThis.fetch = originalFetch; + }); + + it("only triggers the underlying close call once across many releases", async () => { + const release = createPtySessionReleaser( + "https://h.example.com", + "session-a", + "tok", + ); + release(); + release(); + release(); + // Releaser fires `ptyNotifyClosed` asynchronously (void-thened); give it a tick. + await new Promise((r) => setImmediate(r)); + expect(mockFetch).toHaveBeenCalledTimes(1); + const [url, init] = mockFetch.mock.calls[0]; + expect(url).toBe("https://h.example.com/pty/session-a/control"); + expect(init?.method).toBe("POST"); + expect(JSON.parse(init?.body as string)).toEqual({ action: "close" }); + expect( + (init?.headers as Record)["Authorization"], + ).toBe("Bearer tok"); + }); + + it("does not propagate ptyControl failures to the caller", async () => { + mockFetch.mockResolvedValue({ + ok: false, + status: 500, + statusText: "boom", + text: async () => "internal", + } as Response); + + const release = createPtySessionReleaser("https://h.example.com", "s"); + expect(() => release()).not.toThrow(); + await new Promise((r) => setImmediate(r)); + }); +}); + +describe("ptyConnect retry classifier", () => { + const originalFetch = globalThis.fetch; + const mockFetch = jest.fn(); + + beforeEach(() => { + globalThis.fetch = mockFetch; + mockFetch.mockReset(); + }); + + afterEach(() => { + globalThis.fetch = originalFetch; + }); + + it("retries 502 then succeeds", async () => { + mockFetch + .mockResolvedValueOnce({ + ok: false, + status: 502, + statusText: "Bad Gateway", + text: async () => "edge", + } as Response) + .mockResolvedValueOnce({ + ok: true, + json: async () => ({ + session_name: "s", + status: "ready", + protocol_version: "1", + connect_url: "/pty/s/attach", + created: true, + attached: false, + cols: 80, + rows: 24, + idle_ttl_seconds: 300, + }), + } as Response); + + const res = await ptyConnect("https://h.example.com", "s", { + cols: 80, + rows: 24, + }); + expect(res.connect_url).toBe("/pty/s/attach"); + expect(mockFetch).toHaveBeenCalledTimes(2); + }); + + it("does not retry 401 / 404 etc.", async () => { + mockFetch.mockResolvedValueOnce({ + ok: false, + status: 401, + statusText: "Unauthorized", + text: async () => "missing token", + } as Response); + + await expect( + ptyConnect("https://h.example.com", "s", { cols: 80, rows: 24 }), + ).rejects.toThrow(/401/); + expect(mockFetch).toHaveBeenCalledTimes(1); + }); +}); diff --git a/tests/__tests__/lib/pty-ws.test.ts b/tests/__tests__/lib/pty-ws.test.ts new file mode 100644 index 00000000..4e36f84c --- /dev/null +++ b/tests/__tests__/lib/pty-ws.test.ts @@ -0,0 +1,156 @@ +/** + * Tests for the PTY WebSocket connection helper. + * + * Focused on the retry classifier — the bit most likely to silently break + * once people start poking at error message formatting. + * + * `ws` is mocked via `jest.unstable_mockModule`, which requires the module + * under test to be loaded dynamically *after* the mock is registered. We do + * this once in `beforeAll` and store the export at the top of the file so + * individual tests can rely on plain references rather than scattered + * `await import(...)` calls. + */ + +import { + jest, + describe, + it, + expect, + beforeAll, + beforeEach, + afterEach, +} from "@jest/globals"; +import { EventEmitter } from "events"; + +/** + * Minimal fake `ws` that lets a test drive the connect lifecycle (open, + * unexpected-response, error) without touching the network or real timers. + */ +class FakeWs extends EventEmitter { + terminated = false; + terminate() { + this.terminated = true; + } + // Match the `ws` API surface the code under test touches. + close() {} +} + +const fakeWsInstances: FakeWs[] = []; +const wsConstructorMock = jest.fn(() => { + const ws = new FakeWs(); + fakeWsInstances.push(ws); + return ws; +}); + +jest.unstable_mockModule("ws", () => ({ + default: wsConstructorMock, + __esModule: true, +})); + +let openPtyWebSocket: typeof import("@/lib/pty-ws.js").openPtyWebSocket; + +beforeAll(async () => { + ({ openPtyWebSocket } = await import("@/lib/pty-ws.js")); +}); + +describe("openPtyWebSocket", () => { + const originalRetries = process.env.RUNLOOP_PTY_WS_RETRIES; + + beforeEach(() => { + fakeWsInstances.length = 0; + wsConstructorMock.mockClear(); + process.env.RUNLOOP_PTY_WS_RETRIES = "3"; + }); + + afterEach(() => { + if (originalRetries === undefined) { + delete process.env.RUNLOOP_PTY_WS_RETRIES; + } else { + process.env.RUNLOOP_PTY_WS_RETRIES = originalRetries; + } + }); + + it("resolves on open", async () => { + const promise = openPtyWebSocket("wss://h/p", undefined); + + // The constructor is called synchronously inside connectWebSocketOnce; emit + // `open` on the next tick so the listener is in place. + await new Promise((r) => setImmediate(r)); + fakeWsInstances[0].emit("open"); + + const ws = await promise; + expect(ws).toBe(fakeWsInstances[0]); + expect(wsConstructorMock).toHaveBeenCalledTimes(1); + }); + + it("retries on 502 and resolves on the second attempt", async () => { + const promise = openPtyWebSocket("wss://h/p", "token"); + + await new Promise((r) => setImmediate(r)); + fakeWsInstances[0].emit("unexpected-response", null, { + statusCode: 502, + statusMessage: "Bad Gateway", + }); + + // Back-off uses a real timer; allow up to ~1s for the first retry. + for (let i = 0; i < 50 && fakeWsInstances.length < 2; i++) { + await new Promise((r) => setTimeout(r, 20)); + } + expect(fakeWsInstances).toHaveLength(2); + + fakeWsInstances[1].emit("open"); + const ws = await promise; + expect(ws).toBe(fakeWsInstances[1]); + }); + + it("does not retry on 401", async () => { + const promise = openPtyWebSocket("wss://h/p", "token"); + + await new Promise((r) => setImmediate(r)); + fakeWsInstances[0].emit("unexpected-response", null, { + statusCode: 401, + statusMessage: "Unauthorized", + }); + + await expect(promise).rejects.toThrow(/401/); + expect(fakeWsInstances).toHaveLength(1); + }); + + it("does not match 502 / 503 buried in unrelated error text", async () => { + // Regression guard for the old loose `msg.includes("502")` fallback — + // an unrelated message must not trigger a retry just because it contains + // those digits. + const promise = openPtyWebSocket("wss://h/p", undefined); + + await new Promise((r) => setImmediate(r)); + fakeWsInstances[0].emit( + "error", + new Error("ECONNRESET while reading 50234 bytes"), + ); + + await expect(promise).rejects.toThrow(/ECONNRESET/); + expect(fakeWsInstances).toHaveLength(1); + }); + + it("passes the auth token via Sec-WebSocket-Protocol (not the URL)", async () => { + const promise = openPtyWebSocket("wss://h/p", "secret-token"); + + await new Promise((r) => setImmediate(r)); + expect(wsConstructorMock).toHaveBeenCalledWith("wss://h/p", [ + "secret-token", + ]); + + fakeWsInstances[0].emit("open"); + await promise; + }); + + it("passes an empty protocols list when no token is provided", async () => { + const promise = openPtyWebSocket("wss://h/p", undefined); + + await new Promise((r) => setImmediate(r)); + expect(wsConstructorMock).toHaveBeenCalledWith("wss://h/p", []); + + fakeWsInstances[0].emit("open"); + await promise; + }); +});