A Custom hook for automatic, type-safe react-query optimistic updates with proper rollback. Use together with openapi-typescript + openapi-react-query.

use-optimistic-update.ts

import { type paths } from "@/types/api-types";
import { queryClient } from "@/api/api";
import { toast } from "sonner";

type OperationState<TData> = {
  previousData: TData;
  timestamp: number;
};

// Constants for stack management
const MAX_STACK_SIZE_PER_QUERY = 100;
const STALE_OPERATION_THRESHOLD = 30_000; // 30 seconds

const operationStack = new Map<string, OperationState<unknown>[]>();

const cleanupStaleOperations = (queryKeyString: string) => {
  const operations = operationStack.get(queryKeyString);
  if (!operations) return;

  const now = Date.now();
  const validOperations = operations.filter(
    (op) => now - op.timestamp < STALE_OPERATION_THRESHOLD,
  );

  if (validOperations.length === 0) {
    operationStack.delete(queryKeyString);
  } else {
    operationStack.set(queryKeyString, validOperations);
  }
};

type GetOperations = {
  [P in keyof paths]: paths[P]["get"] extends { responses: any } ? P : never;
}[keyof paths];

type GetResponseType<P extends GetOperations> = paths[P]["get"] extends {
  responses: { default: { content: { "application/json": infer T } } };
}
  ? T
  : paths[P]["get"] extends {
        responses: { 200: { content: { "application/json": infer T } } };
      }
    ? T
    : never;

type MutationType = "patch" | "post" | "delete";

type MutationHandlers<TData, TVariables> = {
  /**
   * Custom handler for optimistic update mutations.
   * Executed immediately when mutation is triggered, before the server request.
   *
   * Use this handler when you need to:
   * - Perform additional actions before the optimistic update
   * - Cancel related queries
   * - Modify the cache in a custom way
   * - Store additional context for rollback
   *
   * @param variables - The variables passed to the mutation
   * @param defaultHandler - The default handler that performs standard optimistic update
   *
   * @returns Promise containing the previous data (used for rollback if needed)
   *
   * @example
   * ```typescript
   * onMutate: async (variables, defaultHandler) => {
   *   // Cancel any outgoing refetches that might override our optimistic update
   *   await queryClient.cancelQueries({ queryKey: ['otherData'] });
   *
   *   // Perform the default optimistic update
   *   const result = await defaultHandler(variables);
   *
   *   // You can store additional context if needed
   *   return {
   *     ...result,
   *     additionalContext: 'some-value'
   *   };
   * }
   * ```
   */
  onMutate?: (
    variables: TVariables,
    defaultHandler: (
      variables: TVariables,
    ) => Promise<{ previousData: TData | undefined }>,
  ) => Promise<{ previousData: TData | undefined }>;

  /**
   * Custom error handler for failed mutations.
   * Executed when the server request fails or throws an error.
   *
   * Use this handler when you need to:
   * - Perform custom error handling
   * - Show custom error messages
   * - Execute additional rollback logic
   * - Log errors in a specific way
   *
   * @param error - The error that occurred during mutation
   * @param variables - The variables that were passed to the mutation
   * @param context - The context returned from onMutate (contains previousData)
   * @param defaultHandler - The default handler that performs standard error handling and rollback
   *
   * @example
   * ```typescript
   * onError: (error, variables, context, defaultHandler) => {
   *   // Custom error logging
   *   console.error(`Failed to update item ${variables.params.path.id}`, error);
   *
   *   // Custom error message based on error type
   *   if (error.message.includes('Permission denied')) {
   *     toast.error('You don\'t have permission to perform this action');
   *   } else {
   *     // Use default error handling
   *     defaultHandler(error, variables, context);
   *   }
   *
   *   // Additional cleanup if needed
   *   someCleanupFunction();
   * }
   * ```
   */
  onError?: (
    error: Error,
    variables: TVariables,
    context: any,
    defaultHandler: (error: Error, variables: TVariables, context: any) => void,
  ) => void;
};
type GetPatchVariables<P extends keyof paths> = paths[P]["patch"] extends {
  parameters: infer Params;
  requestBody: { content: { "application/json": infer Body } };
}
  ? { params: Params; body: Body }
  : never;

type GetPostVariables<P extends keyof paths> = paths[P]["post"] extends {
  parameters: { path: infer PathParams };
  requestBody: { content: { "application/json": infer Body } };
}
  ? { params: { path: PathParams }; body: Body }
  : paths[P]["post"] extends {
        requestBody: { content: { "application/json": infer Body } };
      }
    ? Body
    : never;

type GetDeleteVariables<P extends keyof paths> = paths[P]["delete"] extends {
  parameters: infer T;
}
  ? { params: T }
  : never;

type GetMutationVariables<
  P extends keyof paths,
  T extends MutationType,
> = T extends "patch"
  ? GetPatchVariables<P>
  : T extends "post"
    ? GetPostVariables<P>
    : T extends "delete"
      ? GetDeleteVariables<P>
      : never;

/**
 * A hook for handling optimistic updates in TanStack Query mutations.
 *
 * This hook provides automatic optimistic updates for POST, PATCH, and DELETE operations
 * by updating the query cache immediately and rolling back if the mutation fails.
 * It supports both default behavior and custom handlers.
 *
 * Features:
 * - Full TypeScript support with proper type inference
 * - Automatic error handling with rollback
 * - Toast notifications for errors
 * - Support for custom mutation handlers
 *
 * @template GetPath - The API endpoint path for the GET operation
 * @template MutationPath - The API endpoint path for the mutation operation
 * @template TMutationType - The type of mutation ("patch" | "post" | "delete")
 *
 * @param queryKey - The query key array for TanStack Query [method, path, params]
 * @param mutationPath - The path of the mutation endpoint
 * @param config - Configuration object containing update function and optional handlers
 *
 * @example
 * // Example 1: Updating a report status (PATCH)
 * const updateStatusMutation = $api.useMutation("patch", "/reports/{id}/status", {
 *   ...useOptimisticUpdate(
 *     ["get", "/reports", getReportsVariables],
 *     "/reports/{id}/status",
 *     {
 *       type: "patch",
 *       updateFn: (oldData, variables) => ({
 *         ...oldData,
 *         data: oldData.data.map(report =>
 *           report.id === variables.params.path.id
 *             ? { ...report, status: variables.body.status }
 *             : report
 *         )
 *       })
 *     }
 *   )
 * });
 *
 * // Example 2: Adding a new note (POST)
 * const addNoteMutation = $api.useMutation("post", "/feedback/{id}/notes", {
 *   ...useOptimisticUpdate(
 *     ["get", "/feedback", { params: { query: { page: 1, pageSize: 200 } } }],
 *     "/feedback/{id}/notes",
 *     {
 *       type: "post",
 *       updateFn: (oldData, variables) => ({
 *         ...oldData,
 *         data: [
 *           {
 *             id: `temp-${Date.now()}`,
 *             content: variables.body.content,
 *             created_at: new Date().toISOString(),
 *           },
 *           ...oldData.data,
 *         ],
 *         total: oldData.total + 1,
 *       })
 *     }
 *   )
 * });
 *
 * // Example 3: Deleting an item (DELETE)
 * const deleteItemMutation = $api.useMutation("delete", "/items/{id}", {
 *   ...useOptimisticUpdate(
 *     ["get", "/items", queryParams],
 *     "/items/{id}",
 *     {
 *       type: "delete",
 *       updateFn: (oldData, variables) => ({
 *         ...oldData,
 *         data: oldData.data.filter(
 *           item => item.id !== variables.params.path.id
 *         ),
 *         total: oldData.total - 1,
 *       })
 *     }
 *   )
 * });
 *
 * // Example 4: Using custom handlers
 * const mutation = $api.useMutation("patch", "/items/{id}", {
 *   ...useOptimisticUpdate(
 *     ["get", "/items", queryParams],
 *     "/items/{id}",
 *     {
 *       type: "patch",
 *       updateFn: (old, variables) => ({
 *         ...old,
 *         data: old.data.map(item =>
 *           item.id === variables.params.path.id
 *             ? { ...item, ...variables.body }
 *             : item
 *         )
 *       }),
 *       handlers: {
 *         onMutate: async (variables, defaultHandler) => {
 *           // Custom logic before mutation
 *           const result = await defaultHandler(variables);
 *           // Custom logic after mutation
 *           return result;
 *         },
 *         onError: (error, variables, context, defaultHandler) => {
 *           // Custom error handling
 *           defaultHandler(error, variables, context);
 *         }
 *       }
 *     }
 *   )
 * });
 *
 * @returns An object containing onMutate and onError handlers for the mutation
 */
export function useOptimisticUpdate<
  GetPath extends GetOperations,
  MutationPath extends keyof paths,
  TMutationType extends MutationType,
>(
  queryKey: [method: "get", path: GetPath, params: any],
  mutationPath: MutationPath,
  config: {
    type: TMutationType;
    updateFn: (
      old: GetResponseType<GetPath>,
      variables: GetMutationVariables<MutationPath, TMutationType>,
    ) => GetResponseType<GetPath>;
    handlers?: MutationHandlers<
      GetResponseType<GetPath>,
      GetMutationVariables<MutationPath, TMutationType>
    >;
  },
) {
  const { type, updateFn, handlers } = config;
  const queryKeyString = JSON.stringify(queryKey);

  const defaultOnMutate = async (
    variables: GetMutationVariables<MutationPath, TMutationType>,
  ) => {
    cleanupStaleOperations(queryKeyString);

    await queryClient.cancelQueries({ queryKey });
    const previousData =
      queryClient.getQueryData<GetResponseType<GetPath>>(queryKey);

    if (previousData) {
      const operations = operationStack.get(queryKeyString) || [];

      if (operations.length >= MAX_STACK_SIZE_PER_QUERY) {
        const excessCount = operations.length - MAX_STACK_SIZE_PER_QUERY + 1;
        operations.splice(0, excessCount);
        console.warn(
          `Operation stack for ${queryKeyString} exceeded ${MAX_STACK_SIZE_PER_QUERY} items. Removed ${excessCount} oldest operations.`,
        );
      }

      operationStack.set(queryKeyString, [
        ...operations,
        { previousData, timestamp: Date.now() },
      ]);

      queryClient.setQueryData<GetResponseType<GetPath>>(queryKey, (old) => {
        if (!old) return old;
        return updateFn(old, variables);
      });
    }

    return { previousData, timestamp: Date.now() };
  };

  const defaultOnError = (
    error: Error,
    variables: GetMutationVariables<MutationPath, TMutationType>,
    context: { previousData: GetResponseType<GetPath>; timestamp: number },
  ) => {
    const operations = operationStack.get(queryKeyString) || [];

    const operationIndex = operations.findIndex(
      (op) => op.timestamp === context.timestamp,
    );

    if (operationIndex !== -1) {
      const operation = operations[operationIndex];
      if (!operation) return;

      operations.splice(operationIndex);

      if (operations.length === 0) {
        operationStack.delete(queryKeyString);
      } else {
        operationStack.set(queryKeyString, operations);
      }

      queryClient.setQueryData(queryKey, operation.previousData);
    }

    toast.error(`Operation failed`, { description: error.message });
  };

  return {
    onMutate: async (
      variables: GetMutationVariables<MutationPath, TMutationType>,
    ) => {
      if (handlers?.onMutate) {
        return handlers.onMutate(variables, defaultOnMutate);
      }
      return defaultOnMutate(variables);
    },
    onError: (
      error: Error,
      variables: GetMutationVariables<MutationPath, TMutationType>,
      context: any,
    ) => {
      if (handlers?.onError) {
        handlers.onError(error, variables, context, defaultOnError);
      } else {
        defaultOnError(error, variables, context);
      }
    },
    onSettled: (
      data: unknown,
      error: Error | null,
      variables: GetMutationVariables<MutationPath, TMutationType>,
      context: unknown,
    ) => {
      if (
        !context ||
        typeof context !== "object" ||
        !("timestamp" in context)
      ) {
        return;
      }

      const operations = operationStack.get(queryKeyString);
      if (!operations) return;

      const operationIndex = operations.findIndex(
        (op) => op.timestamp === (context as { timestamp: number }).timestamp,
      );

      if (operationIndex !== -1) {
        operations.splice(operationIndex, 1);

        if (operations.length === 0) {
          operationStack.delete(queryKeyString);
        } else {
          operationStack.set(queryKeyString, operations);
        }
      }

      cleanupStaleOperations(queryKeyString);
    },
  };
}

example:

// openapi-fetch client
import { $api } from "@/api/api";
import { useOptimisticUpdate } from "@/hooks/use-optimistic-update";
// ...

const ROWS_PER_PAGE = 20;

function ReportsPanel() {
const [page, setPage] = useState(1);

const getReportsVariables = {
  query: {
    page: page,
    limit: ROWS_PER_PAGE,
  },
};

const { data: reportsData, isLoading } = $api.useQuery("get", "/reports", getReportsVariables);

// - Fully typesafe!
// - Automatic optimistic update
// - Automatic rollback in case of error
// - Support patch, post, delete
// - the internal onMutate and onError are extendable & overridable!
const updateStatusMutation = $api.useMutation(
  "patch",
  "/reports/{id}/status",
  {
    ...useOptimisticUpdate(
      ["get", "/reports", getReportsVariables],
      "/reports/{id}/status",
      {
        type: "patch",
        updateFn: (report, variables) => ({
          ...report,
          data: report.data.map((report) =>
            report.id === variables.params.path.id
              ? { ...report, status: variables.body.status }
              : report,
          ),
        }),
      },
    ),
  },
);

// ...

updateStatusMutation.mutate({
  params: { path: { id: reportId } },
  body: { status: newStatus },
});

// ...