1
/**
2
 * TerminalSocket — manages the WebSocket connection to a sandbox terminal.
3
 *
4
 * Protocol (from Cloudflare Sandbox SDK docs):
5
 *   Binary frames  = terminal I/O  (ANSI/VT byte stream, UTF-8)
6
 *   Text frames    = JSON control / status messages
7
 *
8
 * Lifecycle:
9
 *   1. Client opens WS, sets binaryType = "arraybuffer"
10
 *   2. Server may replay buffered PTY output (binary) before ready
11
 *   3. Server sends { type: "ready" }
12
 *   4. Bidirectional binary I/O
13
 *   5. On disconnect the PTY survives; reconnecting replays buffer
14
 */
15
import type { ConnectionState, ServerStatusMessage } from "../../shared/protocol";
16
 
17
export type SocketState = ConnectionState;
18
 
19
export interface TerminalSocketOptions {
20
  /** WebSocket URL including query params (id, session, token). */
21
  url: string;
22
  /** Called for every binary frame (terminal output). */
23
  onOutput: (data: ArrayBuffer) => void;
24
  /** Called when the connection state changes. */
25
  onStateChange: (state: SocketState, error?: Error) => void;
26
  /** Called once when the server sends the "ready" status. */
27
  onReady: () => void;
28
  /** Called when the shell process exits. */
29
  onExit: (code: number, signal?: string) => void;
30
  /**
31
   * Called after a successful (re)connect with session context.
32
   * `resumed: true` means the server replayed a PTY buffer (existing session).
33
   * `resumed: false` means a fresh shell (container may have restarted).
34
   */
35
  onSessionAttached?: (info: { resumed: boolean; isReconnect: boolean }) => void;
36
  /** Enable automatic reconnection (default true). */
37
  reconnect?: boolean;
38
  /** Max reconnect attempts before giving up (default 10). */
39
  maxReconnectAttempts?: number;
40
}
41
 
42
export class TerminalSocket {
43
  private ws: WebSocket | null = null;
44
  private state: SocketState = "disconnected";
45
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
46
  private reconnectAttempts = 0;
47
  private intentionalClose = false;
48
  private readonly encoder = new TextEncoder();
49
  /** Whether we ever reached "connected" state in this socket's lifetime. */
50
  private wasConnectedBefore = false;
51
  /** Whether binary frames arrived before the "ready" message on this connection. */
52
  private receivedBufferBeforeReady = false;
53
 
54
  constructor(private opts: TerminalSocketOptions) {}
55
 
56
  // --------------- public API ---------------
57
 
58
  connect(): void {
59
    if (this.ws) return;
60
    this.intentionalClose = false;
61
    this.reconnectAttempts = 0;
62
    this.openSocket();
63
  }
64
 
65
  disconnect(): void {
66
    this.intentionalClose = true;
67
    this.cancelReconnect();
68
    if (this.ws) {
69
      this.ws.close();
70
      this.ws = null;
71
    }
72
    this.setState("disconnected");
73
  }
74
 
75
  /** Send user keystrokes as binary UTF-8. */
76
  sendInput(data: string): void {
77
    if (this.ws?.readyState === WebSocket.OPEN) {
78
      this.ws.send(this.encoder.encode(data));
79
    }
80
  }
81
 
82
  /** Send a resize control message. Both cols and rows must be positive. */
83
  sendResize(cols: number, rows: number): void {
84
    if (cols < 1 || rows < 1) return;
85
    if (this.ws?.readyState === WebSocket.OPEN) {
86
      this.ws.send(
87
        JSON.stringify({
88
          type: "resize",
89
          cols: Math.round(cols),
90
          rows: Math.round(rows),
91
        }),
92
      );
93
    }
94
  }
95
 
96
  getState(): SocketState {
97
    return this.state;
98
  }
99
 
100
  // --------------- internals ---------------
101
 
102
  private openSocket(): void {
103
    this.receivedBufferBeforeReady = false;
104
    this.setState(this.reconnectAttempts > 0 ? "reconnecting" : "connecting");
105
 
106
    const ws = new WebSocket(this.opts.url);
107
    ws.binaryType = "arraybuffer";
108
 
109
    ws.addEventListener("open", () => {
110
      this.reconnectAttempts = 0;
111
      // Connected at transport level — wait for "ready" status
112
      // before marking state as "connected".
113
    });
114
 
115
    ws.addEventListener("message", (event: MessageEvent) => {
116
      if (event.data instanceof ArrayBuffer) {
117
        // Binary frame: terminal output (may arrive before "ready").
118
        // Track that we received buffer replay — this means the PTY
119
        // session survived and is being resumed (not a fresh shell).
120
        this.receivedBufferBeforeReady = true;
121
        this.opts.onOutput(event.data);
122
        return;
123
      }
124
 
125
      // Text frame: JSON control / status message.
126
      let msg: ServerStatusMessage;
127
      try {
128
        msg = JSON.parse(event.data as string);
129
      } catch {
130
        return; // ignore malformed frames
131
      }
132
 
133
      switch (msg.type) {
134
        case "ready": {
135
          const isReconnect = this.wasConnectedBefore;
136
          const resumed = this.receivedBufferBeforeReady;
137
          this.wasConnectedBefore = true;
138
 
139
          this.setState("connected");
140
          this.opts.onReady();
141
          this.opts.onSessionAttached?.({ resumed, isReconnect });
142
          break;
143
        }
144
 
145
        case "exit":
146
          this.setState("ended");
147
          this.opts.onExit(msg.code, msg.signal);
148
          break;
149
 
150
        case "error":
151
          this.setState("error", new Error(msg.message));
152
          break;
153
      }
154
    });
155
 
156
    ws.addEventListener("close", () => {
157
      this.ws = null;
158
      if (this.intentionalClose || this.state === "ended") return;
159
      this.maybeReconnect();
160
    });
161
 
162
    ws.addEventListener("error", () => {
163
      // The "close" event always fires after "error", so reconnect
164
      // logic is handled there.
165
    });
166
 
167
    this.ws = ws;
168
  }
169
 
170
  private maybeReconnect(): void {
171
    const maxAttempts = this.opts.maxReconnectAttempts ?? 10;
172
    if (this.opts.reconnect === false || this.reconnectAttempts >= maxAttempts) {
173
      this.setState("error", new Error("Connection lost"));
174
      return;
175
    }
176
    this.scheduleReconnect();
177
  }
178
 
179
  private scheduleReconnect(): void {
180
    if (this.reconnectTimer) return; // prevent stacking timers
181
 
182
    const delay = Math.min(1000 * 2 ** this.reconnectAttempts, 30_000);
183
    this.reconnectAttempts++;
184
    this.setState("reconnecting");
185
 
186
    this.reconnectTimer = setTimeout(() => {
187
      this.reconnectTimer = null;
188
      this.openSocket();
189
    }, delay);
190
  }
191
 
192
  private cancelReconnect(): void {
193
    if (this.reconnectTimer) {
194
      clearTimeout(this.reconnectTimer);
195
      this.reconnectTimer = null;
196
    }
197
  }
198
 
199
  private setState(next: SocketState, error?: Error): void {
200
    if (this.state === next && !error) return;
201
    this.state = next;
202
    this.opts.onStateChange(next, error);
203
  }
204
}