From ea8425b9037e7239e6cb22e8462dd59f0e6ee04f Mon Sep 17 00:00:00 2001
From: Louis Liu <louisliu@Louiss-MacBook-Air-568.local>
Date: Wed, 29 Jul 2026 16:23:35 +0800
Subject: [PATCH] Add relayer state change notifications

---
 src/relayer/WebSocketRelayer.ts  | 44 +++++++++++++++++++++++++++++++-
 src/tests/relayer-stress.test.ts | 23 +++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/src/relayer/WebSocketRelayer.ts b/src/relayer/WebSocketRelayer.ts
index 3b19e176..110a6e88 100644
--- a/src/relayer/WebSocketRelayer.ts
+++ b/src/relayer/WebSocketRelayer.ts
@@ -5,6 +5,7 @@ export interface WebSocketMessage {
 }
 
 export type MessageHandler = (msg: WebSocketMessage) => Promise<void> | void;
+export type StateChangeHandler = (transition: RelayerStateTransition, state: RelayerState) => void;
 
 export interface RelayerState {
   connected: boolean;
@@ -15,10 +16,17 @@ export interface RelayerState {
 
 export type RelayerStateTransition = 'connecting' | 'connected' | 'disconnected' | 'reconnecting' | 'destroyed';
 
+export interface WebSocketRelayerOptions {
+  maxReconnectAttempts?: number;
+  reconnectDelayMs?: number;
+  onStateChange?: StateChangeHandler;
+}
+
 export class WebSocketRelayer {
   private url: string;
   private ws: WebSocket | null = null;
   private handlers: Map<string, Set<MessageHandler>> = new Map();
+  private stateChangeHandlers: Set<StateChangeHandler> = new Set();
   private isDestroyed = false;
   private isLocked = false;
   private lockQueue: Array<() => void> = [];
@@ -28,11 +36,15 @@ export class WebSocketRelayer {
   private maxReconnectAttempts: number;
   private reconnectDelayMs: number;
   private pendingMessages: WebSocketMessage[] = [];
+  private stateTransition: RelayerStateTransition = 'disconnected';
 
-  constructor(url: string, options?: { maxReconnectAttempts?: number; reconnectDelayMs?: number }) {
+  constructor(url: string, options?: WebSocketRelayerOptions) {
     this.url = url;
     this.maxReconnectAttempts = options?.maxReconnectAttempts ?? 5;
     this.reconnectDelayMs = options?.reconnectDelayMs ?? 1000;
+    if (options?.onStateChange) {
+      this.stateChangeHandlers.add(options.onStateChange);
+    }
   }
 
   get state(): RelayerState {
@@ -63,6 +75,20 @@ export class WebSocketRelayer {
     }
   }
 
+  private emitStateChange(transition: RelayerStateTransition): void {
+    if (this.stateTransition === transition) return;
+
+    this.stateTransition = transition;
+    const state = this.state;
+    for (const handler of this.stateChangeHandlers) {
+      try {
+        handler(transition, state);
+      } catch (err) {
+        console.warn('[WebSocketRelayer] state change handler error:', err);
+      }
+    }
+  }
+
   async connect(): Promise<void> {
     if (this.isDestroyed) {
       throw new Error('WebSocketRelayer has been destroyed');
@@ -78,6 +104,7 @@ export class WebSocketRelayer {
         return;
       }
 
+      this.emitStateChange('connecting');
       this.connectPromise = this.establishConnection();
       await this.connectPromise;
     } finally {
@@ -99,6 +126,7 @@ export class WebSocketRelayer {
   private async establishConnection(): Promise<void> {
     const WebSocketCtor = this.wsCtor();
     if (!WebSocketCtor) {
+      this.emitStateChange('disconnected');
       return;
     }
 
@@ -122,6 +150,7 @@ export class WebSocketRelayer {
           settled = true;
           this.reconnectAttempts = 0;
           this.flushPendingMessages();
+          this.emitStateChange('connected');
           resolve();
         };
 
@@ -135,6 +164,7 @@ export class WebSocketRelayer {
             settled = true;
             reject(new Error(`WebSocket connection closed before opening: ${this.url}`));
           }
+          this.emitStateChange('disconnected');
           if (!this.isDestroyed && this.reconnectEnabled) {
             this.attemptReconnect();
           }
@@ -214,6 +244,7 @@ export class WebSocketRelayer {
       if (!this.shouldAttemptReconnect()) return;
 
       this.reconnectAttempts++;
+      this.emitStateChange('reconnecting');
       await new Promise((r) => setTimeout(r, this.reconnectDelayMs * this.reconnectAttempts));
       if (this.isDestroyed) return;
 
@@ -242,6 +273,14 @@ export class WebSocketRelayer {
     };
   }
 
+  onStateChange(handler: StateChangeHandler): () => void {
+    this.stateChangeHandlers.add(handler);
+
+    return () => {
+      this.stateChangeHandlers.delete(handler);
+    };
+  }
+
   off(type: string, handler: MessageHandler): void {
     const handlers = this.handlers.get(type);
     if (handlers) {
@@ -290,6 +329,7 @@ export class WebSocketRelayer {
       }
       this.ws = null;
     }
+    this.emitStateChange('disconnected');
   }
 
   destroy(): void {
@@ -311,7 +351,9 @@ export class WebSocketRelayer {
       this.ws = null;
     }
 
+    this.emitStateChange('destroyed');
     this.handlers.clear();
+    this.stateChangeHandlers.clear();
     this.lockQueue = [];
     this.isLocked = false;
   }
diff --git a/src/tests/relayer-stress.test.ts b/src/tests/relayer-stress.test.ts
index 012efb14..0182be15 100644
--- a/src/tests/relayer-stress.test.ts
+++ b/src/tests/relayer-stress.test.ts
@@ -273,4 +273,27 @@ describe('WebSocketRelayer — High-Concurrency Stress Tests', () => {
 
     expect(relayer.state.pendingCount).toBe(3);
   });
+
+  it('notifies subscribers when connection state changes', async () => {
+    const handler = vi.fn();
+    const unsubscribe = relayer.onStateChange(handler);
+
+    await relayer.connect();
+    relayer.disconnect();
+
+    expect(handler.mock.calls.map(([transition]) => transition)).toEqual([
+      'connecting',
+      'connected',
+      'disconnected',
+    ]);
+    expect(handler.mock.calls[1]![1]).toMatchObject({
+      connected: true,
+      destroyed: false,
+    });
+
+    unsubscribe();
+    relayer.destroy();
+
+    expect(handler).toHaveBeenCalledTimes(3);
+  });
 });