fix(hooks): prevent stale state access in callbacks (#377)

Previously, callbacks passed to ‎`useActionCallbacks` could become stale
if their references changed between renders

Co-authored-by: Edoardo Ranghieri <edoardo@ranghieri.com>

Author: rortan134

apps/playground/src/app/(examples)/state-update/page.tsx

@@ -0,0 +1,53 @@
+"use client";
+
+import { StyledButton } from "@/app/_components/styled-button";
+import { StyledHeading } from "@/app/_components/styled-heading";
+import { useAction } from "next-safe-action/hooks";
+import { useState } from "react";
+import { ResultBox } from "../../_components/result-box";
+import { stateUpdateAction } from "./stateupdate-action";
+
+export default function StateUpdate() {
+	const [count, setCount] = useState(0);
+	// Safe action (`deleteUser`) and optional callbacks passed to `useAction` hook.
+	const { execute, result, status, reset } = useAction(stateUpdateAction, {
+		onSuccess(args) {
+			console.log("onSuccess callback:", args);
+			console.log("Count value:", count);
+		},
+		onError(args) {
+			console.log("onError callback:", args);
+		},
+		onNavigation(args) {
+			console.log("onNavigation callback:", args);
+		},
+		onSettled(args) {
+			console.log("onSettled callback:", args);
+		},
+		onExecute(args) {
+			console.log("onExecute callback:", args);
+		},
+	});
+
+	return (
+		<main className="w-96 max-w-full px-4">
+			<StyledHeading>State update</StyledHeading>
+			<div className="mt-4 flex flex-col gap-2">
+				<StyledButton type="button" onClick={() => setCount(count + 1)}>
+					Increment
+				</StyledButton>
+				<StyledButton type="button" onClick={() => setCount(count - 1)}>
+					Decrement
+				</StyledButton>
+				<StyledButton type="button" onClick={() => execute()}>
+					Execute
+				</StyledButton>
+				<StyledButton type="button" onClick={reset}>
+					Reset
+				</StyledButton>
+			</div>
+			<p className="mt-4">Count value: {count}</p>
+			<ResultBox result={result} status={status} />
+		</main>
+	);
+}

apps/playground/src/app/(examples)/state-update/stateupdate-action.ts

@@ -0,0 +1,11 @@
+"use server";
+
+import { action } from "@/lib/safe-action";
+
+export const stateUpdateAction = action.metadata({ actionName: "stateUpdateAction" }).action(async () => {
+	await new Promise((res) => setTimeout(res, 1000));
+
+	return {
+		message: "Hello, world!",
+	};
+});

apps/playground/src/app/page.tsx

@@ -21,6 +21,7 @@ export default function Home() {
 				<ExampleLink href="/stateful-form">
 					Stateful form (<span className="font-mono">useActionState()</span> hook)
 				</ExampleLink>
+				<ExampleLink href="/state-update">State update</ExampleLink>
 				<ExampleLink href="/navigation">Navigation</ExampleLink>
 				<ExampleLink href="/file-upload">File upload</ExampleLink>
 				<ExampleLink href="/bind-arguments">Bind arguments</ExampleLink>

packages/next-safe-action/src/hooks-utils.ts

@@ -51,6 +51,18 @@ export const getActionShorthandStatusObject = (status: HookActionStatus): HookSh
 	};
 };
 
+/**
+ * Converts a callback to a ref to avoid triggering re-renders when passed as a
+ * prop or avoid re-executing effects when passed as a dependency
+ */
+function useCallbackRef<T extends (arg: any) => any>(callback: T | undefined): T {
+	const callbackRef = React.useRef(callback);
+	React.useEffect(() => {
+		callbackRef.current = callback;
+	});
+	return React.useMemo(() => ((arg) => callbackRef.current?.(arg) as T) as T, []);
+}
+
 export const useActionCallbacks = <ServerError, S extends StandardSchemaV1 | undefined, CVE, Data>({
 	result,
 	input,
@@ -66,20 +78,14 @@ export const useActionCallbacks = <ServerError, S extends StandardSchemaV1 | und
 	navigationError: Error | null;
 	thrownError: Error | null;
 }) => {
-	const onExecuteRef = React.useRef(cb?.onExecute);
-	const onSuccessRef = React.useRef(cb?.onSuccess);
-	const onErrorRef = React.useRef(cb?.onError);
-	const onSettledRef = React.useRef(cb?.onSettled);
-	const onNavigationRef = React.useRef(cb?.onNavigation);
+	const onExecute = useCallbackRef(cb?.onExecute);
+	const onSuccess = useCallbackRef(cb?.onSuccess);
+	const onError = useCallbackRef(cb?.onError);
+	const onSettled = useCallbackRef(cb?.onSettled);
+	const onNavigation = useCallbackRef(cb?.onNavigation);
 
 	// Execute the callback when the action status changes.
-	React.useEffect(() => {
-		const onExecute = onExecuteRef.current;
-		const onSuccess = onSuccessRef.current;
-		const onError = onErrorRef.current;
-		const onSettled = onSettledRef.current;
-		const onNavigation = onNavigationRef.current;
-
+	React.useLayoutEffect(() => {
 		const executeCallbacks = async () => {
 			switch (status) {
 				case "executing":
@@ -99,7 +105,12 @@ export const useActionCallbacks = <ServerError, S extends StandardSchemaV1 | und
 					break;
 				case "hasErrored":
 					await Promise.all([
-						Promise.resolve(onError?.({ error: { ...result, ...(thrownError ? { thrownError } : {}) }, input })),
+						Promise.resolve(
+							onError?.({
+								error: { ...result, ...(thrownError ? { thrownError } : {}) },
+								input,
+							})
+						),
 						Promise.resolve(onSettled?.({ result, input })),
 					]);
 					break;
@@ -128,5 +139,5 @@ export const useActionCallbacks = <ServerError, S extends StandardSchemaV1 | und
 		};
 
 		executeCallbacks().catch(console.error);
-	}, [input, status, result, navigationError, thrownError]);
+	}, [input, status, result, navigationError, thrownError, onExecute, onSuccess, onSettled, onError, onNavigation]);
 };