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"));
+  });
 });

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

@@ -9,6 +9,8 @@ import type { ConnectionState, HandshakeResult } from "../transport/types";
 import { getLabel, getOrCreateInstanceId } from "./instance-id";
 import { compareProtocol, parseProtocolMajor } from "./semver";
 
+const HANDSHAKE_RETRY_DELAY_MS = 1_000;
+
 export interface SnapshotInfo {
   state: ConnectionState;
   instanceId: string;
@@ -43,9 +45,12 @@ export class ConnectionController {
   private lastError: string | null = null;
   private connectionEnabled = true;
   private listeners = new Set<Listener>();
-  private handshakeInFlight = false;
   private lifecycleHooks: ConnectionLifecycleHooks = {};
   private disconnectRecovery: Promise<void> | null = null;
+  private connectionGeneration = 0;
+  private handshakeAbort: AbortController | null = null;
+  private handshakeRetryTimer: ReturnType<typeof setTimeout> | null = null;
+  private suppressDisconnectRecovery = false;
 
   get isConnectionEnabled(): boolean {
     return this.connectionEnabled;
@@ -87,16 +92,25 @@ export class ConnectionController {
 
     transport.onConnectionStateChange((s) => {
       if (s === "disconnected") {
+        // An in-flight handshake belongs to the dead connection — cancel it so
+        // a late resolution can never clobber the next attempt.
+        this.cancelHandshake();
         this.handshake = null;
         if (this.connectionEnabled) {
           this.setState("disconnected");
-          void this.recoverFromDisconnect();
+          if (this.suppressDisconnectRecovery) {
+            // Deliberate disconnect from runHandshake (rejected peer or failed
+            // handshake bounce); the caller owns the reconnect policy there.
+            this.suppressDisconnectRecovery = false;
+          } else {
+            void this.recoverFromDisconnect();
+          }
         }
         return;
       }
       if (!this.connectionEnabled) return;
       if (s === "connected") {
-        void this.runHandshake(browser);
+        this.startHandshake(browser);
         return;
       }
       this.setState(s);
@@ -142,39 +156,66 @@ export class ConnectionController {
     this.fire();
   }
 
-  private async runHandshake(browser: { name: string; version: string }): Promise<void> {
+  private startHandshake(browser: { name: string; version: string }): void {
+    this.cancelHandshake();
+    const generation = ++this.connectionGeneration;
+    const abort = new AbortController();
+    this.handshakeAbort = abort;
+    void this.runHandshake(browser, generation, abort.signal);
+  }
+
+  private async runHandshake(
+    browser: { name: string; version: string },
+    generation: number,
+    signal: AbortSignal,
+  ): Promise<void> {
     if (!this.transport) return;
     if (!this.connectionEnabled) return;
-    if (this.handshakeInFlight) return;
-    this.handshakeInFlight = true;
     this.setState("connecting");
     try {
-      const outcome = await performHandshake(this.transport, {
-        instanceId: this.instanceId,
-        browser,
-        label: this.label,
-      });
+      const outcome = await performHandshake(
+        this.transport,
+        {
+          instanceId: this.instanceId,
+          browser,
+          label: this.label,
+        },
+        { signal },
+      );
+      if (generation !== this.connectionGeneration || signal.aborted) return;
       this.handshake = outcome.result;
       const verdict = computeConnectedState(outcome.result);
       if (verdict.kind === "rejected") {
         this.lastError = `version_too_old: ${verdict.reason}`;
         this.handshake = null;
+        // Rejected peers must not be auto-retried: mark this disconnect as
+        // deliberate so the listener skips unexpected-loss recovery.
+        this.suppressDisconnectRecovery = true;
         await this.transport.disconnect().catch(() => {});
         this.setState("disconnected");
         return;
       }
+      this.clearHandshakeRetry();
       this.lastError = null;
       this.setState(verdict.kind);
     } catch (err) {
+      if (generation !== this.connectionGeneration || signal.aborted || isAbortError(err)) return;
       this.handshake = null;
       this.lastError = err instanceof Error ? err.message : String(err);
+      // Bounce the half-open socket deliberately; the retry timer below owns
+      // the reconnect, so skip unexpected-loss recovery for this disconnect.
+      this.suppressDisconnectRecovery = true;
+      await this.transport.disconnect().catch(() => {});
       this.setState("disconnected");
+      this.scheduleHandshakeRetry(browser);
     } finally {
-      this.handshakeInFlight = false;
+      if (generation === this.connectionGeneration) this.handshakeAbort = null;
     }
   }
 
   private async applyDisabledState(): Promise<void> {
+    this.cancelHandshake();
+    this.clearHandshakeRetry();
     this.handshake = null;
     this.lastError = null;
     await this.lifecycleHooks.beforeDisconnect?.();
@@ -194,6 +235,8 @@ export class ConnectionController {
       // is cancelled before local session teardown begins.
       await Promise.resolve();
       if (!this.connectionEnabled || !this.transport) return;
+      // Another path already reconnected while we yielded — nothing to do.
+      if (this.currentState !== "disconnected") return;
       await this.transport.disconnect().catch(() => {});
       try {
         await this.lifecycleHooks.onDisconnected?.();
@@ -214,6 +257,30 @@ export class ConnectionController {
     return this.disconnectRecovery;
   }
 
+  private cancelHandshake(): void {
+    this.connectionGeneration += 1;
+    this.handshakeAbort?.abort();
+    this.handshakeAbort = null;
+  }
+
+  private scheduleHandshakeRetry(browser: { name: string; version: string }): void {
+    if (this.handshakeRetryTimer || !this.connectionEnabled) return;
+    this.handshakeRetryTimer = setTimeout(() => {
+      this.handshakeRetryTimer = null;
+      if (!this.connectionEnabled || !this.transport) return;
+      void this.transport.connect().catch((err) => {
+        this.lastError = err instanceof Error ? err.message : String(err);
+        this.setState("disconnected");
+      });
+    }, HANDSHAKE_RETRY_DELAY_MS);
+  }
+
+  private clearHandshakeRetry(): void {
+    if (!this.handshakeRetryTimer) return;
+    clearTimeout(this.handshakeRetryTimer);
+    this.handshakeRetryTimer = null;
+  }
+
   private setState(next: ConnectionState): void {
     if (this.currentState === next) return;
     this.currentState = next;
@@ -232,6 +299,12 @@ export class ConnectionController {
   }
 }
 
+function isAbortError(err: unknown): boolean {
+  return (
+    typeof err === "object" && err !== null && (err as { name?: string }).name === "AbortError"
+  );
+}
+
 /**
  * Verdict of the symmetric post-handshake compat check.
  */

apps/extension/src/transport/__tests__/handshake.test.ts

@@ -198,6 +198,25 @@ describe("performHandshake", () => {
       result: { server: "browser-skill-daemon", protocol_version: "1.0" },
     });
   });
+
+  it("stops waiting immediately when the connection generation is aborted", async () => {
+    const { transport } = deferredFakeTransport();
+    const controller = new AbortController();
+    const pending = performHandshake(
+      transport,
+      {
+        instanceId: "x",
+        browser: { name: "chrome", version: "131" },
+        label: "",
+        rpcId: "hs-abort",
+      },
+      { signal: controller.signal, timeoutMs: 60_000 },
+    );
+
+    controller.abort();
+
+    await expect(pending).rejects.toMatchObject({ name: "AbortError" });
+  });
 });
 
 describe("detectBrowserMeta", () => {

apps/extension/src/transport/__tests__/ws-transport.test.ts

@@ -49,6 +49,10 @@ class FakeSocket {
     this.emit("close", { code, reason: "server-gone" });
   }
 
+  emitClose(code = 1006): void {
+    this.emit("close", { code, reason: "delayed-close" });
+  }
+
   private emit(type: string, ev: unknown): void {
     for (const l of this.listeners[type] ?? []) {
       // biome-ignore lint/suspicious/noExplicitAny: minimal fake mirrors WebSocket Event shape
@@ -202,6 +206,34 @@ describe("WSTransport", () => {
     expect(FakeSocket.instances.length).toBe(beforeCount + 1);
   });
 
+  it("cancels pending backoff and ignores delayed events from a replaced socket", async () => {
+    const t = new WSTransport({
+      url: "ws://127.0.0.1:52800",
+      webSocketFactory: (url) => new FakeSocket(url) as unknown as WebSocket,
+    });
+    const initial = t.connect();
+    const first = lastSocket();
+    first.open();
+    await initial;
+
+    first.serverClose();
+    const reconnect = t.connect();
+    expect(FakeSocket.instances).toHaveLength(2);
+    const second = lastSocket();
+    second.open();
+    await reconnect;
+
+    // A real browser can deliver the old close callback after the replacement
+    // socket is already open. It must not clear or disconnect the new socket.
+    first.emitClose();
+    t.send({ id: "current", method: "system.ping" });
+    expect(second.sent).toEqual([JSON.stringify({ id: "current", method: "system.ping" })]);
+    expect(t.state).toBe("connected");
+
+    await vi.advanceTimersByTimeAsync(5_000);
+    expect(FakeSocket.instances).toHaveLength(2);
+  });
+
   it("ignores malformed inbound messages (non-JSON) without throwing", async () => {
     const handler = vi.fn();
     const t = new WSTransport({

apps/extension/src/transport/handshake.ts

@@ -57,7 +57,7 @@ function ridToString(): string {
 export function performHandshake(
   transport: Transport,
   input: HandshakeInput,
-  options: { timeoutMs?: number } = {},
+  options: { timeoutMs?: number; signal?: AbortSignal } = {},
 ): Promise<HandshakeOutcome> {
   const id = input.rpcId ?? ridToString();
   const params: HandshakeParams = {
@@ -75,6 +75,12 @@ export function performHandshake(
   const timeoutMs = options.timeoutMs ?? 10_000;
 
   return new Promise<HandshakeOutcome>((resolve, reject) => {
+    const onAbort = () => {
+      cleanup();
+      const err = new Error("[handshake] aborted because the connection changed");
+      err.name = "AbortError";
+      reject(err);
+    };
     const timer = setTimeout(() => {
       cleanup();
       reject(new Error("[handshake] timed out waiting for daemon response"));
@@ -97,8 +103,15 @@ export function performHandshake(
     function cleanup() {
       clearTimeout(timer);
       sub.dispose();
+      options.signal?.removeEventListener("abort", onAbort);
     }
 
+    if (options.signal?.aborted) {
+      onAbort();
+      return;
+    }
+    options.signal?.addEventListener("abort", onAbort, { once: true });
+
     try {
       transport.send(req);
     } catch (err) {

apps/extension/src/transport/ws-transport.ts

@@ -50,6 +50,7 @@ export class WSTransport implements Transport {
   private explicitlyClosed = false;
   private reconnectAttempt = 0;
   private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
+  private socketGeneration = 0;
   private connectingPromise: Promise<void> | null = null;
   private resolveConnect: ((value: void) => void) | null = null;
   private rejectConnect: ((reason: Error) => void) | null = null;
@@ -69,16 +70,26 @@ export class WSTransport implements Transport {
   }
 
   connect(): Promise<void> {
-    if (this.connectingPromise) return this.connectingPromise;
     if (this.currentState === "connected") return Promise.resolve();
 
     this.explicitlyClosed = false;
-    this.openSocket();
+    if (this.reconnectTimer) {
+      clearTimeout(this.reconnectTimer);
+      this.reconnectTimer = null;
+    }
+
+    if (this.connectingPromise) {
+      // A previous physical attempt closed before reaching OPEN. Keep the
+      // original caller's promise, but start the next socket generation.
+      if (!this.socket) this.openSocket();
+      return this.connectingPromise;
+    }
 
     this.connectingPromise = new Promise<void>((resolve, reject) => {
       this.resolveConnect = resolve;
       this.rejectConnect = reject;
     });
+    this.openSocket();
     return this.connectingPromise;
   }
 
@@ -88,9 +99,12 @@ export class WSTransport implements Transport {
       clearTimeout(this.reconnectTimer);
       this.reconnectTimer = null;
     }
-    if (this.socket) {
+    const socket = this.socket;
+    this.socket = null;
+    this.socketGeneration += 1;
+    if (socket) {
       try {
-        this.socket.close();
+        socket.close();
       } catch {
         // ignore
       }
@@ -135,10 +149,12 @@ export class WSTransport implements Transport {
 
   private openSocket(): void {
     this.setState("connecting");
+    const generation = ++this.socketGeneration;
     const socket = this.factory(this.url);
     this.socket = socket;
 
     socket.addEventListener("open", () => {
+      if (!this.isCurrentSocket(socket, generation)) return;
       this.reconnectAttempt = 0;
       this.setState("connected");
       const resolve = this.resolveConnect;
@@ -149,11 +165,12 @@ export class WSTransport implements Transport {
     });
 
     socket.addEventListener("message", (ev: MessageEvent) => {
+      if (!this.isCurrentSocket(socket, generation)) return;
       this.handleInbound((ev as unknown as MessageLikeEvent).data);
     });
 
     socket.addEventListener("close", (ev: Event) => {
-      this.handleClose(ev as unknown as CloseLikeEvent);
+      this.handleClose(socket, generation, ev as unknown as CloseLikeEvent);
     });
 
     socket.addEventListener("error", () => {
@@ -179,7 +196,8 @@ export class WSTransport implements Transport {
     }
   }
 
-  private handleClose(_ev: CloseLikeEvent): void {
+  private handleClose(socket: WebSocket, generation: number, _ev: CloseLikeEvent): void {
+    if (!this.isCurrentSocket(socket, generation)) return;
     this.socket = null;
     if (this.explicitlyClosed) {
       this.setState("disconnected");
@@ -196,10 +214,16 @@ export class WSTransport implements Transport {
     this.reconnectTimer = setTimeout(() => {
       this.reconnectTimer = null;
       if (this.explicitlyClosed) return;
-      this.openSocket();
+      void this.connect().catch((err) => {
+        console.debug("[WSTransport] reconnect attempt failed", err);
+      });
     }, delay);
   }
 
+  private isCurrentSocket(socket: WebSocket, generation: number): boolean {
+    return this.socket === socket && this.socketGeneration === generation;
+  }
+
   private setState(next: ConnectionState): void {
     if (this.currentState === next) return;
     this.currentState = next;