From a9ef6b6c96ae8fdf891e29ab05c5276773425e1c Mon Sep 17 00:00:00 2001
From: Eduardo San Martin Morote <posva13@gmail.com>
Date: Sat, 18 Dec 2021 10:40:37 +0100
Subject: [PATCH] fix(subscribe): direct mutation should not trigger detached
 subscriptions

Fix #908
---
 .../pinia/__tests__/subscriptions.spec.ts     | 19 ++++++++++++++++---
 packages/pinia/src/store.ts                   | 14 ++++----------
 packages/pinia/src/subscriptions.ts           |  6 +++++-
 packages/pinia/src/types.ts                   |  4 ++--
 packages/pinia/test-dts/state.test-d.ts       | 17 +++++++++++++++++
 5 files changed, 44 insertions(+), 16 deletions(-)

diff --git a/packages/pinia/__tests__/subscriptions.spec.ts b/packages/pinia/__tests__/subscriptions.spec.ts
index 82fb25a6..25db3a57 100644
--- a/packages/pinia/__tests__/subscriptions.spec.ts
+++ b/packages/pinia/__tests__/subscriptions.spec.ts
@@ -124,7 +124,6 @@ describe('Subscriptions', () => {
       expect(spy2).toHaveBeenCalledTimes(0)
 
       s1.name = 'Edu'
-
       expect(spy1).toHaveBeenCalledTimes(1)
       expect(spy2).toHaveBeenCalledTimes(1)
 
@@ -132,12 +131,26 @@ describe('Subscriptions', () => {
       expect(spy1).toHaveBeenCalledTimes(2)
       expect(spy2).toHaveBeenCalledTimes(2)
 
+      s1.$patch((state) => {
+        state.name = 'other'
+      })
+      expect(spy1).toHaveBeenCalledTimes(3)
+      expect(spy2).toHaveBeenCalledTimes(3)
+
       wrapper.unmount()
       await nextTick()
 
       s1.$patch({ name: 'b' })
-      expect(spy1).toHaveBeenCalledTimes(2)
-      expect(spy2).toHaveBeenCalledTimes(3)
+      expect(spy1).toHaveBeenCalledTimes(3)
+      expect(spy2).toHaveBeenCalledTimes(4)
+      s1.$patch((state) => {
+        state.name = 'c'
+      })
+      expect(spy1).toHaveBeenCalledTimes(3)
+      expect(spy2).toHaveBeenCalledTimes(5)
+      s1.name = 'd'
+      expect(spy1).toHaveBeenCalledTimes(3)
+      expect(spy2).toHaveBeenCalledTimes(6)
     })
   })
 
diff --git a/packages/pinia/src/store.ts b/packages/pinia/src/store.ts
index c5697e06..d38becab 100644
--- a/packages/pinia/src/store.ts
+++ b/packages/pinia/src/store.ts
@@ -48,7 +48,7 @@ import {
 import { setActivePinia, piniaSymbol, Pinia, activePinia } from './rootStore'
 import { IS_CLIENT } from './env'
 import { patchObject } from './hmr'
-import { addSubscription, triggerSubscriptions } from './subscriptions'
+import { addSubscription, triggerSubscriptions, noop } from './subscriptions'
 
 type _ArrayType<AT> = AT extends Array<infer T> ? T : never
 
@@ -173,8 +173,6 @@ function createOptionsStore<
   return store as any
 }
 
-const noop = () => {}
-
 function createSetupStore<
   Id extends string,
   SS,
@@ -376,10 +374,11 @@ function createSetupStore<
     $patch,
     $reset,
     $subscribe(callback, options = {}) {
-      const _removeSubscription = addSubscription(
+      const removeSubscription = addSubscription(
         subscriptions,
         callback,
-        options.detached
+        options.detached,
+        () => stopWatcher()
       )
       const stopWatcher = scope.run(() =>
         watch(
@@ -400,11 +399,6 @@ function createSetupStore<
         )
       )!
 
-      const removeSubscription = () => {
-        stopWatcher()
-        _removeSubscription()
-      }
-
       return removeSubscription
     },
     $dispose,
diff --git a/packages/pinia/src/subscriptions.ts b/packages/pinia/src/subscriptions.ts
index 86fef170..5dfd2620 100644
--- a/packages/pinia/src/subscriptions.ts
+++ b/packages/pinia/src/subscriptions.ts
@@ -1,10 +1,13 @@
 import { getCurrentInstance, onUnmounted } from 'vue-demi'
 import { _Method } from './types'
 
+export const noop = () => {}
+
 export function addSubscription<T extends _Method>(
   subscriptions: T[],
   callback: T,
-  detached?: boolean
+  detached?: boolean,
+  onCleanup: () => void = noop
 ) {
   subscriptions.push(callback)
 
@@ -12,6 +15,7 @@ export function addSubscription<T extends _Method>(
     const idx = subscriptions.indexOf(callback)
     if (idx > -1) {
       subscriptions.splice(idx, 1)
+      onCleanup()
     }
   }
 
diff --git a/packages/pinia/src/types.ts b/packages/pinia/src/types.ts
index 356c1f34..d6027582 100644
--- a/packages/pinia/src/types.ts
+++ b/packages/pinia/src/types.ts
@@ -333,11 +333,11 @@ export interface _StoreWithState<
   /**
    * Group multiple changes into one function. Useful when mutating objects like
    * Sets or arrays and applying an object patch isn't practical, e.g. appending
-   * to an array.
+   * to an array. The function passed to `$patch()` **must be synchronous**.
    *
    * @param stateMutator - function that mutates `state`, cannot be async
    */
-  $patch<F extends (state: UnwrapRef<S>) => void>(
+  $patch<F extends (state: UnwrapRef<S>) => any>(
     // this prevents the user from using `async` which isn't allowed
     stateMutator: ReturnType<F> extends Promise<any> ? never : F
   ): void
diff --git a/packages/pinia/test-dts/state.test-d.ts b/packages/pinia/test-dts/state.test-d.ts
index d3b7550e..b69abd4f 100644
--- a/packages/pinia/test-dts/state.test-d.ts
+++ b/packages/pinia/test-dts/state.test-d.ts
@@ -53,6 +53,11 @@ const useStore = defineStore({
 
 const store = useStore()
 
+store.$patch({ counter: 2 })
+store.$patch((state) => {
+  expectType<number>(state.counter)
+})
+
 expectType<number>(store.$state.counter)
 expectType<number>(store.$state.double)
 
@@ -62,3 +67,15 @@ expectType<number>(store.fromARef)
 
 expectType<{ msg: string }>(store.aShallowRef)
 expectType<{ msg: string }>(store.$state.aShallowRef)
+
+const onlyState = defineStore({
+  id: 'main',
+  state: () => ({
+    counter: 0,
+  }),
+})()
+
+onlyState.$patch({ counter: 2 })
+onlyState.$patch((state) => {
+  expectType<number>(state.counter)
+})
-- 
2.47.3