fix(types)!: require complete state if setState's replace flag is set (#2580)

* fix(types): require complete state if `setState`'s `replace` flag is set

* switch to variant 2

* fix type errors

* update setState types for devtools and immer

* make devtools setState non-generic

* add migration guide

* merge migration guides

* run prettier

* Update tests/middlewareTypes.test.tsx

---------

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

docs/guides/migrating-to-v5.mdx

@@ -17,6 +17,7 @@ We highly recommend to update to the latest version of v4, before migrating to v
 - Drop UMD/SystemJS support
 - Organize entry points in the package.json
 - Drop ES5 support
+- Stricter types when setState's replace flag is set
 - Other small improvements (technically breaking changes)
 
 ## Migration Guide
@@ -127,6 +128,44 @@ Alternatively, if you need v4 behavior, `createWithEqualityFn` will do.
 import { createWithEqualityFn as create } from 'zustand/traditional'
 ```
 
+### Stricter types when setState's replace flag is set (Typescript only)
+
+```diff
+- setState:
+-   (partial: T | Partial<T> | ((state: T) => T | Partial<T>), replace?: boolean | undefined) => void;
++ setState:
++   (partial: T | Partial<T> | ((state: T) => T | Partial<T>), replace?: false) => void;
++   (state: T | ((state: T) => T), replace: true) => void;
+```
+
+If you are not using the `replace` flag, no migration is required.
+
+If you are using the `replace` flag and it's set to `true`, you must provide a complete state object.
+This change ensures that `store.setState({}, true)` (which results in an invalid state) is no longer considered valid.
+
+**Examples:**
+
+```ts
+// Partial state update (valid)
+store.setState({ key: 'value' })
+
+// Complete state replacement (valid)
+store.setState({ key: 'value' }, true)
+
+// Incomplete state replacement (invalid)
+store.setState({}, true) // Error
+```
+
+#### Handling Dynamic `replace` Flag
+
+If the value of the `replace` flag is dynamic and determined at runtime, you might face issues. To handle this, you can use a workaround by annotating the `replace` parameter with `as any`:
+
+```ts
+const replaceFlag = Math.random() > 0.5
+store.setState(partialOrFull, replaceFlag as any)
+```
+
 ## Links
 
 - https://github.com/pmndrs/zustand/pull/2138
+- https://github.com/pmndrs/zustand/pull/2580

docs/guides/typescript.md

@@ -232,6 +232,36 @@ For a usual statically typed language, this is impossible. But thanks to TypeScr
 
 If you are eager to know what the answer is to this particular problem then you can [see it here](#middleware-that-changes-the-store-type).
 
+### Handling Dynamic `replace` Flag
+
+If the value of the `replace` flag is not known at compile time and is determined dynamically, you might face issues. To handle this, you can use a workaround by annotating the `replace` parameter with `as any`:
+
+```ts
+const replaceFlag = Math.random() > 0.5
+store.setState(partialOrFull, replaceFlag as any)
+```
+
+#### Example with `as any` Workaround
+
+```ts
+import { create } from 'zustand'
+
+interface BearState {
+  bears: number
+  increase: (by: number) => void
+}
+
+const useBearStore = create<BearState>()((set) => ({
+  bears: 0,
+  increase: (by) => set((state) => ({ bears: state.bears + by })),
+}))
+
+const replaceFlag = Math.random() > 0.5
+useBearStore.setState({ bears: 5 }, replaceFlag as any) // Using the workaround
+```
+
+By following this approach, you can ensure that your code handles dynamic `replace` flags without encountering type issues.
+
 ## Common recipes
 
 ### Middleware that doesn't change the store type

src/middleware/devtools.ts

@@ -49,13 +49,22 @@ type TakeTwo<T> = T extends { length: 0 }
 
 type WithDevtools<S> = Write<S, StoreDevtools<S>>
 
+type Action =
+  | string
+  | {
+      type: string
+      [x: string | number | symbol]: unknown
+    }
 type StoreDevtools<S> = S extends {
-  setState: (...a: infer Sa) => infer Sr
+  setState: {
+    // capture both overloads of setState
+    (...a: infer Sa1): infer Sr1
+    (...a: infer Sa2): infer Sr2
+  }
 }
   ? {
-      setState<A extends string | { type: string }>(
-        ...a: [...a: TakeTwo<Sa>, action?: A]
-      ): Sr
+      setState(...a: [...a: TakeTwo<Sa1>, action?: Action]): Sr1
+      setState(...a: [...a: TakeTwo<Sa2>, action?: Action]): Sr2
     }
   : never
 
@@ -165,8 +174,8 @@ const devtoolsImpl: DevtoolsImpl =
       extractConnectionInformation(store, extensionConnector, options)
 
     let isRecording = true
-    ;(api.setState as NamedSet<S>) = (state, replace, nameOrAction) => {
-      const r = set(state, replace)
+    ;(api.setState as any) = ((state, replace, nameOrAction: Action) => {
+      const r = set(state, replace as any)
       if (!isRecording) return r
       const action: { type: string } =
         nameOrAction === undefined
@@ -189,12 +198,12 @@ const devtoolsImpl: DevtoolsImpl =
         },
       )
       return r
-    }
+    }) as NamedSet<S>
 
     const setStateFromDevtools: StoreApi<S>['setState'] = (...a) => {
       const originalIsRecording = isRecording
       isRecording = false
-      set(...a)
+      set(...(a as Parameters<typeof set>))
       isRecording = originalIsRecording
     }
 

src/middleware/immer.ts

@@ -38,13 +38,21 @@ type StoreImmer<S> = S extends {
   getState: () => infer T
   setState: infer SetState
 }
-  ? SetState extends (...a: infer A) => infer Sr
+  ? SetState extends {
+      (...a: infer A1): infer Sr1
+      (...a: infer A2): infer Sr2
+    }
     ? {
         setState(
           nextStateOrUpdater: T | Partial<T> | ((state: Draft<T>) => void),
-          shouldReplace?: boolean | undefined,
-          ...a: SkipTwo<A>
-        ): Sr
+          shouldReplace?: false,
+          ...a: SkipTwo<A1>
+        ): Sr1
+        setState(
+          nextStateOrUpdater: T | ((state: Draft<T>) => void),
+          shouldReplace: true,
+          ...a: SkipTwo<A2>
+        ): Sr2
       }
     : never
   : never
@@ -61,7 +69,7 @@ const immerImpl: ImmerImpl = (initializer) => (set, get, store) => {
       typeof updater === 'function' ? produce(updater as any) : updater
     ) as ((s: T) => T) | T | Partial<T>
 
-    return set(nextState as any, replace, ...a)
+    return set(nextState, replace as any, ...a)
   }
 
   return initializer(store.setState, get, store)

src/middleware/persist.ts

@@ -196,7 +196,7 @@ const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
         console.warn(
           `[zustand persist middleware] Unable to update item '${options.name}', the given storage is currently unavailable.`,
         )
-        set(...args)
+        set(...(args as Parameters<typeof set>))
       },
       get,
       api,
@@ -214,13 +214,13 @@ const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
   const savedSetState = api.setState
 
   api.setState = (state, replace) => {
-    savedSetState(state, replace)
+    savedSetState(state, replace as any)
     void setItem()
   }
 
   const configResult = config(
     (...args) => {
-      set(...args)
+      set(...(args as Parameters<typeof set>))
       void setItem()
     },
     get,

src/vanilla.ts

@@ -1,8 +1,9 @@
 type SetStateInternal<T> = {
   _(
     partial: T | Partial<T> | { _(state: T): T | Partial<T> }['_'],
-    replace?: boolean | undefined,
+    replace?: false,
   ): void
+  _(state: T | { _(state: T): T }['_'], replace: true): void
 }['_']
 
 export interface StoreApi<T> {