Fix HMR re-run plugins on existing store after HMR to refresh stale closures

packages/pinia/__tests__/hmr.spec.ts

@@ -1,5 +1,5 @@
 import { beforeEach, describe, it, expect, vi } from 'vitest'
-import { computed, reactive, ref, toRefs, watch } from 'vue'
+import { computed, createApp, reactive, ref, toRefs, watch } from 'vue'
 import {
   createPinia,
   defineStore,
@@ -537,7 +537,105 @@ describe('HMR', () => {
   })
 
   describe('both', () => {
-    it.todo('keeps $subscribe subscriptions')
-    it.todo('$onAction subscriptions')
+    const setup = () => {
+      const n = ref(0)
+      function increment(amount = 1) {
+        n.value += amount
+      }
+      return { n, increment }
+    }
+
+    it('keeps $subscribe subscriptions after HMR', () => {
+      const pinia = createPinia()
+      const app = createApp({})
+      app.use(pinia)
+      setActivePinia(pinia)
+
+      const useStore = defineStore('id', setup)
+      const store = useStore()
+      const spy = vi.fn()
+      store.$subscribe(spy, { detached: true })
+
+      defineStore('id', setup)(null, store)
+
+      store.$patch({ n: 1 })
+      expect(spy).toHaveBeenCalledTimes(1)
+    })
+
+    it('keeps $onAction subscriptions after HMR', () => {
+      const pinia = createPinia()
+      const app = createApp({})
+      app.use(pinia)
+      setActivePinia(pinia)
+
+      const useStore = defineStore('id', setup)
+      const store = useStore()
+      const spy = vi.fn()
+      store.$onAction(spy, true)
+
+      defineStore('id', setup)(null, store)
+
+      store.increment()
+      expect(spy).toHaveBeenCalledTimes(1)
+    })
+
+    // _hotUpdate re-runs plugins on the existing store so subscriptions capture fresh closures from the new module.
+    it('plugin $subscribe uses fresh closures after HMR', () => {
+      const pinia = createPinia()
+      const app = createApp({})
+      app.use(pinia)
+
+      let closureValue = 'original'
+      let capturedValue = ''
+      pinia.use(({ store }) => {
+        const current = closureValue // captured at plugin-run time
+        store.$subscribe(() => {
+          capturedValue = current
+        })
+      })
+
+      setActivePinia(pinia)
+      const useStore = defineStore('id', setup)
+      const store = useStore()
+
+      store.$patch({ n: 1 })
+      expect(capturedValue).toBe('original')
+
+      // simulate the module updating (new version of the store file loaded by HMR)
+      closureValue = 'updated'
+      defineStore('id', setup)(null, store)
+
+      store.$patch({ n: 2 })
+      // _hotUpdate re-runs plugins on the existing store with the new closure → capturedValue becomes 'updated'
+      expect(capturedValue).toBe('updated')
+    })
+
+    it('plugin $onAction uses fresh closures after HMR', () => {
+      const pinia = createPinia()
+      const app = createApp({})
+      app.use(pinia)
+
+      let closureValue = 'original'
+      let capturedValue = ''
+      pinia.use(({ store }) => {
+        const current = closureValue
+        store.$onAction(() => {
+          capturedValue = current
+        })
+      })
+
+      setActivePinia(pinia)
+      const useStore = defineStore('id', setup)
+      const store = useStore()
+
+      store.increment()
+      expect(capturedValue).toBe('original')
+
+      closureValue = 'updated'
+      defineStore('id', setup)(null, store)
+
+      store.increment()
+      expect(capturedValue).toBe('updated')
+    })
   })
 })

packages/pinia/src/store.ts

@@ -268,6 +268,7 @@ function createSetupStore<
   let subscriptions: Set<SubscriptionCallback<S>> = new Set()
   let actionSubscriptions: Set<StoreOnActionListener<Id, S, G, A>> = new Set()
   let debuggerEvents: DebuggerEvent[] | DebuggerEvent
+  let pluginScope: EffectScope | undefined
   const initialState = pinia.state.value[$id] as UnwrapRef<S> | undefined
 
   // avoid setting the state for option stores if it is set
@@ -451,24 +452,28 @@ function createSetupStore<
         options.detached,
         () => stopWatcher()
       )
-      const stopWatcher = scope.run(() =>
-        watch(
-          () => pinia.state.value[$id] as UnwrapRef<S>,
-          (state) => {
-            if (options.flush === 'sync' ? isSyncListening : isListening) {
-              callback(
-                {
-                  storeId: $id,
-                  type: MutationType.direct,
-                  events: debuggerEvents as DebuggerEvent,
-                },
-                state
-              )
-            }
-          },
-          assign({}, $subscribeOptions, options)
-        )
-      )!
+      // hot stores are temporary; skip the watcher to avoid spurious callbacks
+      // when the hot store's state entry is deleted after HMR completes
+      const stopWatcher = hot
+        ? noop
+        : scope.run(() =>
+            watch(
+              () => pinia.state.value[$id] as UnwrapRef<S>,
+              (state) => {
+                if (options.flush === 'sync' ? isSyncListening : isListening) {
+                  callback(
+                    {
+                      storeId: $id,
+                      type: MutationType.direct,
+                      events: debuggerEvents as DebuggerEvent,
+                    },
+                    state
+                  )
+                }
+              },
+              assign({}, $subscribeOptions, options)
+            )
+          )!
 
       return removeSubscription
     },
@@ -689,6 +694,8 @@ function createSetupStore<
       // update the values used in devtools and to allow deleting new properties later on
       store._hmrPayload = newStore._hmrPayload
       store._getters = newStore._getters
+      // re-run plugins so their subscriptions use fresh closures from updated modules
+      runPlugins()
       store._hotUpdating = false
     })
   }
@@ -713,45 +720,53 @@ function createSetupStore<
     )
   }
 
-  // apply all plugins
-  pinia._p.forEach((extender) => {
-    const extensions = scope.run(() =>
-      extender({
-        store: store as Store,
-        app: pinia._a,
-        pinia,
-        options: optionsForPlugin,
-      })
-    )!
+  // apply all plugins; extracted so _hotUpdate can re-run them with fresh closures
+  function runPlugins() {
+    if (pluginScope) pluginScope.stop()
+    pluginScope = scope.run(() => effectScope())!
+    pinia._p.forEach((extender) => {
+      const extensions = pluginScope!.run(() =>
+        extender({
+          store: store as Store,
+          app: pinia._a,
+          pinia,
+          options: optionsForPlugin,
+        })
+      )!
 
-    /* istanbul ignore else */
-    if (__USE_DEVTOOLS__ && IS_CLIENT) {
-      Object.keys(extensions || {}).forEach((key) =>
-        store._customProperties.add(key)
-      )
-    }
+      /* istanbul ignore else */
+      if (__USE_DEVTOOLS__ && IS_CLIENT) {
+        Object.keys(extensions || {}).forEach((key) =>
+          store._customProperties.add(key)
+        )
+      }
 
-    // Check properties that are not properly configured. We check the values
-    // as the plugin returned them: once assigned to the store, a `reactive()`
-    // value is unwrapped and indistinguishable from a plain object.
-    if (__DEV__) {
-      for (const key in extensions) {
-        const value = (extensions as any)[key]
-        // `null` is included (typeof null === 'object'). refs, reactive objects,
-        // primitives, functions and markRaw() values are skipped.
-        if (
-          typeof value === 'object' &&
-          !isRef(value) &&
-          !isReactive(value) &&
-          !value?.__v_skip
-        ) {
-          diagnostics.PINIA_R1006({ key, id: $id })
+      // Check properties that are not properly configured. We check the values
+      // as the plugin returned them: once assigned to the store, a `reactive()`
+      // value is unwrapped and indistinguishable from a plain object.
+      if (__DEV__) {
+        for (const key in extensions) {
+          const value = (extensions as any)[key]
+          // `null` is included (typeof null === 'object'). refs, reactive objects,
+          // primitives, functions and markRaw() values are skipped.
+          if (
+            typeof value === 'object' &&
+            !isRef(value) &&
+            !isReactive(value) &&
+            !value?.__v_skip
+          ) {
+            diagnostics.PINIA_R1006({ key, id: $id })
+          }
         }
       }
-    }
 
-    assign(store, extensions)
-  })
+      assign(store, extensions)
+    })
+  }
+  // skip plugins for hot stores; _hotUpdate re-runs them on the existing store
+  if (!hot) {
+    runPlugins()
+  }
 
   if (
     __DEV__ &&