feat: add devtools.cleanup() method (#3111)

* feat(devtools): add cleanup method

* docs(devtools): add cleanup section

* reduce lines

* test(devtools): test if the connection removed after cleanup

---------

Co-authored-by: daishi <daishi@axlight.com>
Co-authored-by: Daishi Kato <dai-shi@users.noreply.github.com>

docs/middlewares/devtools.md

@@ -24,6 +24,7 @@ const nextStateCreatorFn = devtools(stateCreatorFn, devtoolsOptions)
 - [Usage](#usage)
   - [Debugging a store](#debugging-a-store)
   - [Debugging a Slices pattern based store](#debugging-a-slices-pattern-based-store)
+  - [Cleanup](#cleanup)
 - [Troubleshooting](#troubleshooting)
   - [Only one store is displayed](#only-one-store-is-displayed)
   - [Action names are labeled as 'anonymous'](#all-action-names-are-labeled-as-anonymous)
@@ -156,6 +157,27 @@ const useJungleStore = create<JungleStore>()(
 )
 ```
 
+### Cleanup
+
+When a store is no longer needed, you can clean up the Redux DevTools connection by calling the `cleanup` method on the store:
+
+```ts
+import { create } from 'zustand'
+import { devtools } from 'zustand/middleware'
+
+const useStore = create(
+  devtools((set) => ({
+    count: 0,
+    increment: () => set((state) => ({ count: state.count + 1 })),
+  })),
+)
+
+// When you're done with the store, clean it up
+useStore.devtools.cleanup()
+```
+
+This is particularly useful in applications that wrap store in context or create multiple stores dynamically.
+
 ## Troubleshooting
 
 ### Only one store is displayed

src/middleware/devtools.ts

@@ -65,6 +65,9 @@ type StoreDevtools<S> = S extends {
   ? {
       setState(...args: [...args: TakeTwo<Sa1>, action?: Action]): Sr1
       setState(...args: [...args: TakeTwo<Sa2>, action?: Action]): Sr2
+      devtools: {
+        cleanup: () => void
+      }
     }
   : never
 
@@ -146,6 +149,19 @@ const extractConnectionInformation = (
   return { type: 'tracked' as const, store, ...newConnection }
 }
 
+const removeStoreFromTrackedConnections = (
+  name: string | undefined,
+  store: string | undefined,
+) => {
+  if (store === undefined) return
+  const connectionInfo = trackedConnections.get(name)
+  if (!connectionInfo) return
+  delete connectionInfo.stores[store]
+  if (Object.keys(connectionInfo.stores).length === 0) {
+    trackedConnections.delete(name)
+  }
+}
+
 const devtoolsImpl: DevtoolsImpl =
   (fn, devtoolsOptions = {}) =>
   (set, get, api) => {
@@ -200,6 +216,17 @@ const devtoolsImpl: DevtoolsImpl =
       )
       return r
     }) as NamedSet<S>
+    ;(api as StoreApi<S> & StoreDevtools<S>).devtools = {
+      cleanup: () => {
+        if (
+          connection &&
+          typeof (connection as any).unsubscribe === 'function'
+        ) {
+          ;(connection as any).unsubscribe()
+        }
+        removeStoreFromTrackedConnections(options.name, store)
+      },
+    }
 
     const setStateFromDevtools: StoreApi<S>['setState'] = (...a) => {
       const originalIsRecording = isRecording

tests/devtools.test.tsx

@@ -98,7 +98,11 @@ const extensionConnector = {
         subscribers.push(f)
         return () => {}
       }),
-      unsubscribe: vi.fn(),
+      unsubscribe: vi.fn(() => {
+        connectionMap.delete(
+          areNameUndefinedMapsNeeded ? options.testConnectionId : key,
+        )
+      }),
       send: vi.fn(),
       init: vi.fn(),
       error: vi.fn(),
@@ -2448,3 +2452,40 @@ describe('when create devtools was called multiple times with `name` and `store`
     })
   })
 })
+
+describe('cleanup', () => {
+  it('should unsubscribe from devtools when cleanup is called', async () => {
+    const options = { name: 'test' }
+    const store = createStore(devtools(() => ({ count: 0 }), options))
+    const [connection] = getNamedConnectionApis(options.name)
+    store.devtools.cleanup()
+
+    expect(connection.unsubscribe).toHaveBeenCalledTimes(1)
+  })
+
+  it('should remove store from tracked connection after cleanup', async () => {
+    const options = {
+      name: 'test-store-name',
+      store: 'test-store-id',
+      enabled: true,
+    }
+
+    const store1 = createStore(devtools(() => ({ count: 0 }), options))
+    store1.devtools.cleanup()
+    const store2 = createStore(devtools(() => ({ count: 0 }), options))
+
+    const [connection] = getNamedConnectionApis(options.name)
+
+    store2.setState({ count: 15 }, false, 'updateCount')
+    expect(connection.send).toHaveBeenLastCalledWith(
+      { type: `${options.store}/updateCount` },
+      { [options.store]: { count: 15 } },
+    )
+
+    store1.setState({ count: 20 }, false, 'ignoredAction')
+    expect(connection.send).not.toHaveBeenLastCalledWith(
+      { type: `${options.store}/ignoredAction` },
+      expect.anything(),
+    )
+  })
+})