fix(core-backend): improve traces

Kriys94 · Kriys94 · commit 03b0cc2a5aeb · 2025-11-12T19:12:40.000+01:00
diff --git a/packages/core-backend/CHANGELOG.md b/packages/core-backend/CHANGELOG.md
@@ -7,6 +7,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- Improve WebSocket connection lifecycle tracing in `BackendWebSocketService` ([#7101](https://github.com/MetaMask/core/pull/7101))
+  - WebSocket connection duration is now properly reflected in trace span duration instead of only in custom data
+  - Trace all disconnections (both manual and unexpected) to provide complete connection lifecycle visibility in traces
+  - Omit `connectionDuration_ms` from disconnection traces when connection never established (onClose without onOpen)
+- Update `BackendWebSocketService` default exponential backoff options for reconnection ([#7101](https://github.com/MetaMask/core/pull/7101))
+  - Increase default `reconnectDelay` from 500 milliseconds to 10 seconds
+  - Increase default `maxReconnectDelay` from 30 seconds to 60 seconds
+- Simplify WebSocket disconnection code in `BackendWebSocketService` ([#7101](https://github.com/MetaMask/core/pull/7101))
+  - Centralize all disconnection logic in `ws.onclose` handler for single source of truth
+  - Centralize all state changes within `#establishConnection` method - state transitions only occur in `onopen` (CONNECTING → CONNECTED) and `onclose` (any state → DISCONNECTED)
+  - Add `MANUAL_DISCONNECT_CODE` (4999) and `MANUAL_DISCONNECT_REASON` constants to distinguish manual from unexpected disconnects
+
+### Removed
+
+- Remove `BackendWebSocketService Channel Message` trace as it provided no useful performance insights ([#7101](https://github.com/MetaMask/core/pull/7101))
+
 ## [4.0.0]
 
 ### Changed
diff --git a/packages/core-backend/README.md b/packages/core-backend/README.md
@@ -311,7 +311,7 @@ sequenceDiagram
 #### Key Flow Characteristics
 
 1. **Initial Setup**: BackendWebSocketService establishes connection, then AccountActivityService subscribes to selected account. Backend automatically sends a system notification with all chains that are currently up. AccountActivityService tracks these chains internally and notifies TokenBalancesController, which increases polling interval to 5 min
-2. **Chain Status Tracking**: AccountActivityService maintains an internal set of chains that are 'up' based on system notifications. On disconnect/error, it marks all tracked chains as 'down' before clearing the set
+2. **Chain Status Tracking**: AccountActivityService maintains an internal set of chains that are 'up' based on system notifications. On disconnect, it marks all tracked chains as 'down' before clearing the set
 3. **System Notifications**: Backend automatically sends chain status updates (up/down) upon subscription and when status changes. AccountActivityService forwards these to TokenBalancesController, which adjusts polling intervals (up: 5min, down: 30s + immediate fetch)
 4. **User Account Changes**: When users switch accounts, AccountActivityService unsubscribes from old account and subscribes to new account. Backend sends fresh system notification with current chain status for the new account
 5. **Connection Resilience**: On reconnection, AccountActivityService resubscribes to selected account and receives fresh chain status via system notification. Automatic reconnection with exponential backoff
diff --git a/packages/core-backend/src/AccountActivityService.test.ts b/packages/core-backend/src/AccountActivityService.test.ts
@@ -686,11 +686,11 @@ describe('AccountActivityService', () => {
             timestamp: 1760344704595,
           });
 
-          // Publish WebSocket ERROR state event - should flush tracked chains as down
+          // Publish WebSocket DISCONNECTED state event - should flush tracked chains as down
           rootMessenger.publish(
             'BackendWebSocketService:connectionStateChanged',
             {
-              state: WebSocketState.ERROR,
+              state: WebSocketState.DISCONNECTED,
               url: 'ws://test',
               reconnectAttempts: 2,
               timeout: 10000,
@@ -701,7 +701,7 @@ describe('AccountActivityService', () => {
           );
           await completeAsyncOperations(100);
 
-          // Verify that the ERROR state triggered the status change for tracked chains
+          // Verify that the DISCONNECTED state triggered the status change for tracked chains
           expect(statusChangedEventListener).toHaveBeenCalledWith({
             chainIds: ['eip155:1', 'eip155:137', 'eip155:56'],
             status: 'down',
@@ -720,11 +720,11 @@ describe('AccountActivityService', () => {
 
           mocks.getSelectedAccount.mockReturnValue(null);
 
-          // Publish WebSocket ERROR state event without any tracked chains
+          // Publish WebSocket DISCONNECTED state event without any tracked chains
           rootMessenger.publish(
             'BackendWebSocketService:connectionStateChanged',
             {
-              state: WebSocketState.ERROR,
+              state: WebSocketState.DISCONNECTED,
               url: 'ws://test',
               reconnectAttempts: 2,
               timeout: 10000,
diff --git a/packages/core-backend/src/AccountActivityService.ts b/packages/core-backend/src/AccountActivityService.ts
@@ -451,6 +451,15 @@ export class AccountActivityService {
       status: data.status,
       timestamp,
     });
+
+    log(
+      `WebSocket status change - Published tracked chains as ${data.status}`,
+      {
+        count: data.chainIds.length,
+        chains: data.chainIds,
+        status: data.status,
+      },
+    );
   }
 
   /**
@@ -467,11 +476,8 @@ export class AccountActivityService {
       // WebSocket connected - resubscribe to selected account
       // The system notification will automatically provide the list of chains that are up
       await this.#subscribeToSelectedAccount();
-    } else if (
-      state === WebSocketState.DISCONNECTED ||
-      state === WebSocketState.ERROR
-    ) {
-      // On disconnect/error, flush all tracked chains as down
+    } else if (state === WebSocketState.DISCONNECTED) {
+      // On disconnect, flush all tracked chains as down
       const chainsToMarkDown = Array.from(this.#chainsUp);
 
       if (chainsToMarkDown.length > 0) {
@@ -481,13 +487,10 @@ export class AccountActivityService {
           timestamp: Date.now(),
         });
 
-        log(
-          'WebSocket error/disconnection - Published tracked chains as down',
-          {
-            count: chainsToMarkDown.length,
-            chains: chainsToMarkDown,
-          },
-        );
+        log('WebSocket disconnection - Published tracked chains as down', {
+          count: chainsToMarkDown.length,
+          chains: chainsToMarkDown,
+        });
 
         // Clear the tracking set since all chains are now down
         this.#chainsUp.clear();
diff --git a/packages/core-backend/src/BackendWebSocketService.test.ts b/packages/core-backend/src/BackendWebSocketService.test.ts
@@ -102,8 +102,13 @@ class MockWebSocket extends EventTarget {
     MockWebSocket.instanceCount += 1;
     this.url = url;
     // TypeScript has issues with jest.spyOn on WebSocket methods, so using direct assignment
+    // Store reference to simulateClose for use in close()
+    const simulateCloseFn = this.simulateClose.bind(this);
     // eslint-disable-next-line jest/prefer-spy-on
-    this.close = jest.fn().mockImplementation();
+    this.close = jest.fn().mockImplementation((code = 1000, reason = '') => {
+      // When close() is called, trigger the close event to simulate real WebSocket behavior
+      simulateCloseFn(code, reason);
+    });
     // eslint-disable-next-line jest/prefer-spy-on
     this.send = jest.fn().mockImplementation((data: string) => {
       this._lastSentMessage = data;
@@ -506,8 +511,8 @@ describe('BackendWebSocketService', () => {
         async ({ service }) => {
           expect(service.getConnectionInfo().url).toBe('ws://test.example.com');
           expect(service.getConnectionInfo().timeout).toBe(10000);
-          expect(service.getConnectionInfo().reconnectDelay).toBe(500);
-          expect(service.getConnectionInfo().maxReconnectDelay).toBe(5000);
+          expect(service.getConnectionInfo().reconnectDelay).toBe(10000);
+          expect(service.getConnectionInfo().maxReconnectDelay).toBe(60000);
           expect(service.getConnectionInfo().requestTimeout).toBe(30000);
           expect(service).toBeInstanceOf(BackendWebSocketService);
         },
@@ -877,14 +882,19 @@ describe('BackendWebSocketService', () => {
           mockWebSocketOptions: { autoConnect: false },
         },
         async ({ service, completeAsyncOperations }) => {
+          // Mock Math.random to make Cockatiel's jitter deterministic
+          jest.spyOn(Math, 'random').mockReturnValue(0);
+
           // eslint-disable-next-line @typescript-eslint/no-floating-promises
           service.connect();
 
           // Advance time past the timeout
-          await completeAsyncOperations(150);
+          await completeAsyncOperations(110);
 
-          // Should have transitioned to ERROR state after timeout
-          expect(service.getConnectionInfo().state).toBe(WebSocketState.ERROR);
+          // Should have transitioned to DISCONNECTED state after timeout
+          expect(service.getConnectionInfo().state).toBe(
+            WebSocketState.DISCONNECTED,
+          );
         },
       );
     });
@@ -939,13 +949,15 @@ describe('BackendWebSocketService', () => {
           mockWs.simulateClose(1006, 'Connection failed');
           await completeAsyncOperations(0);
 
-          // Should schedule reconnect and be in ERROR state
-          expect(service.getConnectionInfo().state).toBe(WebSocketState.ERROR);
+          // Should schedule reconnect and be in DISCONNECTED state
+          expect(service.getConnectionInfo().state).toBe(
+            WebSocketState.DISCONNECTED,
+          );
         },
       );
     });
 
-    it('should clear connection timeout in handleClose when timeout occurs then close fires', async () => {
+    it('should clear connection timeout when timeout occurs then close fires', async () => {
       await withService(
         {
           options: { timeout: 100 },
@@ -965,24 +977,23 @@ describe('BackendWebSocketService', () => {
             WebSocketState.CONNECTING,
           );
 
-          // Let timeout fire (closes WebSocket and sets state to ERROR)
-          await completeAsyncOperations(150);
+          // Let timeout fire (closes WebSocket and sets state to DISCONNECTED)
+          // Advance time past timeout but before reconnect would fire
+          await completeAsyncOperations(100);
 
-          // State should be ERROR or DISCONNECTED after timeout
-          const stateAfterTimeout = service.getConnectionInfo().state;
-          expect([WebSocketState.ERROR, WebSocketState.DISCONNECTED]).toContain(
-            stateAfterTimeout,
+          // State should be DISCONNECTED after timeout
+          expect(service.getConnectionInfo().state).toBe(
+            WebSocketState.DISCONNECTED,
           );
 
           // Now manually trigger close event
-          // Since state is ERROR (not CONNECTING), onclose will call handleClose
-          // which will clear connectionTimeout
+          // Since state is DISCONNECTED, onclose will return early due to idempotency guard
           mockWs.simulateClose(1006, 'Close after timeout');
           await completeAsyncOperations(0);
 
-          // State should still be ERROR or DISCONNECTED
-          expect([WebSocketState.ERROR, WebSocketState.DISCONNECTED]).toContain(
-            service.getConnectionInfo().state,
+          // State should still be DISCONNECTED
+          expect(service.getConnectionInfo().state).toBe(
+            WebSocketState.DISCONNECTED,
           );
         },
       );
@@ -1072,10 +1083,13 @@ describe('BackendWebSocketService', () => {
           await completeAsyncOperations(10);
 
           const mockWs = getMockWebSocket();
-          mockWs.simulateError();
+          // Simulate connection failure (error is always followed by close in real WebSocket)
+          mockWs.simulateClose(1006, 'Connection failed');
           await completeAsyncOperations(0);
 
           const attemptsBefore = service.getConnectionInfo().reconnectAttempts;
+          // Should be 1 after the failure
+          expect(attemptsBefore).toBe(1);
 
           // Try to force reconnection while timer is already scheduled
           await service.forceReconnection();
@@ -1087,6 +1101,134 @@ describe('BackendWebSocketService', () => {
         },
       );
     });
+
+    it('should clear reconnect timer when feature is disabled', async () => {
+      let isEnabled = true;
+      await withService(
+        {
+          options: {
+            isEnabled: () => isEnabled,
+          },
+          mockWebSocketOptions: { autoConnect: false },
+        },
+        async ({ service, getMockWebSocket, completeAsyncOperations }) => {
+          // Mock Math.random to make Cockatiel's jitter deterministic
+          jest.spyOn(Math, 'random').mockReturnValue(0);
+
+          // Trigger a connection failure to schedule a reconnect
+          // eslint-disable-next-line @typescript-eslint/no-floating-promises
+          service.connect();
+          await completeAsyncOperations(10);
+
+          const mockWs = getMockWebSocket();
+          // Simulate connection failure to trigger reconnect timer
+          mockWs.simulateClose(1006, 'Connection failed');
+          await completeAsyncOperations(0);
+
+          // Verify reconnect timer is scheduled
+          expect(service.getConnectionInfo().reconnectAttempts).toBe(1);
+
+          // Now disable the feature
+          isEnabled = false;
+
+          // Try to connect again with feature disabled
+          await service.connect();
+
+          // Reconnect attempts should be reset to 0 (timers cleared)
+          expect(service.getConnectionInfo().reconnectAttempts).toBe(0);
+
+          // Advance time to ensure the old reconnect timer doesn't fire
+          await completeAsyncOperations(10000);
+
+          // Should still be 0, confirming timer was cleared
+          expect(service.getConnectionInfo().reconnectAttempts).toBe(0);
+        },
+      );
+    });
+
+    it('should include connectionDuration_ms in trace when connection was established', async () => {
+      const mockTraceFn = jest.fn((_request, fn) => fn?.());
+      await withService(
+        {
+          options: { traceFn: mockTraceFn },
+        },
+        async ({ service, getMockWebSocket, completeAsyncOperations }) => {
+          // Connect and let it establish
+          await service.connect();
+          await completeAsyncOperations(10);
+
+          expect(service.getConnectionInfo().state).toBe(
+            WebSocketState.CONNECTED,
+          );
+
+          // Clear previous trace calls to focus on disconnection trace
+          mockTraceFn.mockClear();
+
+          // Trigger unexpected close after connection was established
+          const mockWs = getMockWebSocket();
+          mockWs.simulateClose(1006, 'Abnormal closure');
+          await completeAsyncOperations(10);
+
+          // Find the Disconnection trace call
+          const disconnectionTrace = mockTraceFn.mock.calls.find(
+            (call) => call[0]?.name === 'BackendWebSocketService Disconnection',
+          );
+
+          expect(disconnectionTrace).toBeDefined();
+          expect(disconnectionTrace?.[0]?.data).toHaveProperty(
+            'connectionDuration_ms',
+          );
+          expect(
+            disconnectionTrace?.[0]?.data?.connectionDuration_ms,
+          ).toBeGreaterThan(0);
+        },
+      );
+    });
+
+    it('should omit connectionDuration_ms in trace when connection never established', async () => {
+      const mockTraceFn = jest.fn((_request, fn) => fn?.());
+      await withService(
+        {
+          options: { traceFn: mockTraceFn },
+          mockWebSocketOptions: { autoConnect: false },
+        },
+        async ({ service, getMockWebSocket, completeAsyncOperations }) => {
+          // Start connecting
+          // eslint-disable-next-line @typescript-eslint/no-floating-promises
+          service.connect();
+          await completeAsyncOperations(0);
+
+          expect(service.getConnectionInfo().state).toBe(
+            WebSocketState.CONNECTING,
+          );
+
+          // Manually disconnect before connection establishes
+          // This will set state to DISCONNECTING and close the WebSocket
+          service.disconnect();
+          await completeAsyncOperations(0);
+
+          // Clear trace calls from connection attempt
+          mockTraceFn.mockClear();
+
+          // Simulate the close event after disconnect (state is now DISCONNECTING, not CONNECTING)
+          // This will trigger #handleClose with connectedAt = 0
+          const mockWs = getMockWebSocket();
+          mockWs.simulateClose(1000, 'Normal closure');
+          await completeAsyncOperations(10);
+
+          // Find the Disconnection trace call
+          const disconnectionTrace = mockTraceFn.mock.calls.find(
+            (call) => call[0]?.name === 'BackendWebSocketService Disconnection',
+          );
+
+          // Trace should exist but should NOT have connectionDuration_ms
+          expect(disconnectionTrace).toBeDefined();
+          expect(disconnectionTrace?.[0]?.data).not.toHaveProperty(
+            'connectionDuration_ms',
+          );
+        },
+      );
+    });
   });
 
   // =====================================================
@@ -1437,7 +1579,9 @@ describe('BackendWebSocketService', () => {
 
         service.disconnect();
 
-        await expect(requestPromise).rejects.toThrow('WebSocket disconnected');
+        await expect(requestPromise).rejects.toThrow(
+          'WebSocket connection closed: 4999 Internal: Manual disconnect',
+        );
       });
     });
 
@@ -1840,7 +1984,9 @@ describe('BackendWebSocketService', () => {
         service.destroy();
 
         // Pending request should be rejected
-        await expect(requestPromise).rejects.toThrow('Service cleanup');
+        await expect(requestPromise).rejects.toThrow(
+          'WebSocket connection closed: 4999 Internal: Manual disconnect',
+        );
 
         // Verify service is in disconnected state
         expect(service.getConnectionInfo().state).toBe(
diff --git a/packages/core-backend/src/BackendWebSocketService.ts b/packages/core-backend/src/BackendWebSocketService.ts
