43 lines · 1.4 KB
import { useMutation } from "@tanstack/react-query";
import { useRouter } from "@tanstack/react-router";
import { useRef, useState } from "react";
import { errorMessage } from "~/lib/error-message";

/**
 * Optimistic mutation: the expected result shows immediately and stays until every in-flight call
 * has finished and the route loaders have refreshed, so toggles like star/follow/watch never flash
 * the stale value. On failure the optimistic value is dropped, the loaders reload the real state,
 * and the error is exposed.
 */
export function useOptimisticMutation<TState, TInput>(
  serverState: TState,
  reduce: (state: TState, input: TInput) => TState,
  action: (input: TInput) => Promise<unknown>,
) {
  const router = useRouter();
  const [optimistic, setOptimistic] = useState<{ value: TState } | null>(null);
  const inFlight = useRef(0);
  const mutation = useMutation({ mutationFn: action, onSettled: () => settle() });
  const state = optimistic ? optimistic.value : serverState;

  async function settle() {
    inFlight.current -= 1;
    if (inFlight.current > 0) return;
    await router.invalidate({ sync: true });
    if (inFlight.current === 0) setOptimistic(null);
  }

  function run(input: TInput) {
    inFlight.current += 1;
    setOptimistic({ value: reduce(state, input) });
    mutation.mutate(input);
  }

  return {
    state,
    run,
    pending: mutation.isPending,
    error: mutation.isError ? errorMessage(mutation.error) : null,