From feddc0c210c555fc58d9c9a7100545603f568cf5 Mon Sep 17 00:00:00 2001
From: Daishi Kato <dai-shi@users.noreply.github.com>
Date: Wed, 20 Aug 2025 08:08:35 +0900
Subject: [PATCH] feat(middleare/persist): return storage promise from setState
 (#3206)

* feat(middleare/persist): return storage promise from setState

* refactor types (technically breaking)

* another breaking change in types

* make public types not breaking
---
 src/middleware/persist.ts   | 73 ++++++++++++++++++++++---------------
 tests/persistAsync.test.tsx |  8 +++-
 2 files changed, 49 insertions(+), 32 deletions(-)

diff --git a/src/middleware/persist.ts b/src/middleware/persist.ts
index 357f3e62..5a2dfe5a 100644
--- a/src/middleware/persist.ts
+++ b/src/middleware/persist.ts
@@ -4,10 +4,10 @@ import type {
   StoreMutatorIdentifier,
 } from '../vanilla.ts'
 
-export interface StateStorage {
+export interface StateStorage<R = unknown> {
   getItem: (name: string) => string | null | Promise<string | null>
-  setItem: (name: string, value: string) => unknown | Promise<unknown>
-  removeItem: (name: string) => unknown | Promise<unknown>
+  setItem: (name: string, value: string) => R
+  removeItem: (name: string) => R
 }
 
 export type StorageValue<S> = {
@@ -15,12 +15,12 @@ export type StorageValue<S> = {
   version?: number
 }
 
-export interface PersistStorage<S> {
+export interface PersistStorage<S, R = unknown> {
   getItem: (
     name: string,
   ) => StorageValue<S> | null | Promise<StorageValue<S> | null>
-  setItem: (name: string, value: StorageValue<S>) => unknown | Promise<unknown>
-  removeItem: (name: string) => unknown | Promise<unknown>
+  setItem: (name: string, value: StorageValue<S>) => R
+  removeItem: (name: string) => R
 }
 
 type JsonStorageOptions = {
@@ -28,18 +28,18 @@ type JsonStorageOptions = {
   replacer?: (key: string, value: unknown) => unknown
 }
 
-export function createJSONStorage<S>(
-  getStorage: () => StateStorage,
+export function createJSONStorage<S, R = unknown>(
+  getStorage: () => StateStorage<R>,
   options?: JsonStorageOptions,
-): PersistStorage<S> | undefined {
-  let storage: StateStorage | undefined
+): PersistStorage<S, unknown> | undefined {
+  let storage: StateStorage<R> | undefined
   try {
     storage = getStorage()
   } catch {
     // prevent error if the storage is not defined (e.g. when server side rendering a page)
     return
   }
-  const persistStorage: PersistStorage<S> = {
+  const persistStorage: PersistStorage<S, R> = {
     getItem: (name) => {
       const parse = (str: string | null) => {
         if (str === null) {
@@ -60,7 +60,11 @@ export function createJSONStorage<S>(
   return persistStorage
 }
 
-export interface PersistOptions<S, PersistedState = S> {
+export interface PersistOptions<
+  S,
+  PersistedState = S,
+  PersistReturn = unknown,
+> {
   /** Name of the storage (must be unique) */
   name: string
   /**
@@ -71,7 +75,7 @@ export interface PersistOptions<S, PersistedState = S> {
    *
    * @default createJSONStorage(() => localStorage)
    */
-  storage?: PersistStorage<PersistedState> | undefined
+  storage?: PersistStorage<PersistedState, PersistReturn> | undefined
   /**
    * Filter the persisted value.
    *
@@ -118,17 +122,28 @@ export interface PersistOptions<S, PersistedState = S> {
 
 type PersistListener<S> = (state: S) => void
 
-type StorePersist<S, Ps> = {
-  persist: {
-    setOptions: (options: Partial<PersistOptions<S, Ps>>) => void
-    clearStorage: () => void
-    rehydrate: () => Promise<void> | void
-    hasHydrated: () => boolean
-    onHydrate: (fn: PersistListener<S>) => () => void
-    onFinishHydration: (fn: PersistListener<S>) => () => void
-    getOptions: () => Partial<PersistOptions<S, Ps>>
+type StorePersist<S, Ps, Pr> = S extends {
+  getState: () => infer T
+  setState: {
+    // capture both overloads of setState
+    (...args: infer Sa1): infer Sr1
+    (...args: infer Sa2): infer Sr2
   }
 }
+  ? {
+      setState(...args: Sa1): Sr1 | Pr
+      setState(...args: Sa2): Sr2 | Pr
+      persist: {
+        setOptions: (options: Partial<PersistOptions<T, Ps, Pr>>) => void
+        clearStorage: () => void
+        rehydrate: () => Promise<void> | void
+        hasHydrated: () => boolean
+        onHydrate: (fn: PersistListener<T>) => () => void
+        onFinishHydration: (fn: PersistListener<T>) => () => void
+        getOptions: () => Partial<PersistOptions<T, Ps, Pr>>
+      }
+    }
+  : never
 
 type Thenable<Value> = {
   then<V>(
@@ -172,7 +187,7 @@ const toThenable =
 const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
   type S = ReturnType<typeof config>
   let options = {
-    storage: createJSONStorage<S>(() => localStorage),
+    storage: createJSONStorage<S, void>(() => localStorage),
     partialize: (state: S) => state,
     version: 0,
     merge: (persistedState: unknown, currentState: S) => ({
@@ -202,7 +217,7 @@ const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
 
   const setItem = () => {
     const state = options.partialize({ ...get() })
-    return (storage as PersistStorage<S>).setItem(options.name, {
+    return (storage as PersistStorage<S, unknown>).setItem(options.name, {
       state,
       version: options.version,
     })
@@ -212,13 +227,13 @@ const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
 
   api.setState = (state, replace) => {
     savedSetState(state, replace as any)
-    void setItem()
+    return setItem()
   }
 
   const configResult = config(
     (...args) => {
       set(...(args as Parameters<typeof set>))
-      void setItem()
+      return setItem()
     },
     get,
     api,
@@ -307,7 +322,7 @@ const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
       })
   }
 
-  ;(api as StoreApi<S> & StorePersist<S, S>).persist = {
+  ;(api as StoreApi<S> & StorePersist<StoreApi<S>, S, unknown>).persist = {
     setOptions: (newOptions) => {
       options = {
         ...options,
@@ -365,9 +380,7 @@ declare module '../vanilla' {
 
 type Write<T, U> = Omit<T, keyof U> & U
 
-type WithPersist<S, A> = S extends { getState: () => infer T }
-  ? Write<S, StorePersist<T, A>>
-  : never
+type WithPersist<S, A> = Write<S, StorePersist<S, A, unknown>>
 
 type PersistImpl = <T>(
   storeInitializer: StateCreator<T, [], []>,
diff --git a/tests/persistAsync.test.tsx b/tests/persistAsync.test.tsx
index 4e113389..37ded00b 100644
--- a/tests/persistAsync.test.tsx
+++ b/tests/persistAsync.test.tsx
@@ -165,7 +165,9 @@ describe('persist middleware with async configuration', () => {
     })
 
     // Write something to the store
-    act(() => useBoundStore.setState({ count: 42 }))
+    act(() => {
+      useBoundStore.setState({ count: 42 })
+    })
     expect(await screen.findByText('count: 42')).toBeInTheDocument()
     expect(setItemSpy).toBeCalledWith(
       'test-storage',
@@ -788,7 +790,9 @@ describe('persist middleware with async configuration', () => {
 
     // Write something to the store
     const updatedMap = new Map(map).set('foo', 'bar')
-    act(() => useBoundStore.setState({ map: updatedMap }))
+    act(() => {
+      useBoundStore.setState({ map: updatedMap })
+    })
     expect(await screen.findByText('map-content: bar')).toBeInTheDocument()
 
     expect(setItemSpy).toBeCalledWith(