fix(WebSocketClientManager): clear inMemoryClients on localStorage removal

Author: kettanaito

src/browser/setupWorker/stop/createStop.ts

@@ -1,4 +1,5 @@
 import { devUtils } from '~/core/utils/internal/devUtils'
+import { MSW_WEBSOCKET_CLIENTS_KEY } from '~/core/ws/WebSocketClientManager'
 import { SetupWorkerInternalContext, StopHandler } from '../glossary'
 import { printStopMessage } from './utils/printStopMessage'
 
@@ -24,6 +25,9 @@ export const createStop = (
     context.isMockingEnabled = false
     window.clearInterval(context.keepAliveInterval)
 
+    // Clear the WebSocket clients from the shared storage.
+    localStorage.removeItem(MSW_WEBSOCKET_CLIENTS_KEY)
+
     printStopMessage({ quiet: context.startOptions?.quiet })
   }
 }

src/core/ws/WebSocketClientManager.ts

@@ -6,7 +6,7 @@ import type {
 } from '@mswjs/interceptors/WebSocket'
 import { matchRequestUrl, type Path } from '../utils/matching/matchRequestUrl'
 
-const MSW_WEBSOCKET_CLIENTS_KEY = 'msw:ws:clients'
+export const MSW_WEBSOCKET_CLIENTS_KEY = 'msw:ws:clients'
 
 export type WebSocketBroadcastChannelMessage =
   | {
@@ -42,6 +42,22 @@ export class WebSocketClientManager {
     private url: Path,
   ) {
     this.inMemoryClients = new Set()
+
+    if (typeof localStorage !== 'undefined') {
+      // When the worker clears the local storage key in "worker.stop()",
+      // also clear the in-memory clients map.
+      localStorage.removeItem = new Proxy(localStorage.removeItem, {
+        apply: (target, thisArg, args) => {
+          const [key] = args
+
+          if (key === MSW_WEBSOCKET_CLIENTS_KEY) {
+            this.inMemoryClients.clear()
+          }
+
+          return Reflect.apply(target, thisArg, args)
+        },
+      })
+    }
   }
 
   /**
@@ -53,6 +69,8 @@ export class WebSocketClientManager {
     if (typeof localStorage !== 'undefined') {
       const inMemoryClients = Array.from(this.inMemoryClients)
 
+      console.log('get clients()', inMemoryClients, this.getSerializedClients())
+
       return new Set(
         inMemoryClients.concat(
           this.getSerializedClients()

test/browser/ws-api/ws.clients.browser.test.ts

@@ -1,5 +1,5 @@
 import type { WebSocketLink, ws } from 'msw'
-import type { setupWorker } from 'msw/browser'
+import type { SetupWorker, setupWorker } from 'msw/browser'
 import { test, expect } from '../playwright.extend'
 
 declare global {
@@ -8,6 +8,7 @@ declare global {
       ws: typeof ws
       setupWorker: typeof setupWorker
     }
+    worker: SetupWorker
     link: WebSocketLink
     ws: WebSocket
     messages: string[]
@@ -164,3 +165,41 @@ test('broadcasts messages across runtimes', async ({
     'hi from two',
   ])
 })
+
+test('clears the list of clients when the worker is stopped', async ({
+  loadExample,
+  page,
+}) => {
+  await loadExample(require.resolve('./ws.runtime.js'), {
+    skipActivation: true,
+  })
+
+  await page.evaluate(async () => {
+    const { setupWorker, ws } = window.msw
+    const api = ws.link('wss://example.com')
+    const worker = setupWorker(api.on('connection', () => {}))
+    window.link = api
+    window.worker = worker
+    await worker.start()
+  })
+
+  await page.evaluate(async () => {
+    const ws = new WebSocket('wss://example.com')
+    await new Promise((done) => (ws.onopen = done))
+  })
+
+  // Must return 1 after a single client joined.
+  expect(
+    await page.evaluate(() => {
+      return window.link.clients.size
+    }),
+  ).toBe(1)
+
+  await page.evaluate(() => {
+    window.worker.stop()
+  })
+
+  // Must return 0.
+  // The localStorage has been purged, and the in-memory manager clients too.
+  expect(await page.evaluate(() => window.link.clients.size)).toBe(0)
+})