Update getLoadContext types for middleware context

Author: brophdawg11

packages/react-router-architect/server.ts

@@ -1,4 +1,9 @@
-import type { AppLoadContext, ServerBuild } from "react-router";
+import type {
+  AppLoadContext,
+  UNSAFE_MiddlewareEnabled as MiddlewareEnabled,
+  ServerBuild,
+  unstable_InitialContext,
+} from "react-router";
 import { createRequestHandler as createReactRouterRequestHandler } from "react-router";
 import { readableStreamToString } from "@react-router/node";
 import type {
@@ -10,6 +15,8 @@ import type {
 
 import { isBinaryType } from "./binaryTypes";
 
+type MaybePromise<T> = T | Promise<T>;
+
 /**
  * A function that returns the value to use as `context` in route `loader` and
  * `action` functions.
@@ -19,7 +26,9 @@ import { isBinaryType } from "./binaryTypes";
  */
 export type GetLoadContextFunction = (
   event: APIGatewayProxyEventV2
-) => Promise<AppLoadContext> | AppLoadContext;
+) => MiddlewareEnabled extends true
+  ? MaybePromise<unstable_InitialContext>
+  : MaybePromise<AppLoadContext>;
 
 export type RequestHandler = APIGatewayProxyHandlerV2;
 

packages/react-router-cloudflare/worker.ts

@@ -1,7 +1,14 @@
-import type { AppLoadContext, ServerBuild } from "react-router";
+import type {
+  AppLoadContext,
+  UNSAFE_MiddlewareEnabled as MiddlewareEnabled,
+  ServerBuild,
+  unstable_InitialContext,
+} from "react-router";
 import { createRequestHandler as createReactRouterRequestHandler } from "react-router";
 import { type CacheStorage } from "@cloudflare/workers-types";
 
+type MaybePromise<T> = T | Promise<T>;
+
 /**
  * A function that returns the value to use as `context` in route `loader` and
  * `action` functions.
@@ -29,7 +36,9 @@ export type GetLoadContextFunction<
       caches: CacheStorage;
     };
   };
-}) => AppLoadContext | Promise<AppLoadContext>;
+}) => MiddlewareEnabled extends true
+  ? MaybePromise<unstable_InitialContext>
+  : MaybePromise<AppLoadContext>;
 
 export type RequestHandler<Env = any> = PagesFunction<Env>;
 

packages/react-router-express/server.ts

@@ -3,13 +3,20 @@
 /// <reference lib="DOM.Iterable" />
 
 import type * as express from "express";
-import type { AppLoadContext, ServerBuild } from "react-router";
+import type {
+  AppLoadContext,
+  ServerBuild,
+  UNSAFE_MiddlewareEnabled as MiddlewareEnabled,
+  unstable_InitialContext,
+} from "react-router";
 import { createRequestHandler as createRemixRequestHandler } from "react-router";
 import {
   createReadableStreamFromReadable,
   writeReadableStreamToWritable,
 } from "@react-router/node";
 
+type MaybePromise<T> = T | Promise<T>;
+
 /**
  * A function that returns the value to use as `context` in route `loader` and
  * `action` functions.
@@ -21,7 +28,9 @@ import {
 export type GetLoadContextFunction = (
   req: express.Request,
   res: express.Response
-) => Promise<AppLoadContext> | AppLoadContext;
+) => MiddlewareEnabled extends true
+  ? MaybePromise<unstable_InitialContext>
+  : MaybePromise<AppLoadContext>;
 
 export type RequestHandler = (
   req: express.Request,

packages/react-router/index.ts

@@ -269,7 +269,10 @@ export type {
   FlashSessionData,
 } from "./lib/server-runtime/sessions";
 
-export type { Future } from "./lib/types/future.ts";
+export type {
+  Future,
+  MiddlewareEnabled as UNSAFE_MiddlewareEnabled,
+} from "./lib/types/future.ts";
 export type { unstable_SerializesTo } from "./lib/types/serializes-to.ts";
 export type { Register } from "./lib/types/register";
 export { href } from "./lib/href";

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

@@ -1,5 +1,5 @@
 import type { StaticHandler, StaticHandlerContext } from "../router/router";
-import type { ErrorResponse } from "../router/utils";
+import type { ErrorResponse, unstable_InitialContext } from "../router/utils";
 import { unstable_RouterContextProvider } from "../router/utils";
 import {
   isRouteErrorResponse,
@@ -90,11 +90,12 @@ export const createRequestHandler: CreateRequestHandlerFunction = (
   return async function requestHandler(request, _loadContext) {
     _build = typeof build === "function" ? await build() : build;
 
-    // TODO: Accept `initialContext` from `getLoadContext` when the flag is enabled
     let loadContext: AppLoadContext = _build.future.unstable_middleware
       ? // @ts-expect-error This type changes when middleware is enabled
-        (new unstable_RouterContextProvider() as AppLoadContext)
-      : _loadContext ?? {};
+        (new unstable_RouterContextProvider(
+          _loadContext as unknown as unstable_InitialContext
+        ) as AppLoadContext)
+      : _loadContext || {};
 
     if (typeof build === "function") {
       let derived = derive(_build, mode);