diff --git a/jsr.json b/jsr.json
index 835a348d014e3df039d914cd2fa553a6c81c73d8..df5f363d0ad989c18637e77086deec33b96d00d0 100644
--- a/jsr.json
+++ b/jsr.json
@@ -1,6 +1,6 @@
 {
   "name": "@clo/react-mutation",
-  "version": "1.0.0-beta.8",
+  "version": "1.0.0-beta.9",
   "exports": {
     ".": "./src/mod.ts",
     "./tanstack-query.ts": "./src/tanstack-query.ts",
diff --git a/src/blocking.ts b/src/blocking.ts
index 55c376e6de21cde4d1aa24b932d10e79bd61dae9..bfd4e94d017bc5f8a033b304038c55f650fd5a6e 100644
--- a/src/blocking.ts
+++ b/src/blocking.ts
@@ -237,12 +237,12 @@ export class BlockingMutation<
     }
 
     const args = array.slice() as Args;
-    const { onSuccess, onSuccessDataOnly, onError, onSettled } = args
+    const { onSuccess, onSuccessDataOnly, onError, onSettled, onRestore } = args
       .pop() as RunOptions<Result>;
     const suppressGlobalSuccess = onSuccess !== undefined;
     const suppressGlobalError = onError !== undefined;
 
-    const promise = this.runAsPromise(...args);
+    const promise = this.#runAsPromiseWithOptions(args, { onRestore });
     promise.then((result) => {
       // Call user handlers
       onSuccess?.(result);
@@ -275,6 +275,13 @@ export class BlockingMutation<
 
   /** Calls the mutation, treating the errors as promise rejection. */
   runAsPromise(...args: Args): Promise<Result> {
+    return this.#runAsPromiseWithOptions(args, {});
+  }
+
+  #runAsPromiseWithOptions(
+    args: Args,
+    { onRestore: userOnRestore }: Pick<RunOptions<Result>, "onRestore">,
+  ): Promise<Result> {
     if (!this.#client.enabled) {
       throw new Error(
         "MutationClient was passed enabled: false. Are you trying to perform a mutation from SSR?",
@@ -285,7 +292,7 @@ export class BlockingMutation<
 
     // Check if debouncing is enabled
     if (this.#options.debounceMs !== undefined) {
-      return this.#runDebouncedAndReturn(args, key, channel);
+      return this.#runDebouncedAndReturn(args, key, channel, userOnRestore);
     }
 
     // Create shared optimistic helpers instance for the channel if it doesn't exist
@@ -315,6 +322,12 @@ export class BlockingMutation<
       rollbacks += 1;
     };
 
+    // Register user's onRestore callback if provided
+    if (userOnRestore) {
+      channel.rollbacks.push(userOnRestore);
+      rollbacks += 1;
+    }
+
     try {
       this.#options.optimistic({
         args,
@@ -452,6 +465,7 @@ export class BlockingMutation<
     args: Args,
     key: string,
     channel: Channel<Args, Result, Config["optimisticHelpers"]>,
+    userOnRestore?: () => void,
   ): Promise<Result> {
     // If there's a pending debounced call, roll it back
     if (channel.pendingDebounced) {
@@ -485,6 +499,12 @@ export class BlockingMutation<
       rollbacks += 1;
     };
 
+    // Register user's onRestore callback if provided
+    if (userOnRestore) {
+      channel.rollbacks.push(userOnRestore);
+      rollbacks += 1;
+    }
+
     try {
       this.#options.optimistic({
         args,
diff --git a/src/debounced.ts b/src/debounced.ts
index 910f0c8322ad5e0f50253e3136cc1ef1259a1b84..4b4cf388bf50da5e3ab03acce4083186dbed27d8 100644
--- a/src/debounced.ts
+++ b/src/debounced.ts
@@ -285,7 +285,7 @@ export class DebouncedMutation<
         "MutationClient was passed enabled: false. Are you trying to perform a mutation from SSR?",
       );
     }
-    this.#runAndReturn(args, true).catch((error) => {
+    this.#runAndReturn(args, true, undefined).catch((error) => {
       const message = `Failed to ${this.describe(...args)}: ${
         errMessage(error)
       }`;
@@ -300,12 +300,12 @@ export class DebouncedMutation<
       );
     }
     const args = array.slice() as Args;
-    const { onSuccess, onSuccessDataOnly, onError, onSettled } = args
+    const { onSuccess, onSuccessDataOnly, onError, onSettled, onRestore } = args
       .pop() as RunOptions<Result>;
     const suppressGlobalSuccess = onSuccess !== undefined;
     const suppressGlobalError = onError !== undefined;
 
-    const promise = this.#runAndReturn(args, !suppressGlobalSuccess);
+    const promise = this.#runAndReturn(args, !suppressGlobalSuccess, onRestore);
 
     promise.then((result) => {
       // Call user handlers
@@ -334,10 +334,14 @@ export class DebouncedMutation<
         "MutationClient was passed enabled: false. Are you trying to perform a mutation from SSR?",
       );
     }
-    return this.#runAndReturn(args, false);
+    return this.#runAndReturn(args, false, undefined);
   }
 
-  #runAndReturn(args: Args, reportSuccessGlobally: boolean): Promise<Result> {
+  #runAndReturn(
+    args: Args,
+    reportSuccessGlobally: boolean,
+    userOnRestore?: () => void,
+  ): Promise<Result> {
     const key = this.key(args);
     const channel = this.#getOrPutChannel(key);
 
@@ -374,6 +378,11 @@ export class DebouncedMutation<
       channel.rollbacks.push(cb);
     };
 
+    // Register user's onRestore callback if provided
+    if (userOnRestore) {
+      channel.rollbacks.push(userOnRestore);
+    }
+
     try {
       this.#options.optimistic(
         {
diff --git a/src/react.ts b/src/react.ts
index 15caaee76e046ecc7823d371e48eacc462e5094a..aaf73159b18f7e394c43ea52bd7467449c190c33 100644
--- a/src/react.ts
+++ b/src/react.ts
@@ -45,6 +45,7 @@ export interface UseMutateResultBase<Args extends unknown[], Result> {
     ..._: [...args: Args, options: RunOptions<Result>]
   ) => Promise<Result>;
   clear: () => void;
+  setError: (error: unknown) => void;
   args: Args | undefined;
 }
 
@@ -98,7 +99,7 @@ export interface UseMutateIdle {
   isOptimisticData: boolean;
 }
 
-type AnyMutationStateWithoutRun<Result> =
+type AnyMutationStateWithoutRun<Args extends unknown[], Result> =
   & Omit<
     UseMutateIdle,
     "status" | "result" | "error" | "isSuccess" | "isError" | "errorMessage"
@@ -110,10 +111,11 @@ type AnyMutationStateWithoutRun<Result> =
     errorMessage: undefined | string;
     isSuccess: boolean;
     isError: boolean;
+    args: Args | undefined;
   };
 
 export type AnyMutationState<Args extends unknown[], Result> =
-  & AnyMutationStateWithoutRun<Result>
+  & AnyMutationStateWithoutRun<Args, Result>
   & UseMutateResultBase<Args, Result>;
 
 function initialState() {
@@ -127,6 +129,7 @@ function initialState() {
     isSuccess: false,
     isError: false,
     isOptimisticData: false,
+    args: undefined,
   } as const;
 }
 
@@ -135,15 +138,14 @@ class Observer<Args extends unknown[], Result> {
   mutation: Mutation<Args, Result> | null = null;
   unsubscribe: (() => void) | null = null;
   currentKey: string | null = null;
-  currentArgs: Args | null = null;
 
   constructor(setRerender: (fn: number) => void) {
     this.setRerender = setRerender;
   }
 
   watched: Set<string> = new Set();
-  state: AnyMutationStateWithoutRun<Result> = initialState();
-  setState(newState: Partial<AnyMutationStateWithoutRun<Result>>) {
+  state: AnyMutationStateWithoutRun<Args, Result> = initialState();
+  setState(newState: Partial<AnyMutationStateWithoutRun<Args, Result>>) {
     let updateUi = false;
     const current: Record<string, unknown> = this.state;
     for (const [key, value] of Object.entries(newState)) {
@@ -167,8 +169,8 @@ class Observer<Args extends unknown[], Result> {
   computeErrorMessage(error: unknown): string | undefined {
     if (!error) return undefined;
     const mutation = this.mutation;
-    if (!mutation || !this.currentArgs) return errMessage(error);
-    return `Failed to ${mutation.describe(...this.currentArgs)}: ${
+    if (!mutation || !this.state.args) return errMessage(error);
+    return `Failed to ${mutation.describe(...this.state.args)}: ${
       errMessage(error)
     }`;
   }
@@ -176,11 +178,7 @@ class Observer<Args extends unknown[], Result> {
   run(...args: Args) {
     const mutation = this.mutation;
     if (!mutation) return;
-    const argsChanged = this.currentArgs !== args;
-    this.currentArgs = args;
-    if (argsChanged && this.watched.has("args")) {
-      this.setRerender(Math.random());
-    }
+    this.setState({ args });
     const key = mutation.key(args);
     if (key !== this.currentKey) {
       this.currentKey = key;
@@ -216,6 +214,7 @@ class Observer<Args extends unknown[], Result> {
             isError: hasError,
             isOptimisticData: status === "waiting" || status === "mutating" ||
               status === "refetching",
+            args: hasError || hasResult ? undefined : this.state.args,
           });
         },
       );
@@ -254,11 +253,7 @@ class Observer<Args extends unknown[], Result> {
     const args = array.slice() as Args;
     const options = args.pop() as RunOptions<Result>;
 
-    const argsChanged = this.currentArgs !== args;
-    this.currentArgs = args;
-    if (argsChanged && this.watched.has("args")) {
-      this.setRerender(Math.random());
-    }
+    this.setState({ args });
     const key = mutation.key(args);
 
     // Set up subscription if key changed
@@ -296,6 +291,7 @@ class Observer<Args extends unknown[], Result> {
             isError: hasError,
             isOptimisticData: status === "waiting" || status === "mutating" ||
               status === "refetching",
+            args: hasError || hasResult ? undefined : this.state.args,
           });
         },
       );
@@ -320,6 +316,16 @@ class Observer<Args extends unknown[], Result> {
         result: undefined,
       });
     },
+    setError(error: unknown) {
+      self.setState({
+        status: "error",
+        error,
+        errorMessage: errMessage(error),
+        isError: true,
+        isSuccess: false,
+        result: undefined,
+      });
+    },
     get status() {
       self.watched.add("status");
       return self.state.status;
@@ -358,7 +364,7 @@ class Observer<Args extends unknown[], Result> {
     },
     get args() {
       self.watched.add("args");
-      return self.currentArgs ?? undefined;
+      return self.state.args;
     },
   } as UseMutateResult<Args, Result>))(this);
 }
diff --git a/src/types.ts b/src/types.ts
index 527d087474b91bd64fbdef8eb429f0281de0bc02..89f3d15415d8b5f01b84d8573dd76078a520ce52 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -33,6 +33,8 @@ export interface RunOptions<Result> {
       | { status: "success"; result: Result }
       | { status: "error"; error: unknown },
   ) => void;
+  /** Called when optimistic state is being restored/rolled back */
+  onRestore?: () => void;
 }
 
 export interface MutationEvent<Result> {