refactor: 404 handling for static sites

packages/waku/src/minimal/client.ts

@@ -58,19 +58,32 @@ const checkStatus = async (
 
 type Elements = Record<string, unknown>;
 
+// HACK I'm not super happy with this hack
+const erroredElementsPromiseMap = new WeakMap<
+  Promise<Elements>,
+  Promise<Elements>
+>();
+
 const getCached = <T>(c: () => T, m: WeakMap<object, T>, k: object): T =>
   (m.has(k) ? m : m.set(k, c())).get(k) as T;
 const cache1 = new WeakMap();
 const mergeElementsPromise = (
   a: Promise<Elements>,
   b: Promise<Elements>,
 ): Promise<Elements> => {
-  const getResult = () =>
-    Promise.all([a, b]).then(([a, b]) => {
-      const nextElements = { ...a, ...b };
-      delete nextElements._value;
-      return nextElements;
-    });
+  const getResult = () => {
+    const p = Promise.all([erroredElementsPromiseMap.get(a) || a, b])
+      .then(([a, b]) => {
+        const nextElements = { ...a, ...b };
+        delete nextElements._value;
+        return nextElements;
+      })
+      .catch((err) => {
+        erroredElementsPromiseMap.set(p, a);
+        throw err;
+      });
+    return p;
+  };
   const cache2 = getCached(() => new WeakMap(), cache1, a);
   return getCached(getResult, cache2, b);
 };
@@ -248,6 +261,27 @@ export const useRefetch = () => use(RefetchContext);
 
 const ChildrenContext = createContext<ReactNode>(undefined);
 const ChildrenContextProvider = memo(ChildrenContext.Provider);
+const ErrorContext = createContext<
+  [error: unknown, reset: () => void] | undefined
+>(undefined);
+const ErrorContextProvider = memo(ErrorContext.Provider);
+
+export const Children = () => use(ChildrenContext);
+
+export const ThrowError_UNSTABLE = () => {
+  const errAndReset = use(ErrorContext);
+  if (errAndReset) {
+    throw errAndReset[0];
+  }
+  return null;
+};
+
+export const useResetError_UNSTABLE = () => {
+  const errAndReset = use(ErrorContext);
+  if (errAndReset) {
+    return errAndReset[1];
+  }
+};
 
 export const useElement = (id: string) => {
   const elementsPromise = use(ElementsContext);
@@ -264,22 +298,22 @@ export const useElement = (id: string) => {
 const InnerSlot = ({
   id,
   children,
-  setFallback,
+  setValidElement,
   unstable_fallback,
 }: {
   id: string;
   children?: ReactNode;
-  setFallback?: (fallback: ReactNode) => void;
+  setValidElement?: (element: ReactNode) => void;
   unstable_fallback?: ReactNode;
 }) => {
   const element = useElement(id);
   const isValidElement = element !== undefined;
   useEffect(() => {
-    if (isValidElement && setFallback) {
+    if (isValidElement && setValidElement) {
       // FIXME is there `isReactNode` type checker?
-      setFallback(element as ReactNode);
+      setValidElement(element as ReactNode);
     }
-  }, [isValidElement, element, setFallback]);
+  }, [isValidElement, element, setValidElement]);
   if (!isValidElement) {
     if (unstable_fallback) {
       return unstable_fallback;
@@ -294,31 +328,32 @@ const InnerSlot = ({
   );
 };
 
-const ThrowError = ({ error }: { error: unknown }) => {
-  throw error;
-};
-
-class Fallback extends Component<
-  { children: ReactNode; fallback: ReactNode },
-  { error?: unknown }
+class GeneralErrorHandler extends Component<
+  { children?: ReactNode; errorHandler: ReactNode },
+  { error: unknown | null }
 > {
-  constructor(props: { children: ReactNode; fallback: ReactNode }) {
+  constructor(props: { children?: ReactNode; errorHandler: ReactNode }) {
     super(props);
-    this.state = {};
+    this.state = { error: null };
+    this.reset = this.reset.bind(this);
   }
   static getDerivedStateFromError(error: unknown) {
     return { error };
   }
+  reset() {
+    this.setState({ error: null });
+  }
   render() {
-    if ('error' in this.state) {
-      if (this.props.fallback) {
+    const { error } = this.state;
+    if (error !== null) {
+      if (this.props.errorHandler) {
         return createElement(
-          ChildrenContextProvider,
-          { value: createElement(ThrowError, { error: this.state.error }) },
-          this.props.fallback,
+          ErrorContextProvider,
+          { value: [error, this.reset] },
+          this.props.errorHandler,
         );
       }
-      throw this.state.error;
+      throw error;
     }
     return this.props.children;
   }
@@ -341,27 +376,36 @@ class Fallback extends Component<
 export const Slot = ({
   id,
   children,
-  unstable_fallbackToPrev,
+  unstable_handleError,
   unstable_fallback,
 }: {
   id: string;
   children?: ReactNode;
-  unstable_fallbackToPrev?: boolean;
+  unstable_handleError?: ReactNode;
   unstable_fallback?: ReactNode;
 }) => {
-  const [fallback, setFallback] = useState<ReactNode>();
-  if (unstable_fallbackToPrev) {
+  const [errorHandler, setErrorHandler] = useState<ReactNode>();
+  const setValidElement = useCallback(
+    (element: ReactNode) =>
+      setErrorHandler(
+        createElement(
+          ChildrenContextProvider,
+          { value: unstable_handleError },
+          element,
+        ),
+      ),
+    [unstable_handleError],
+  );
+  if (unstable_handleError !== undefined) {
     return createElement(
-      Fallback,
-      { fallback } as never,
-      createElement(InnerSlot, { id, setFallback }, children),
+      GeneralErrorHandler,
+      { errorHandler },
+      createElement(InnerSlot, { id, setValidElement }, children),
     );
   }
   return createElement(InnerSlot, { id, unstable_fallback }, children);
 };
 
-export const Children = () => use(ChildrenContext);
-
 /**
  * ServerRoot for SSR
  * This is not a public API.

packages/waku/src/router/client.ts

@@ -22,11 +22,12 @@ import type {
 } from 'react';
 
 import {
-  fetchRsc,
   prefetchRsc,
   Root,
   Slot,
   useRefetch,
+  ThrowError_UNSTABLE as ThrowError,
+  useResetError_UNSTABLE as useResetError,
 } from '../minimal/client.js';
 import {
   encodeRoutePath,
@@ -350,6 +351,7 @@ const NotFound = ({
   has404: boolean;
   reset: () => void;
 }) => {
+  const resetError = useResetError();
   const router = useContext(RouterContext);
   if (!router) {
     throw new Error('Missing Router');
@@ -359,13 +361,15 @@ const NotFound = ({
     if (has404) {
       const url = new URL('/404', window.location.href);
       changeRoute(parseRoute(url), { shouldScroll: true });
+      resetError?.();
       reset();
     }
-  }, [has404, reset, changeRoute]);
-  return createElement('h1', null, 'Not Found');
+  }, [has404, resetError, reset, changeRoute]);
+  return has404 ? null : createElement('h1', null, 'Not Found');
 };
 
 const Redirect = ({ to, reset }: { to: string; reset: () => void }) => {
+  const resetError = useResetError();
   const router = useContext(RouterContext);
   if (!router) {
     throw new Error('Missing Router');
@@ -388,8 +392,9 @@ const Redirect = ({ to, reset }: { to: string; reset: () => void }) => {
       url,
     );
     changeRoute(parseRoute(url), { shouldScroll: newPath });
+    resetError?.();
     reset();
-  }, [to, reset, changeRoute]);
+  }, [to, resetError, reset, changeRoute]);
   return null;
 };
 
@@ -541,7 +546,14 @@ const InnerRouter = ({
   const routeElement = createElement(Slot, { id: getRouteSlotId(route.path) });
   const rootElement = createElement(
     Slot,
-    { id: 'root', unstable_fallbackToPrev: true },
+    {
+      id: 'root',
+      unstable_handleError: createElement(
+        CustomErrorHandler,
+        { has404 },
+        createElement(ThrowError),
+      ),
+    },
     createElement(CustomErrorHandler, { has404 }, routeElement),
   );
   return createElement(
@@ -588,13 +600,6 @@ export function Router({
       ) => Promise<Record<string, unknown>>,
     ) =>
     async (responsePromise: Promise<Response>) => {
-      const has404 = (routerData[3] ||= false);
-      const response = await responsePromise;
-      if (response.status === 404 && has404) {
-        // HACK this is still an experimental logic. It's very fragile.
-        // FIXME we should cache it if 404.txt is static.
-        return fetchRsc(encodeRoutePath('/404'));
-      }
       const data = createData(responsePromise);
       Promise.resolve(data)
         .then((data) => {
@@ -654,7 +659,7 @@ export function INTERNAL_ServerRouter({ route }: { route: RouteProps }) {
   const routeElement = createElement(Slot, { id: getRouteSlotId(route.path) });
   const rootElement = createElement(
     Slot,
-    { id: 'root', unstable_fallbackToPrev: true },
+    { id: 'root', unstable_handleError: null },
     routeElement,
   );
   return createElement(