Hook to execute a server action and handle its state

action.ts

export async function action() {
  await new Promise((resolve) =>
    setTimeout(resolve, 1000 + Math.random() * 1000),
  )

  if (Math.random() > 0.5) {
    return {
      error: 'There was an error fetching the data',
    }
  }

  return {
    data: Math.round(Math.random() * 10000),
  }
}

page.tsx

'use client'

import { useServerAction } from './use-server-action'
import { action } from './action'

export default function UseServerActionPage() {
  const { idle, data, error, pending, execute, reset } = useServerAction(action)

  useEffect(() => {
    execute().then(({ data, error }) => {
      console.log('Server action executed:', { data, error })
    })
  }, [execute])

  if (idle) {
    return (
      <div>Idle</div>
    )
  }

  if (pending) {
    return (
      <div>Pending</div>
    )
  }

  if (error) {
    return (
      <div>Error</div>
    )
  }

  if (data) {
    return (
      <div>Data</div>
    )
  }
}

use-server-action.ts

'use client'

import { useCallback, useRef, useState } from 'react'

type ServerAction = (
  // biome-ignore lint/suspicious/noExplicitAny: it is required to allow any type here
  ...args: any[]
  // biome-ignore lint/suspicious/noExplicitAny: it is required to allow any type here
) => Promise<ServerActionResult<any, any> | undefined>

type ServerActionResult<Data, Error> =
  | { data: Data; error?: never }
  | { data?: never; error: Error }

type InferData<T> = T extends { data: infer D } ? D : never
type InferError<T> = T extends { error: infer E } ? E : never

/**
 * @description Hook to execute a server action and handle its state
 *
 * @param fn Function that executes the server action and returns a promise with the result
 *
 * @returns An object with the following properties:
 *
 * - `idle`: A boolean indicating if the action is idle
 * - `pending`: A boolean indicating if the action is pending
 * - `data`: The data returned by the action, if any.
 * - `error`: The error thrown by the action, if any
 * - `execute`: A function to execute the action that returns a promise with the result
 * - `reset`: A function to reset the internal state
 *
 * @example
 * ```tsx
 * const { execute, reset, pending, data, error } = useServerAction(someServerAction)
 *
 * return (
 *   <div>
 *     {pending && <p>Loading...</p>}
 *     {data && <p>Data: {data}</p>}
 *     {error && <p>Error: {error}</p>}
 *     <button onClick={async () => { await execute({ some: 'data' }) }}>Submit</button>
 *     <button onClick={reset}>Reset</button>
 *   </div>
 * )
 * ```
 */
export function useServerAction<Fn extends ServerAction>(fn: Fn) {
  type FnArgs = Parameters<Fn>
  type Result = Awaited<ReturnType<Fn>>
  type Data = InferData<Result>
  type Error = InferError<Result>

  const [state, setState] = useState<
    | { status: 'idle' }
    | { status: 'pending' }
    | { status: 'data'; data: Data }
    | { status: 'error'; error: Error }
  >({ status: 'idle' })

  const requestIdRef = useRef<string | null>(null)

  const execute = useCallback(
    async (...args: FnArgs): Promise<ServerActionResult<Data, Error>> => {
      const requestId = Math.random().toString(36).slice(2)
      requestIdRef.current = requestId

      setState({ status: 'pending' })

      const res = (await fn(...args)) ?? { data: undefined, error: undefined }

      if (requestIdRef.current !== requestId) {
        return {} as ServerActionResult<Data, Error>
      }

      if ('error' in res) {
        setState({ status: 'error', error: res.error })
        return { error: res.error }
      }

      setState({ status: 'data', data: res?.data })
      return { data: res.data }
    },
    [fn],
  )

  const reset = useCallback(() => {
    requestIdRef.current = null
    setState({ status: 'idle' })
  }, [])

  return {
    idle: state.status === 'idle',
    data: state.status === 'data' ? state.data : undefined,
    error: state.status === 'error' ? state.error : undefined,
    pending: state.status === 'pending',
    execute,
    reset,
  }
}