diff --git a/docs/guide/demos/Redo/store.ts b/docs/guide/demos/Redo/store.ts
index febae34a..aa0ef1f0 100644
--- a/docs/guide/demos/Redo/store.ts
+++ b/docs/guide/demos/Redo/store.ts
@@ -1,6 +1,6 @@
 import { proEditorMiddleware, ProEditorOptions } from '@ant-design/pro-editor';
 import { create, StateCreator } from 'zustand';
-import { devtools } from 'zustand/middleware';
+import { devtools, subscribeWithSelector } from 'zustand/middleware';
 
 interface Store {
   tabs: string;
@@ -49,5 +49,9 @@ const proEditorOptions: ProEditorOptions<Store, ProEditorStore> = {
 };
 
 export const useStore = create<Store>()(
-  devtools(proEditorMiddleware(createStore, proEditorOptions), { name: storeName }),
+  devtools(proEditorMiddleware(subscribeWithSelector(createStore), proEditorOptions), {
+    name: storeName,
+  }),
 );
+
+useStore.subscribe((s) => s.data, console.log);
diff --git a/docs/guide/redo-undo.md b/docs/guide/redo-undo.md
index 0d7adb69..758eb316 100644
--- a/docs/guide/redo-undo.md
+++ b/docs/guide/redo-undo.md
@@ -19,14 +19,13 @@ order: 2
 1. 外层包裹 ProEditorProvider，传入相应的 zustand store
 
 ```tsx | pure
-import { ProEditorProvider, ProEditorStoreUpdater } from '@ant-design/pro-editor';
+import { ProEditorProvider } from '@ant-design/pro-editor';
 
 import { useStore } from './store';
 
-
 export default () => {
   return (
-    <ProEditorProvider store=[useStore]>
+    <ProEditorProvider store={[useStore]}>
       <App />
     </ProEditorProvider>
   );
@@ -56,15 +55,14 @@ export const useStore = create<Store>()(
 多个 Store 使用的方式：
 
 ```tsx | pure
-import { ProEditorProvider, ProEditorStoreUpdater } from '@ant-design/pro-editor';
+import { ProEditorProvider } from '@ant-design/pro-editor';
 
 import { useAStore } from './storeA';
 import { useBStore } from './storeB';
 
-
 export default () => {
   return (
-    <ProEditorProvider  store=[useAStore,useBStore]>
+    <ProEditorProvider store={[useAStore, useBStore]}>
       <App />
     </ProEditorProvider>
   );
@@ -75,11 +73,11 @@ export default () => {
 
 ```tsx | pure
 <div>
-  <ProEditorProvider store=[cStore]>
+  <ProEditorProvider store={[cStore]}>
     <A />
   </ProEditorProvider>
 
-  <ProEditorProvider store=[dStore]>
+  <ProEditorProvider store={[dStore]}>
     <B />
   </ProEditorProvider>
 </div>
diff --git a/package.json b/package.json
index 33b2b1be..fa68c600 100644
--- a/package.json
+++ b/package.json
@@ -121,7 +121,7 @@
   },
   "devDependencies": {
     "@emotion/jest": "^11.11.0",
-    "@testing-library/jest-dom": "^6.1.2",
+    "@testing-library/jest-dom": "^6.1.3",
     "@testing-library/react": "^13.4.0",
     "@testing-library/user-event": "^14.5.1",
     "@types/color": "^3.0.4",
diff --git a/src/ProEditor/middleware/pro-editor/index.ts b/src/ProEditor/middleware/pro-editor/index.ts
index 2ca0a3f6..ff8dbf35 100644
--- a/src/ProEditor/middleware/pro-editor/index.ts
+++ b/src/ProEditor/middleware/pro-editor/index.ts
@@ -33,6 +33,17 @@ const middleware: ProEditorImpl = (storeInitializer, options) => (set, get, api)
       { trigger: 'proEditorMiddleware', ...action },
     );
   };
+
+  /**
+   * handle setState function
+   */
+  const savedSetState = api.setState;
+  api.setState = (partial, replace, action) => {
+    savedSetState(partial, replace, action);
+
+    updateInProEditor((action as any) || {});
+  };
+
   /*
    * Capture the initial state so that we can initialize the pro editor store to the
    * same values as the initial values of the Zustand store.
@@ -47,15 +58,7 @@ const middleware: ProEditorImpl = (storeInitializer, options) => (set, get, api)
       updateInProEditor((action as any) || {});
     },
     get,
-    {
-      ...api,
-      // Create a new setState function as we did with set.
-      setState: (partial, replace, action) => {
-        api.setState(partial, replace, action);
-
-        updateInProEditor((action as any) || {});
-      },
-    },
+    api,
   );
 
   // Return the initial state to create or the next middleware.