withSwr.ts

withSwr.ts

// Sentinel values
const REVALIDATING = Symbol.for("@@withSwr/REVALIDATING");
const READ_ERROR = Symbol.for("@@withSwr/READ_ERROR");
const INITIAL_READ_ERROR = Symbol.for("@@withSwr/INITIAL_READ_ERROR");

// Control signals
const RESET_INITIAL_READ = Symbol.for("@@withSwr/RESET_INITIAL_READ");
const FINISH_INITIAL_READ = Symbol.for("@@withSwr/FINISH_INITIAL_READ");
const REVALIDATE = Symbol.for("@@withSwr/REVALIDATE");
const RESET = RESET_INITIAL_READ;

type SwrMeta<T> =
  | { status: "error"; error: Error; value: null }
  | { status: "success"; error: null; value: T }
  | { status: "stale"; error: null; value: T };

type SwrAtom<T, Args extends any[] = [] | [symbol]> = WritableAtom<
  T,
  Args,
  void
> & {
  isStale: Atom<boolean>;
  error: Atom<Error | null>;
  meta: Atom<SwrMeta<T>>;
};

type ErrorBoundaryStrategy = "never" | "initial-only" | "always";

interface WithSwrOptions {
  useErrorBoundary?: ErrorBoundaryStrategy;
  preserveStaleIfError?: boolean;
}

function withSwr<T>(
  asyncAtom: WritableAtom<Promise<T>, any[], any>,
  options?: WithSwrOptions
): SwrAtom<T>;
function withSwr<T>(
  asyncAtom: Atom<Promise<T>>,
  options?: WithSwrOptions
): SwrAtom<T>;
function withSwr<T>(
  asyncAtom: WritableAtom<Promise<T>, any[], any> | Atom<Promise<T>>,
  {
    useErrorBoundary = "initial-only",
    preserveStaleIfError = true,
  }: WithSwrOptions = {}
): SwrAtom<T> {
  let latestValueBeforeError: T | undefined;
  const asyncNoErrorsAtom = atom(async (get) => {
    try {
      latestValueBeforeError = await get(asyncAtom);
      return latestValueBeforeError;
    } catch (err) {
      if (preserveStaleIfError) {
        return latestValueBeforeError ?? INITIAL_READ_ERROR;
      } else {
        return latestValueBeforeError ? READ_ERROR : INITIAL_READ_ERROR;
      }
    }
  });
  asyncNoErrorsAtom.debugPrivate = true;

  const errorAtom = unwrap(
    atom(async (get) => {
      try {
        await get(asyncAtom);
      } catch (err) {
        if (!(err instanceof Error)) {
          return null;
        }
        return err;
      }
      return null;
    }),
    () => null
  );

  const isInitialReadAtom = atom(true);
  isInitialReadAtom.debugPrivate = true;

  const currentValueAtom = unwrap(asyncNoErrorsAtom, () => REVALIDATING);
  currentValueAtom.debugPrivate = true;

  const latestValueAtom = unwrap(asyncNoErrorsAtom, (previous) => previous!);
  latestValueAtom.debugPrivate = true;

  const isStaleAtom = atom((get) => {
    const isInitialRead = get(isInitialReadAtom);
    const currentValue = get(currentValueAtom);

    if (isInitialRead) {
      return false;
    }

    return currentValue === REVALIDATING;
  });

  const metaAtom = atom((get): SwrMeta<T> => {
    try {
      const value = get(latestValueAtom);
      const isStale = get(isStaleAtom);
      const error = get(errorAtom);

      if (
        error instanceof Error ||
        value === INITIAL_READ_ERROR ||
        value === READ_ERROR
      ) {
        if (!error) {
          throw new Error(
            "Invariant failed: missing error in INITIAL_READ_ERROR state"
          );
        }
        return { status: "error", error, value: null };
      }

      if (isStale) {
        return { status: "stale", error: null, value };
      }

      return { status: "success", error: null, value };
    } catch (err) {
      if (err instanceof Promise) {
        return { status: "stale", error: null, value: undefined as any };
      }

      return { status: "success", error: err, value: null };
    }
  });

  // Reads from this atom are always synchronous.
  // But we use a trick to "suspend" the atom on first boot.
  const swrAtom = atom(
    (get, opts) => {
      const isInitialRead = get(isInitialReadAtom);

      // If we're still in the initial read:
      if (isInitialRead) {
        // We will throw the promise of the underlying async atom to suspend
        // without making the atom asynchronous.
        //
        // Note: the promise is unstable and doesn't have a cache but
        //       it doesn't require extra machinery because it is only
        //       used for control flow. It's safe because `.finally`
        //       flips `isInitialRead` before re-render.
        try {
          // Note: this may throw a promise from an upstream `withSwr`.
          const value = get(asyncAtom);
          if (value instanceof Promise) {
            throw value.finally(() => {
              opts.setSelf(FINISH_INITIAL_READ);
            });
          }

          return value;
        } catch (err) {
          if (err instanceof Promise) {
            throw err.finally(() => {
              opts.setSelf(FINISH_INITIAL_READ);
            });
          }

          throw err;
        }
      }

      const latestValue = get(latestValueAtom);
      const error = get(errorAtom);

      if (useErrorBoundary !== "never") {
        // Synchronously throw any errors
        if (error instanceof Error) {
          // On initial read.
          if (latestValue === INITIAL_READ_ERROR) {
            throw error;
          }
          // Or if always required.
          if (useErrorBoundary === "always") {
            throw error;
          }
        }
      } else if (
        /* useErrorBoundary === "never" && */
        latestValue === INITIAL_READ_ERROR
      ) {
        return undefined;
      }

      // If we're not in the initial read we never throw, never suspend, and
      // always return the latest value synchronously.
      return latestValue;
    },
    (get, set, ...args): void => {
      const maybeSymbol = args[0];
      if (maybeSymbol === FINISH_INITIAL_READ) {
        set(isInitialReadAtom, false);
        return;
      }

      if ("write" in asyncAtom) {
        switch (maybeSymbol) {
          // The default behaviour of refreshing is revalidation.
          default:
          case undefined: {
            asyncAtom.write(get, set, ...args);
            return;
          }
          // But you can explicitly request revalidation if you wish to.
          case REVALIDATE: {
            asyncAtom.write(get, set, ...args.slice(1));
            return;
          }
          // It's also possible to manually reset an atom in order to
          // cause a suspense state on the next "initial read".
          case RESET_INITIAL_READ: {
            set(isInitialReadAtom, true);
            asyncAtom.write(get, set, ...args.slice(1));
            return;
          }
        }
      }
    }
  ) as SwrAtom<T>;
  if ("onMount" in asyncAtom) {
    swrAtom.onMount = asyncAtom.onMount;
  }

  swrAtom.isStale = isStaleAtom;
  swrAtom.error = errorAtom;
  swrAtom.meta = metaAtom;

  return swrAtom;
}

See: https://codesandbox.io/p/sandbox/zdjsgx