diff --git a/package.json b/package.json
index fc39474..38c3941 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "zundo",
-  "version": "2.2.0",
+  "version": "3.0.0-beta.0",
   "private": false,
   "description": "🍜 undo/redo middleware for zustand",
   "keywords": [
diff --git a/src/index.ts b/src/index.ts
index 5fe86d2..454b28b 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -50,11 +50,38 @@ export const temporal = (<TState>(
         temporalStateCreator(set, get, options),
     );
 
+    const handleSet = (
+      pastState: TState,
+      // `replace` will likely be deprecated and removed in the future
+      replace: boolean | undefined,
+      currentState: TState,
+      deltaState?: Partial<TState>,
+    ) => {
+      if (store.temporal.getState().isTracking) {
+        // This naively assumes that only one new state can be added at a time
+        if (
+          options?.limit &&
+          store.temporal.getState().pastStates.length >= options?.limit
+        ) {
+          store.temporal.getState().pastStates.shift();
+        }
+
+        (store.temporal.getState() as _TemporalState<TState>)._onSave?.(
+          pastState,
+          currentState,
+        );
+        store.temporal.setState({
+          pastStates: store.temporal
+            .getState()
+            .pastStates.concat(deltaState || pastState),
+          futureStates: [],
+        });
+      }
+    };
+
     const curriedHandleSet =
-      options?.handleSet?.(
-        (store.temporal.getState() as _TemporalState<TState>)
-          ._handleSet as StoreApi<TState>['setState'],
-      ) || (store.temporal.getState() as _TemporalState<TState>)._handleSet;
+      options?.handleSet?.(handleSet as StoreApi<TState>['setState']) ||
+      handleSet;
 
     const temporalHandleSet = (pastState: TState) => {
       if (!store.temporal.getState().isTracking) return;
diff --git a/src/temporal.ts b/src/temporal.ts
index 2043fa8..cac92be 100644
--- a/src/temporal.ts
+++ b/src/temporal.ts
@@ -58,18 +58,6 @@ export const temporalStateCreator = <TState>(
       setOnSave: (_onSave) => set({ _onSave }),
       // Internal properties
       _onSave: options?.onSave,
-      _handleSet: (pastState, replace, currentState, deltaState) => {
-        // This naively assumes that only one new state can be added at a time
-        if (options?.limit && get().pastStates.length >= options?.limit) {
-          get().pastStates.shift();
-        }
-
-        get()._onSave?.(pastState, currentState);
-        set({
-          pastStates: get().pastStates.concat(deltaState || pastState),
-          futureStates: [],
-        });
-      },
     };
   };
 
diff --git a/src/types.ts b/src/types.ts
index e15162e..75b68c2 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -18,13 +18,6 @@ export interface _TemporalState<TState> {
 
   setOnSave: (onSave: onSave<TState>) => void;
   _onSave: onSave<TState>;
-  _handleSet: (
-    pastState: TState,
-    // `replace` will likely be deprecated and removed in the future
-    replace: Parameters<StoreApi<TState>['setState']>[1],
-    currentState: TState,
-    deltaState?: Partial<TState> | null,
-  ) => void;
 }
 
 export interface ZundoOptions<TState, PartialTState = TState> {
@@ -60,7 +53,4 @@ export interface ZundoOptions<TState, PartialTState = TState> {
 
 export type Write<T, U> = Omit<T, keyof U> & U;
 
-export type TemporalState<TState> = Omit<
-  _TemporalState<TState>,
-  '_onSave' | '_handleSet'
->;
+export type TemporalState<TState> = Omit<_TemporalState<TState>, '_onSave'>;
diff --git a/tests/__tests__/options.test.ts b/tests/__tests__/options.test.ts
index e294590..562405e 100644
--- a/tests/__tests__/options.test.ts
+++ b/tests/__tests__/options.test.ts
@@ -629,6 +629,7 @@ describe('Middleware options', () => {
               console.info('handleSet called');
               _set(partial, replace);
             };
+            store.setState = set;
             return config(set, get, store);
           };
         },
@@ -661,6 +662,28 @@ describe('Middleware options', () => {
       expect(console.info).toHaveBeenCalledTimes(3);
     });
 
+    it('should not call function if `store.setState` is not assigned to `set` (wrapTemporal)', () => {
+      global.console.info = vi.fn();
+      const storeWithHandleSet = createVanillaStore({
+        wrapTemporal: (config) => {
+          return (_set, get, store) => {
+            const set: typeof _set = (partial, replace) => {
+              console.info('handleSet called');
+              _set(partial, replace);
+            };
+            // intentionally commented out
+            // store.setState = set;
+            return config(set, get, store);
+          };
+        },
+      });
+      const { increment } = storeWithHandleSet.getState();
+      act(() => {
+        increment();
+      });
+      expect(console.info).toHaveBeenCalledTimes(0);
+    });
+
     it('should correctly use throttling', () => {
       global.console.error = vi.fn();
       vi.useFakeTimers();
@@ -722,6 +745,7 @@ describe('Middleware options', () => {
               },
               1000,
             );
+            store.setState = set;
             return config(set, get, store);
           };
         },
@@ -964,9 +988,7 @@ describe('Middleware options', () => {
 
   describe('secret internals', () => {
     it('should have a secret internal state', () => {
-      const { _handleSet, _onSave } =
-        store.temporal.getState() as _TemporalState<MyState>;
-      expect(_handleSet).toBeInstanceOf(Function);
+      const { _onSave } = store.temporal.getState() as _TemporalState<MyState>;
       expect(_onSave).toBe(undefined);
     });
     describe('onSave', () => {
@@ -1035,19 +1057,6 @@ describe('Middleware options', () => {
         expect(console.trace).toHaveBeenCalledTimes(1);
       });
     });
-
-    describe('handleUserSet', () => {
-      it('should update the temporal store with the pastState when called', () => {
-        const { _handleSet } =
-          store.temporal.getState() as _TemporalState<MyState>;
-        act(() => {
-          _handleSet(store.getState(), undefined, store.getState(), null);
-        });
-        expect(store.temporal.getState().pastStates.length).toBe(1);
-      });
-
-      // TODO: should this check the equality function, limit, and call onSave? These are already tested but indirectly.
-    });
   });
 
   describe('init pastStates', () => {