fix(ws): recover terminal replay after silent disconnects

Author: pallyoung

packages/server/src/__tests__/ws-hub.test.ts

@@ -13,6 +13,7 @@ import type { FastifyRequest } from "fastify";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import WebSocket from "ws";
 import { EventBus } from "../bus/event-bus.js";
+import "../commands/connection.js";
 import "../commands/terminal.js";
 import { clearPendingTerminalInput } from "../commands/terminal.js";
 import type { ServerConfig } from "../config.js";
@@ -34,6 +35,7 @@ type MockSocket = {
 
 type MessageHandler = ((data: Buffer, isBinary?: boolean) => void) | undefined;
 type CloseHandler = (() => void) | undefined;
+type PongHandler = (() => void) | undefined;
 
 type ResultMessage = Extract<ServerToClient, { kind: "result" }>;
 
@@ -90,6 +92,11 @@ const getCloseHandler = (socket: MockSocket): CloseHandler =>
     | CloseHandler
     | undefined;
 
+const getPongHandler = (socket: MockSocket): PongHandler =>
+  socket.on.mock.calls.find((call: unknown[]) => call[0] === "pong")?.[1] as
+    | PongHandler
+    | undefined;
+
 const subscribeToAllTopics = (socket: MockSocket) => {
   const messageHandler = getMessageHandler(socket);
 
@@ -435,6 +442,86 @@ describe("WsHub", () => {
     expect(socket.ping).toHaveBeenCalled();
   });
 
+  it("handles connection.probe commands", async () => {
+    const socket = createMockSocket();
+    hub.handleConnection(socket as never, createMockRequest());
+    const messageHandler = getMessageHandler(socket);
+    socket.send.mockClear();
+
+    messageHandler?.(
+      Buffer.from(
+        JSON.stringify({
+          kind: "command",
+          id: "00000000-0000-4000-8000-000000000001",
+          op: "connection.probe",
+          args: {},
+        })
+      )
+    );
+
+    await Promise.resolve();
+    await new Promise((resolve) => setImmediate(resolve));
+
+    expect(findResultMessage(socket, "00000000-0000-4000-8000-000000000001")).toMatchObject({
+      ok: true,
+      data: {
+        ok: true,
+      },
+    });
+  });
+
+  it("closes clients that stay unresponsive across keepalive sweeps", () => {
+    const socket = createMockSocket();
+    hub.handleConnection(socket as never, createMockRequest());
+
+    hub.pingAll();
+    expect(socket.ping).toHaveBeenCalledTimes(1);
+    expect(socket.close).not.toHaveBeenCalled();
+
+    hub.pingAll();
+    expect(socket.close).toHaveBeenCalledTimes(1);
+  });
+
+  it("does not close clients that answer the previous keepalive ping", () => {
+    const socket = createMockSocket();
+    hub.handleConnection(socket as never, createMockRequest());
+    const pongHandler = getPongHandler(socket);
+
+    hub.pingAll();
+    pongHandler?.();
+    hub.pingAll();
+
+    expect(socket.ping).toHaveBeenCalledTimes(2);
+    expect(socket.close).not.toHaveBeenCalled();
+  });
+
+  it("treats inbound commands as proof of life before the next keepalive sweep", async () => {
+    const socket = createMockSocket();
+    hub.handleConnection(socket as never, createMockRequest());
+    const messageHandler = getMessageHandler(socket);
+
+    hub.pingAll();
+
+    messageHandler?.(
+      Buffer.from(
+        JSON.stringify({
+          kind: "command",
+          id: "00000000-0000-4000-8000-000000000002",
+          op: "connection.probe",
+          args: {},
+        })
+      )
+    );
+
+    await Promise.resolve();
+    await new Promise((resolve) => setImmediate(resolve));
+
+    hub.pingAll();
+
+    expect(socket.close).not.toHaveBeenCalled();
+    expect(socket.ping).toHaveBeenCalledTimes(2);
+  });
+
   it("should return null for writer (deprecated - use FencingManager)", () => {
     const socket = createMockSocket();
     hub.handleConnection(socket as never, createMockRequest());

packages/server/src/commands/connection.ts

@@ -0,0 +1,6 @@
+import { z } from "zod";
+import { registerCommand } from "../ws/dispatch.js";
+
+registerCommand("connection.probe", z.object({}).default({}), async () => {
+  return { ok: true as const };
+});

packages/server/src/commands/index.ts

@@ -6,6 +6,7 @@
 
 import "./workspace.js";
 import "./workspace-activity.js";
+import "./connection.js";
 import "./session.js";
 import "./terminal.js";
 import "./file.js";

packages/server/src/server.ts

@@ -44,6 +44,8 @@ import { WsHub } from "./ws/hub.js";
 
 import "./commands/index.js";
 
+const WS_KEEPALIVE_INTERVAL_MS = 15_000;
+
 export interface Server {
   app: FastifyInstance;
   stop: () => Promise<void>;
@@ -235,12 +237,18 @@ export async function createServer(
   }, STARTUP_GC_DELAY_MS);
   gcTimer.unref();
 
+  const wsKeepaliveTimer = setInterval(() => {
+    wsHub.pingAll();
+  }, WS_KEEPALIVE_INTERVAL_MS);
+  wsKeepaliveTimer.unref();
+
   let stopped = false;
   const stopServer = async () => {
     if (stopped) return;
     stopped = true;
 
     clearTimeout(gcTimer);
+    clearInterval(wsKeepaliveTimer);
     await app.close();
     autoFetch.stop();
     supervisorMgr.stop();

packages/server/src/ws/client.ts

@@ -64,8 +64,14 @@ export class WsClient {
     this.setupSocketHandlers();
   }
 
+  private markAlive(): void {
+    this.isAlive = true;
+  }
+
   private setupSocketHandlers(): void {
     this.socket.on("message", (data: Buffer, isBinary: boolean) => {
+      this.markAlive();
+
       if (isBinary) {
         this.messageHandler?.(data);
         return;
@@ -87,7 +93,7 @@ export class WsClient {
     });
 
     this.socket.on("pong", () => {
-      this.isAlive = true;
+      this.markAlive();
     });
   }
 

packages/server/src/ws/hub.ts

@@ -320,6 +320,10 @@ export class WsHub implements Broadcaster {
    */
   pingAll(): void {
     for (const client of this.clients.values()) {
+      if (!client.alive) {
+        client.close(1011, "keepalive_timeout");
+        continue;
+      }
       client.ping();
     }
   }

packages/web/src/features/terminal-panel/__tests__/xterm-host.test.tsx

@@ -4449,6 +4449,78 @@ describe("XtermHost", () => {
     });
   });
 
+  it("replays terminal history when websocket recovery succeeds without a status transition", async () => {
+    const store = createStore();
+    const initialReplayChunk = new TextEncoder().encode("initial replay\n");
+    const recoveredReplayChunk = new TextEncoder().encode("recovered after probe\n");
+    let recoveryHandler:
+      | ((trigger: "visibility_resume" | "network_online" | "manual_retry" | "reconnected") => void)
+      | undefined;
+    let replayCount = 0;
+    const sendCommand = vi
+      .fn()
+      .mockImplementation((op: string, args: { terminalId?: string; lastSeq?: number }) => {
+        if (op !== "terminal.replay") {
+          return Promise.resolve({ status: "ok" });
+        }
+
+        replayCount += 1;
+        if (replayCount === 1) {
+          expect(args.lastSeq).toBe(0);
+          return Promise.resolve({
+            status: "ok",
+            transport: "binary",
+            streamId: 980,
+            size: initialReplayChunk.byteLength,
+            seq: 100,
+            bytes: initialReplayChunk,
+          } satisfies TerminalReplayPayload);
+        }
+
+        expect(args.lastSeq).toBe(100);
+        return Promise.resolve({
+          status: "ok",
+          transport: "binary",
+          streamId: 981,
+          size: recoveredReplayChunk.byteLength,
+          seq: 126,
+          bytes: recoveredReplayChunk,
+        } satisfies TerminalReplayPayload);
+      });
+
+    store.set(wsClientAtom, {
+      sendCommand,
+      subscribe: vi.fn(() => vi.fn()),
+      getStatus: vi.fn(() => "connected"),
+      onStatus: vi.fn(() => () => {}),
+      onRecovery: vi.fn((handler: typeof recoveryHandler) => {
+        recoveryHandler = handler;
+        return () => {};
+      }),
+    } as never);
+
+    render(
+      <Provider store={store}>
+        <XtermHost terminalId="probe-recovery-terminal" workspaceId="test-workspace" />
+      </Provider>
+    );
+
+    await waitFor(() => {
+      expectTerminalWriteData(initialReplayChunk);
+    });
+
+    await act(async () => {
+      recoveryHandler?.("network_online");
+      await Promise.resolve();
+      await Promise.resolve();
+    });
+
+    await waitFor(() => {
+      expect(replayCount).toBe(2);
+      expectTerminalWriteData(recoveredReplayChunk);
+    });
+  });
+
   it("waits for reconnect replay bytes to finish rendering before starting another reconnect recovery", async () => {
     const store = createStore();
     const initialReplayChunk = new TextEncoder().encode("initial replay\n");
@@ -4573,6 +4645,125 @@ describe("XtermHost", () => {
     });
   });
 
+  it("queues probe-triggered recovery that arrives while historical recovery writes are still in flight", async () => {
+    const store = createStore();
+    const initialReplayChunk = new TextEncoder().encode("initial replay\n");
+    const delayedReconnectReplay = new TextEncoder().encode("delayed reconnect replay\n");
+    const queuedProbeReplay = new TextEncoder().encode("queued probe recovery\n");
+    let recoveryHandler:
+      | ((trigger: "visibility_resume" | "network_online" | "manual_retry" | "reconnected") => void)
+      | undefined;
+    let replayCount = 0;
+    let releaseDelayedReconnectWrite: (() => void) | undefined;
+
+    mockTerminal.write.mockImplementation((data: Uint8Array | string, callback?: () => void) => {
+      if (data === delayedReconnectReplay) {
+        releaseDelayedReconnectWrite = callback;
+        return;
+      }
+      callback?.();
+    });
+
+    const sendCommand = vi
+      .fn()
+      .mockImplementation((op: string, args: { terminalId?: string; lastSeq?: number }) => {
+        if (op !== "terminal.replay") {
+          return Promise.resolve({ status: "ok" });
+        }
+
+        replayCount += 1;
+        if (replayCount === 1) {
+          expect(args.lastSeq).toBe(0);
+          return Promise.resolve({
+            status: "ok",
+            transport: "binary",
+            streamId: 982,
+            size: initialReplayChunk.byteLength,
+            seq: 100,
+            bytes: initialReplayChunk,
+          } satisfies TerminalReplayPayload);
+        }
+
+        if (replayCount === 2) {
+          expect(args.lastSeq).toBe(100);
+          return Promise.resolve({
+            status: "ok",
+            transport: "binary",
+            streamId: 983,
+            size: delayedReconnectReplay.byteLength,
+            seq: 200,
+            bytes: delayedReconnectReplay,
+          } satisfies TerminalReplayPayload);
+        }
+
+        expect(args.lastSeq).toBe(200);
+        return Promise.resolve({
+          status: "ok",
+          transport: "binary",
+          streamId: 984,
+          size: queuedProbeReplay.byteLength,
+          seq: 240,
+          bytes: queuedProbeReplay,
+        } satisfies TerminalReplayPayload);
+      });
+
+    mockTerminal.cols = 132;
+    mockTerminal.rows = 36;
+
+    store.set(wsClientAtom, {
+      sendCommand,
+      subscribe: vi.fn(() => vi.fn()),
+      getStatus: vi.fn(() => "connected"),
+      onStatus: vi.fn(() => () => {}),
+      onRecovery: vi.fn((handler: typeof recoveryHandler) => {
+        recoveryHandler = handler;
+        return () => {};
+      }),
+    } as never);
+
+    render(
+      <Provider store={store}>
+        <XtermHost terminalId="queued-probe-recovery-terminal" workspaceId="test-workspace" />
+      </Provider>
+    );
+
+    await waitFor(() => {
+      expectTerminalWriteData(initialReplayChunk);
+    });
+
+    await act(async () => {
+      recoveryHandler?.("network_online");
+      await Promise.resolve();
+      await Promise.resolve();
+    });
+
+    await waitFor(() => {
+      expect(replayCount).toBe(2);
+      expectTerminalWriteData(delayedReconnectReplay);
+    });
+
+    expect(typeof releaseDelayedReconnectWrite).toBe("function");
+
+    await act(async () => {
+      recoveryHandler?.("visibility_resume");
+      await Promise.resolve();
+      await Promise.resolve();
+    });
+
+    expect(replayCount).toBe(2);
+
+    await act(async () => {
+      releaseDelayedReconnectWrite?.();
+      await Promise.resolve();
+      await Promise.resolve();
+    });
+
+    await waitFor(() => {
+      expect(replayCount).toBe(3);
+      expectTerminalWriteData(queuedProbeReplay);
+    });
+  });
+
   it("uses the flushed pending chunk seq for a later reconnect after a successful recovery", async () => {
     const store = createStore();
     const initialReplayChunk = new TextEncoder().encode("initial replay\n");