fix: react-server Await component (#14261)

chore: remove compression from integration test server to help streaming tests run smoother

Author: jacob-ebey

.changeset/cool-readers-attack.md

@@ -0,0 +1,5 @@
+---
+"react-router": patch
+---
+
+Add react-server Await component implementation

integration/helpers/rsc-vite/server.js

@@ -1,13 +1,11 @@
 import { parseArgs } from "node:util";
 import { createRequestListener } from "@mjackson/node-fetch-server";
-import compression from "compression";
 import express from "express";
 
 import rscRequestHandler from "./dist/rsc/index.js";
 
 const app = express();
 
-app.use(compression());
 app.use(express.static("dist/client"));
 
 app.get("/.well-known/appspecific/com.chrome.devtools.json", (req, res) => {

integration/rsc/rsc-test.ts

@@ -482,6 +482,27 @@ implementations.forEach((implementation) => {
                       path: "no-revalidate-server-action",
                       lazy: () => import("./routes/no-revalidate-server-action/home"),
                     },
+                    {
+                      id: "await-component",
+                      path: "await-component",
+                      children: [
+                        {
+                          id: "await-component.home",
+                          index: true,
+                          lazy: () => import("./routes/await-component/home"),
+                        },
+                        {
+                          id: "await-component.reject",
+                          path: "reject",
+                          lazy: () => import("./routes/await-component/reject"),
+                        },
+                        {
+                          id: "await-component.api",
+                          path: "api",
+                          lazy: () => import("./routes/await-component/api"),
+                        }
+                      ]
+                    }
                   ],
                 },
               ] satisfies RSCRouteConfig;
@@ -903,7 +924,6 @@ implementations.forEach((implementation) => {
               import { Counter } from "./home.client";
 
               export default function HomeRoute(props) {
-                console.log({props});
                 return (
                   <div>
                     <form action={redirectAction}>
@@ -1155,7 +1175,7 @@ implementations.forEach((implementation) => {
               import ClientHomeRoute from "./home.client";
 
               export function loader() {
-                console.log("loader");
+                console.log("THIS SHOULD NOT BE LOGGED!!!");
               }
 
               export default function HomeRoute() {
@@ -1184,6 +1204,90 @@ implementations.forEach((implementation) => {
                 );
               }
             `,
+
+            "src/routes/await-component/events.ts": js`
+              import EventEmitter from 'node:events'
+
+              export const events = new EventEmitter();
+            `,
+            "src/routes/await-component/api.ts": js`
+              import { events } from "./events";
+              export async function action({ request }) {
+                const event = await request.text()
+                events.emit(event);
+                return Response.json(event);
+              }
+            `,
+            "src/routes/await-component/client.tsx": js`
+              "use client";
+              import { useAsyncError, useAsyncValue } from "react-router";
+
+              export function ClientValue() {
+                const value = useAsyncValue();
+                return <div data-resolved>{value}</div>;
+              }
+
+              export function ClientError() {
+                const error = useAsyncError();
+                return <div data-rejected>{error.message}</div>;
+              }
+            `,
+            "src/routes/await-component/home.tsx": js`
+              import { Suspense } from "react";
+              import { Await } from "react-router";
+
+              import { ClientValue } from "./client";
+              import { events } from "./events";
+
+              export default function AwaitResolveTest() {
+                const promise = new Promise(resolve => {
+                  events.on("resolve", () => {
+                    resolve("Async Data");
+                  });
+                });
+
+                return (
+                  <>
+                    <Suspense fallback={<p data-fallback>Loading...</p>}>
+                      <Await resolve={promise}>
+                        <ClientValue />
+                      </Await>
+                    </Suspense>
+                    {Array.from({ length: 100 }, (_, i) => (
+                      <p key={i}>Item {i}</p>
+                    ))}
+                  </>
+                );
+              }
+            `,
+            "src/routes/await-component/reject.tsx": js`
+              import { Suspense } from "react";
+              import { Await } from "react-router";
+
+              import { ClientError } from "./client";
+              import { events } from "./events";
+
+              export default function AwaitRejectTest() {
+                const promise = new Promise((_, reject) => {
+                  events.on("reject", () => {
+                    reject(new Error("Async Error"));
+                  });
+                });
+
+                return (
+                  <>
+                    <Suspense fallback={<p data-fallback>Loading...</p>}>
+                      <Await resolve={promise} errorElement={<ClientError />}>
+                        {(data) => (<p data-resolved>{data}</p>)}
+                      </Await>
+                    </Suspense>
+                    {Array.from({ length: 100 }, (_, i) => (
+                      <p key={i}>Item {i}</p>
+                    ))}
+                  </>
+                );
+              }
+            `,
           },
         });
       });
@@ -1432,6 +1536,36 @@ implementations.forEach((implementation) => {
           await page.goto(`http://localhost:${port}/resource-error-handling/`);
           validateRSCHtml(await page.content());
         });
+
+        test("Supports Await component resolve", async ({ page }) => {
+          await page.goto(`http://localhost:${port}/await-component`, {
+            waitUntil: "commit",
+          });
+          await page.waitForSelector("[data-fallback]");
+          await fetch(`http://localhost:${port}/await-component/api`, {
+            method: "POST",
+            headers: { "Content-Type": "text/plain" },
+            body: "resolve",
+          });
+          const resolved = await page.waitForSelector("[data-resolved]");
+          expect(await resolved.innerText()).toContain("Async Data");
+        });
+
+        test("Supports Await component rejection", async ({ page }) => {
+          await page.goto(`http://localhost:${port}/await-component/reject`, {
+            waitUntil: "commit",
+          });
+          await page.waitForSelector("[data-fallback]");
+          await fetch(`http://localhost:${port}/await-component/api`, {
+            method: "POST",
+            headers: { "Content-Type": "text/plain" },
+            body: "reject",
+          });
+          const rejected = await page.waitForSelector("[data-rejected]");
+          expect(await rejected.innerText()).toContain(
+            "An error occurred in the Server Components render.",
+          );
+        });
       });
 
       test.describe("Server Actions", () => {

packages/react-router/index-react-server-client.ts

@@ -1,7 +1,7 @@
 "use client";
 
+export { AwaitContextProvider as UNSAFE_AwaitContextProvider } from "./lib/context";
 export {
-  Await,
   MemoryRouter,
   Navigate,
   Outlet,

packages/react-router/index-react-server.ts

@@ -17,11 +17,15 @@ export type {
 } from "./lib/rsc/server.rsc";
 
 // RSC implementation of agnostic APIs
-export { redirect, redirectDocument, replace } from "./lib/rsc/server.rsc";
+export {
+  Await,
+  redirect,
+  redirectDocument,
+  replace,
+} from "./lib/rsc/server.rsc";
 
 // Client references
 export {
-  Await,
   BrowserRouter,
   Form,
   HashRouter,

packages/react-router/index.ts

@@ -91,6 +91,7 @@ export type {
   RouteMatch,
   RouteObject,
 } from "./lib/context";
+export { AwaitContextProvider as UNSAFE_AwaitContextProvider } from "./lib/context";
 export type {
   AwaitProps,
   IndexRouteProps,

packages/react-router/lib/context.ts

@@ -137,6 +137,10 @@ FetchersContext.displayName = "Fetchers";
 export const AwaitContext = React.createContext<TrackedPromise | null>(null);
 AwaitContext.displayName = "Await";
 
+export const AwaitContextProvider = (
+  props: React.ComponentProps<typeof AwaitContext.Provider>,
+) => React.createElement(AwaitContext.Provider, props);
+
 export interface NavigateOptions {
   /** Replace the current entry in the history stack instead of pushing a new one */
   replace?: boolean;

packages/react-router/lib/rsc/server.rsc.ts

@@ -25,6 +25,7 @@ import {
   type Params,
   type ShouldRevalidateFunction,
   type RouterContextProvider,
+  type TrackedPromise,
   isRouteErrorResponse,
   matchRoutes,
   prependBasename,
@@ -41,6 +42,7 @@ import invariant from "../server-runtime/invariant";
 
 import {
   Outlet as UNTYPED_Outlet,
+  UNSAFE_AwaitContextProvider,
   UNSAFE_WithComponentProps,
   UNSAFE_WithHydrateFallbackProps,
   UNSAFE_WithErrorBoundaryProps,
@@ -49,6 +51,7 @@ import {
   // TSConfig, it breaks the Parcel build within this repo.
 } from "react-router/internal/react-server-client";
 import type {
+  Await as AwaitType,
   Outlet as OutletType,
   WithComponentProps as WithComponentPropsType,
   WithErrorBoundaryProps as WithErrorBoundaryPropsType,
@@ -110,6 +113,41 @@ export const replace: typeof baseReplace = (...args) => {
   return response;
 };
 
+const cachedResolvePromise: <T>(
+  resolve: T,
+) => Promise<PromiseSettledResult<Awaited<T>>> =
+  // @ts-expect-error - on 18 types, requires 19.
+  React.cache(async <T>(resolve: T) => {
+    return Promise.allSettled([resolve]).then((r) => r[0]);
+  });
+
+export const Await: typeof AwaitType = (async ({
+  children,
+  resolve,
+  errorElement,
+}: React.ComponentProps<typeof AwaitType>) => {
+  let promise = cachedResolvePromise(resolve);
+  let resolved: Awaited<typeof promise> = await promise;
+
+  if (resolved.status === "rejected" && !errorElement) {
+    throw resolved.reason;
+  }
+  if (resolved.status === "rejected") {
+    return React.createElement(UNSAFE_AwaitContextProvider, {
+      children: React.createElement(React.Fragment, null, errorElement),
+      value: { _tracked: true, _error: resolved.reason } as TrackedPromise,
+    });
+  }
+
+  const toRender =
+    typeof children === "function" ? children(resolved.value) : children;
+
+  return React.createElement(UNSAFE_AwaitContextProvider, {
+    children: toRender,
+    value: { _tracked: true, _data: resolved.value } as TrackedPromise,
+  });
+}) as any;
+
 type RSCRouteConfigEntryBase = {
   action?: ActionFunction;
   clientAction?: ClientActionFunction;