How are others retrying an async atom from an error boundary?

[sebinsua]

Hello.

How are others retrying an atom from an error boundary?

This is what I'm currently doing:

import { useCallback } from 'react';
import { atom, useStore, type WritableAtom } from 'jotai';

type RetryableAtom<Value, Args extends unknown[], Result> = WritableAtom<
  Value,
  Args,
  Result
>;

export class AtomWithRetryError extends Error {
  constructor(
    message: string,
    public atom: RetryableAtom<any, any[], any>,
    options: { cause?: Error } = {}
  ) {
    super(message, options);
  }
}
export function isAtomWithRetryError(
  error: unknown
): error is AtomWithRetryError {
  return error instanceof AtomWithRetryError;
}

// The types for this are tricky to get right and were borrowed from Jotai itself.
// See: https://github.com/pmndrs/jotai/blob/663e054d0841ae4a74c1519a7618049fd1cb4ecc/src/vanilla/utils/atomWithRefresh.ts#L24
type Read<Value, Args extends unknown[], Result> = WritableAtom<
  Value,
  Args,
  Result
>['read'];
type Write<Value, Args extends unknown[], Result> = WritableAtom<
  Value,
  Args,
  Result
>['write'];

/**
 * This function creates a retryable atom that will place its atom in
 * any errors it throws in order that they can be retried within an
 * error boundary or wherever the error is caught.
 */
export function atomWithRetry<Value, Args extends unknown[], Result>(
  read: Read<Value, Args, Result>,
  write: Write<Value, Args, Result>
): RetryableAtom<Value, Args | [], Result | void>;
export function atomWithRetry<Value>(
  read: Read<Value, [], void>
): RetryableAtom<Value, [], void>;
export function atomWithRetry<Value, Args extends unknown[], Result>(
  read: Read<Value, Args, Result>,
  write?: Write<Value, Args, Result>
) {
  const refreshAtom = atom(0)
  if (import.meta.env?.MODE !== 'production') {
    refreshAtom.debugPrivate = true
  }
  const retryAtom = atom(
    async (get, options) => {
      get(refreshAtom);
      try {
        return await read(get, options as never);
      } catch (e) {
        if (e instanceof Error) {
          throw new AtomWithRetryError(e.message, retryAtom, { cause: e });
        }
        throw e;
      }
    },
    (get, set, ...args: Args) => {
      if (args.length === 0) {
        set(refreshAtom, (c) => c + 1);
      } else if (write) {
        return write(get, set, ...args);
      }
    }
  );
  return retryAtom;
}

/**
 * This hook provides a function that can be used to retry a retryable
 * atom. Instead of interacting with an atom setter via a hook, it
 * provides a function that takes the atom at call-time, so you don't
 * need to know the atom at render-time.
 */
export function useRetryAtom() {
  const store = useStore();
  return useCallback(
    <Value>(atom: RetryableAtom<Value, unknown[], unknown>) => store.set(atom),
    [store]
  );
}

And then within my error boundaries, I write something like this:

const markAtomToRetry = useRetryAtom();

// ...

// If the error was an `AtomWithRetryError` containing a
// writable atom then we can 'set' this atom on the store
// to cause it to retry on the next read.
if (isAtomWithRetryError(error)) {
  markAtomToRetry(error.atom);
}