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

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

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.
  */