fix(relay): avoid stale client timer resetting control socket

2026-02-13 19:01:51 +07:00
parent 67620569ab
commit c41815dcb7
2 changed files with 106 additions and 0 deletions
@@ -0,0 +1,96 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { RelayDurableObject } from "./cloudflare-adapter.js";
+
+type MockSocket = WebSocket & {
+  send: ReturnType<typeof vi.fn>;
+  close: ReturnType<typeof vi.fn>;
+  serializeAttachment: ReturnType<typeof vi.fn>;
+  deserializeAttachment: ReturnType<typeof vi.fn>;
+};
+
+function createMockSocket(attachment: unknown = null): MockSocket {
+  let storedAttachment = attachment;
+  return {
+    send: vi.fn(),
+    close: vi.fn(),
+    serializeAttachment: vi.fn((value: unknown) => {
+      storedAttachment = value;
+    }),
+    deserializeAttachment: vi.fn(() => storedAttachment),
+  } as unknown as MockSocket;
+}
+
+function createMockState() {
+  const socketsByTag = new Map<string, WebSocket[]>();
+  const state = {
+    acceptWebSocket: vi.fn(),
+    getWebSockets: vi.fn((tag?: string): WebSocket[] => {
+      if (!tag) {
+        const out: WebSocket[] = [];
+        for (const sockets of socketsByTag.values()) out.push(...sockets);
+        return out;
+      }
+      return socketsByTag.get(tag) ?? [];
+    }),
+  };
+
+  return {
+    state,
+    setTagSockets: (tag: string, sockets: WebSocket[]) => {
+      socketsByTag.set(tag, sockets);
+    },
+  };
+}
+
+describe("RelayDurableObject control nudge/reset behavior", () => {
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it("does not nudge or reset control after the client already disconnected", () => {
+    vi.useFakeTimers();
+    const clientId = "clt_stale_timer";
+    const control = createMockSocket();
+    const { state, setTagSockets } = createMockState();
+
+    setTagSockets("server-control", [control]);
+    setTagSockets("client", []);
+    setTagSockets(`client:${clientId}`, []);
+    setTagSockets(`server:${clientId}`, []);
+
+    const relay = new RelayDurableObject(state as any);
+    (relay as any).nudgeOrResetControlForClient(clientId);
+
+    vi.advanceTimersByTime(15_000);
+
+    expect(control.send).not.toHaveBeenCalled();
+    expect(control.close).not.toHaveBeenCalled();
+  });
+
+  it("resets control when the client remains connected but no server-data socket appears", () => {
+    vi.useFakeTimers();
+    const clientId = "clt_waiting_for_daemon";
+    const control = createMockSocket();
+    const client = createMockSocket({
+      role: "client",
+      clientId,
+      serverId: "srv_test",
+      createdAt: Date.now(),
+    });
+    const { state, setTagSockets } = createMockState();
+
+    setTagSockets("server-control", [control]);
+    setTagSockets("client", [client]);
+    setTagSockets(`client:${clientId}`, [client]);
+    setTagSockets(`server:${clientId}`, []);
+
+    const relay = new RelayDurableObject(state as any);
+    (relay as any).nudgeOrResetControlForClient(clientId);
+
+    vi.advanceTimersByTime(10_000);
+    expect(control.send).toHaveBeenCalledTimes(1);
+
+    vi.advanceTimersByTime(5_000);
+    expect(control.close).toHaveBeenCalledWith(1011, "Control unresponsive");
+  });
+});