fix(extension): harden websocket reconnect lifecycle (#82)

Co-authored-by: NianJiuZst <3235467914@qq.com>

apps/extension/src/lib/__tests__/connection-controller.test.ts

@@ -1,7 +1,7 @@
-import { beforeEach, describe, expect, it, vi } from "vitest";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { MIN_COMPATIBLE_PROTOCOL } from "../../transport/handshake";
-import type { ConnectionStateHandler, Transport } from "../../transport/transport";
-import type { ConnectionState, HandshakeResult } from "../../transport/types";
+import type { ConnectionStateHandler, FrameHandler, Transport } from "../../transport/transport";
+import type { ConnectionState, HandshakeResult, ProtocolFrame } from "../../transport/types";
 import { __testing__, ConnectionController } from "../connection-controller";
 
 vi.mock("../instance-id", () => ({
@@ -97,6 +97,7 @@ describe("computeConnectedState (protocol-based compat)", () => {
 function makeMockTransport(initialState: ConnectionState = "disconnected") {
   let state = initialState;
   const stateHandlers = new Set<ConnectionStateHandler>();
+  const messageHandlers = new Set<FrameHandler>();
   const transport = {
     get state() {
       return state;
@@ -110,7 +111,10 @@ function makeMockTransport(initialState: ConnectionState = "disconnected") {
       for (const h of stateHandlers) h("disconnected");
     }),
     send: vi.fn(),
-    onMessage: vi.fn(() => ({ dispose: () => {} })),
+    onMessage: vi.fn((handler: FrameHandler) => {
+      messageHandlers.add(handler);
+      return { dispose: () => messageHandlers.delete(handler) };
+    }),
     onConnectionStateChange: vi.fn((handler: ConnectionStateHandler) => {
       stateHandlers.add(handler);
       return { dispose: () => stateHandlers.delete(handler) };
@@ -119,6 +123,9 @@ function makeMockTransport(initialState: ConnectionState = "disconnected") {
       state = next;
       for (const h of stateHandlers) h(next);
     },
+    emitMessage(frame: ProtocolFrame) {
+      for (const handler of messageHandlers) handler(frame);
+    },
   };
   return transport as typeof transport & Transport;
 }
@@ -128,6 +135,10 @@ describe("ConnectionController connectionEnabled", () => {
     vi.clearAllMocks();
   });
 
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
   it("does not connect on attach when connection is disabled", async () => {
     const controller = new ConnectionController();
     const transport = makeMockTransport();
@@ -201,4 +212,41 @@ describe("ConnectionController connectionEnabled", () => {
 
     expect(controller.snapshot().state).toBe("disconnected");
   });
+
+  it("disconnects and retries when handshake fails while the socket is still open", async () => {
+    vi.useFakeTimers();
+    const controller = new ConnectionController();
+    const transport = makeMockTransport();
+    await controller.attach(transport, { name: "Chrome", version: "120" }, true);
+    const request = transport.send.mock.calls[0]?.[0] as { id: string };
+
+    transport.emitMessage({
+      id: request.id,
+      error: { code: "protocol_error", message: "bad handshake" },
+    });
+    await vi.waitFor(() => expect(transport.disconnect).toHaveBeenCalledTimes(1));
+    expect(controller.snapshot().state).toBe("disconnected");
+
+    await vi.advanceTimersByTimeAsync(1_000);
+    expect(transport.connect).toHaveBeenCalledTimes(2);
+  });
+
+  it("binds each handshake to the connection that initiated it", async () => {
+    const controller = new ConnectionController();
+    const transport = makeMockTransport();
+    await controller.attach(transport, { name: "Chrome", version: "120" }, true);
+    const first = transport.send.mock.calls[0]?.[0] as { id: string };
+
+    transport.emitState("disconnected");
+    transport.emitState("connected");
+    const second = transport.send.mock.calls[1]?.[0] as { id: string };
+    expect(second.id).not.toBe(first.id);
+
+    transport.emitMessage({ id: first.id, result: handshake("1.0", "1.0") });
+    await Promise.resolve();
+    expect(controller.snapshot().state).not.toBe("connected");
+
+    transport.emitMessage({ id: second.id, result: handshake("1.0", "1.0") });
+    await vi.waitFor(() => expect(controller.snapshot().state).toBe("connected"));
+  });
 });