feat(middleware/devtools): Better redux devtools. One connection for …

…all zustand stores (#1435)

* add example project for zustand devonly devtools with better rewind feature of redux devtools

* fix yarn install for demo

* remove lodash from demo

* adjust utils

* remove assets from demo project

* add codesandbox

* devtools for groups of stores

* remove command from .yml

* add tests

* cleanup

* Existing test cleanup

* Slight test cleanup; extract iife

* Slight cleanup

* Non-recursive S

* Remove initialState

Co-authored-by: Corin (CWStra) <cwstra@gmail.com>

.codesandbox/ci.json

@@ -6,7 +6,8 @@
     "simple-react-browserify-x9yni",
     "simple-snowpack-react-o1gmx",
     "react-parcel-onewf",
-    "next-js-uo1h0"
+    "next-js-uo1h0",
+    "pavlobu-zustand-demo-frutec"
   ],
   "node": "14"
 }

readme.md

@@ -389,6 +389,21 @@ const usePlainStore = create(devtools(store))
 const useReduxStore = create(devtools(redux(reducer, initialState)))
 ```
 
+One redux devtools connection for multiple stores
+
+```jsx
+import { devtools } from 'zustand/middleware'
+
+// Usage with a plain action store, it will log actions as "setState"
+const usePlainStore1 = create(devtools(store), { name, store: storeName1 })
+const usePlainStore2 = create(devtools(store), { name, store: storeName2 })
+// Usage with a redux store, it will log full action types
+const useReduxStore = create(devtools(redux(reducer, initialState)), , { name, store: storeName3 })
+const useReduxStore = create(devtools(redux(reducer, initialState)), , { name, store: storeName4 })
+```
+
+Assingning different connection names, will separate stores in redux devtools. This also helps grouping different stores into separate redux devtools connections.
+
 devtools takes the store function as its first argument, optionally you can name the store or configure [serialize](https://github.com/zalmoxisus/redux-devtools-extension/blob/master/docs/API/Arguments.md#serialize) options with a second argument.
 
 Name store: `devtools(store, {name: "MyStore"})`, which will create a separate instance named "MyStore" in the devtools.

src/middleware/devtools.ts

@@ -122,6 +122,7 @@ type StoreDevtools<S> = S extends {
 export interface DevtoolsOptions extends Config {
   enabled?: boolean
   anonymousActionType?: string
+  store?: string
 }
 
 type Devtools = <
@@ -147,13 +148,63 @@ type DevtoolsImpl = <T>(
 
 export type NamedSet<T> = WithDevtools<StoreApi<T>>['setState']
 
+type Connection = ReturnType<
+  NonNullable<Window['__REDUX_DEVTOOLS_EXTENSION__']>['connect']
+>
+type ConnectionName = string | undefined
+type StoreName = string
+type StoreInformation = StoreApi<unknown>
+type ConnectionInformation = {
+  connection: Connection
+  stores: Record<StoreName, StoreInformation>
+}
+const trackedConnections: Map<ConnectionName, ConnectionInformation> = new Map()
+
+const getTrackedConnectionState = (
+  name: string | undefined
+): Record<string, any> => {
+  const api = trackedConnections.get(name)
+  if (!api) return {}
+  return Object.fromEntries(
+    Object.entries(api.stores).map(([key, api]) => [key, api.getState()])
+  )
+}
+
+const extractConnectionInformation = (
+  store: string | undefined,
+  extensionConnector: NonNullable<
+    typeof window['__REDUX_DEVTOOLS_EXTENSION__']
+  >,
+  options: Omit<DevtoolsOptions, 'enabled' | 'anonymousActionType' | 'store'>
+) => {
+  if (store === undefined) {
+    return {
+      type: 'untracked' as const,
+      connection: extensionConnector.connect(options),
+    }
+  }
+  const existingConnection = trackedConnections.get(options.name)
+  if (existingConnection) {
+    return { type: 'tracked' as const, store, ...existingConnection }
+  }
+  const newConnection: ConnectionInformation = {
+    connection: extensionConnector.connect(options),
+    stores: {},
+  }
+  trackedConnections.set(options.name, newConnection)
+  return { type: 'tracked' as const, store, ...newConnection }
+}
+
 const devtoolsImpl: DevtoolsImpl =
   (fn, devtoolsOptions = {}) =>
   (set, get, api) => {
-    type S = ReturnType<typeof fn>
+    const { enabled, anonymousActionType, store, ...options } = devtoolsOptions
+
+    type S = ReturnType<typeof fn> & {
+      [store: string]: ReturnType<typeof fn>
+    }
     type PartialState = Partial<S> | ((s: S) => Partial<S>)
 
-    const { enabled, anonymousActionType, ...options } = devtoolsOptions
     let extensionConnector:
       | typeof window['__REDUX_DEVTOOLS_EXTENSION__']
       | false
@@ -173,22 +224,36 @@ const devtoolsImpl: DevtoolsImpl =
       return fn(set, get, api)
     }
 
-    const extension = extensionConnector.connect(options)
+    const { connection, ...connectionInformation } =
+      extractConnectionInformation(store, extensionConnector, options)
 
     let isRecording = true
     ;(api.setState as NamedSet<S>) = (state, replace, nameOrAction) => {
       const r = set(state, replace)
       if (!isRecording) return r
-      extension.send(
+      const action: Action<unknown> =
         nameOrAction === undefined
           ? { type: anonymousActionType || 'anonymous' }
           : typeof nameOrAction === 'string'
           ? { type: nameOrAction }
-          : nameOrAction,
-        get()
+          : nameOrAction
+      if (store === undefined) {
+        connection?.send(action, get())
+        return r
+      }
+      connection?.send(
+        {
+          ...action,
+          type: `${store}/${action.type}`,
+        },
+        {
+          ...getTrackedConnectionState(options.name),
+          [store]: api.getState(),
+        }
       )
       return r
     }
+
     const setStateFromDevtools: StoreApi<S>['setState'] = (...a) => {
       const originalIsRecording = isRecording
       isRecording = false
@@ -197,7 +262,21 @@ const devtoolsImpl: DevtoolsImpl =
     }
 
     const initialState = fn(api.setState, get, api)
-    extension.init(initialState)
+    if (connectionInformation.type === 'untracked') {
+      connection?.init(initialState)
+    } else {
+      connectionInformation.stores[connectionInformation.store] = api
+      connection?.init(
+        Object.fromEntries(
+          Object.entries(connectionInformation.stores).map(([key, store]) => [
+            key,
+            key === connectionInformation.store
+              ? initialState
+              : store.getState(),
+          ])
+        )
+      )
+    }
 
     if (
       (api as any).dispatchFromDevtools &&
@@ -222,7 +301,7 @@ const devtoolsImpl: DevtoolsImpl =
     }
 
     ;(
-      extension as unknown as {
+      connection as unknown as {
         // FIXME https://github.com/reduxjs/redux-devtools/issues/1097
         subscribe: (
           listener: (message: Message) => void
@@ -241,7 +320,32 @@ const devtoolsImpl: DevtoolsImpl =
             message.payload,
             (action) => {
               if (action.type === '__setState') {
-                setStateFromDevtools(action.state as PartialState)
+                if (store === undefined) {
+                  setStateFromDevtools(action.state as PartialState)
+                  return
+                }
+                if (Object.keys(action.state as S).length !== 1) {
+                  console.error(
+                    `
+                    [zustand devtools middleware] Unsupported __setState action format. 
+                    When using 'store' option in devtools(), the 'state' should have only one key, which is a value of 'store' that was passed in devtools(),
+                    and value of this only key should be a state object. Example: { "type": "__setState", "state": { "abc123Store": { "foo": "bar" } } }
+                    `
+                  )
+                }
+                const stateFromDevtools = (action.state as S)[store]
+                if (
+                  stateFromDevtools === undefined ||
+                  stateFromDevtools === null
+                ) {
+                  return
+                }
+                if (
+                  JSON.stringify(api.getState()) !==
+                  JSON.stringify(stateFromDevtools)
+                ) {
+                  setStateFromDevtools(stateFromDevtools)
+                }
                 return
               }
 
@@ -254,31 +358,56 @@ const devtoolsImpl: DevtoolsImpl =
         case 'DISPATCH':
           switch (message.payload.type) {
             case 'RESET':
-              setStateFromDevtools(initialState)
-              return extension.init(api.getState())
+              setStateFromDevtools(initialState as S)
+              if (store === undefined) {
+                return connection?.init(api.getState())
+              }
+              return connection?.init(getTrackedConnectionState(options.name))
 
             case 'COMMIT':
-              return extension.init(api.getState())
+              if (store === undefined) {
+                connection?.init(api.getState())
+                return
+              }
+              return connection?.init(getTrackedConnectionState(options.name))
 
             case 'ROLLBACK':
               return parseJsonThen<S>(message.state, (state) => {
-                setStateFromDevtools(state)
-                extension.init(api.getState())
+                if (store === undefined) {
+                  setStateFromDevtools(state)
+                  connection?.init(api.getState())
+                  return
+                }
+                setStateFromDevtools(state[store] as S)
+                connection?.init(getTrackedConnectionState(options.name))
               })
 
             case 'JUMP_TO_STATE':
             case 'JUMP_TO_ACTION':
               return parseJsonThen<S>(message.state, (state) => {
-                setStateFromDevtools(state)
+                if (store === undefined) {
+                  setStateFromDevtools(state)
+                  return
+                }
+                if (
+                  JSON.stringify(api.getState()) !==
+                  JSON.stringify(state[store])
+                ) {
+                  setStateFromDevtools(state[store] as S)
+                }
               })
 
             case 'IMPORT_STATE': {
               const { nextLiftedState } = message.payload
               const lastComputedState =
                 nextLiftedState.computedStates.slice(-1)[0]?.state
               if (!lastComputedState) return
-              setStateFromDevtools(lastComputedState)
-              extension.send(
+              if (store === undefined) {
+                setStateFromDevtools(lastComputedState)
+              } else {
+                setStateFromDevtools(lastComputedState[store])
+              }
+              connection?.send(
                 null as any, // FIXME no-any
                 nextLiftedState
               )