more type acrobatics

Author: chargome

packages/react-router/src/server/createSentryHandleRequest.tsx

@@ -53,6 +53,22 @@ export interface SentryHandleRequestOptions {
   botRegex?: RegExp;
 }
 
+type HandleRequestWithoutMiddleware = (
+  request: Request,
+  responseStatusCode: number,
+  responseHeaders: Headers,
+  routerContext: EntryContext,
+  loadContext: AppLoadContext,
+) => Promise<unknown>;
+
+type HandleRequestWithMiddleware = (
+  request: Request,
+  responseStatusCode: number,
+  responseHeaders: Headers,
+  routerContext: EntryContext,
+  loadContext: RouterContextProvider,
+) => Promise<unknown>;
+
 /**
  * A complete Sentry-instrumented handleRequest implementation that handles both
  * route parametrization and trace meta tag injection.
@@ -62,13 +78,7 @@ export interface SentryHandleRequestOptions {
  */
 export function createSentryHandleRequest(
   options: SentryHandleRequestOptions,
-): (
-  request: Request,
-  responseStatusCode: number,
-  responseHeaders: Headers,
-  routerContext: EntryContext,
-  loadContext: AppLoadContext | RouterContextProvider,
-) => Promise<unknown> {
+): HandleRequestWithoutMiddleware & HandleRequestWithMiddleware {
   const {
     streamTimeout = 10000,
     renderToPipeableStream,
@@ -135,5 +145,6 @@ export function createSentryHandleRequest(
   };
 
   // Wrap the handle request function for request parametrization
-  return wrapSentryHandleRequest(handleRequest);
+  return wrapSentryHandleRequest(handleRequest as HandleRequestWithoutMiddleware) as HandleRequestWithoutMiddleware &
+    HandleRequestWithMiddleware;
 }

packages/react-router/src/server/wrapSentryHandleRequest.ts

@@ -26,15 +26,33 @@ type OriginalHandleRequestWithMiddleware = (
   loadContext: RouterContextProvider,
 ) => Promise<unknown>;
 
-type OriginalHandleRequest = OriginalHandleRequestWithoutMiddleware | OriginalHandleRequestWithMiddleware;
-
 /**
  * Wraps the original handleRequest function to add Sentry instrumentation.
  *
  * @param originalHandle - The original handleRequest function to wrap
  * @returns A wrapped version of the handle request function with Sentry instrumentation
  */
-export function wrapSentryHandleRequest(originalHandle: OriginalHandleRequest): OriginalHandleRequest {
+export function wrapSentryHandleRequest(
+  originalHandle: OriginalHandleRequestWithoutMiddleware,
+): OriginalHandleRequestWithoutMiddleware;
+/**
+ * Wraps the original handleRequest function to add Sentry instrumentation.
+ *
+ * @param originalHandle - The original handleRequest function to wrap
+ * @returns A wrapped version of the handle request function with Sentry instrumentation
+ */
+export function wrapSentryHandleRequest(
+  originalHandle: OriginalHandleRequestWithMiddleware,
+): OriginalHandleRequestWithMiddleware;
+/**
+ * Wraps the original handleRequest function to add Sentry instrumentation.
+ *
+ * @param originalHandle - The original handleRequest function to wrap
+ * @returns A wrapped version of the handle request function with Sentry instrumentation
+ */
+export function wrapSentryHandleRequest(
+  originalHandle: OriginalHandleRequestWithoutMiddleware | OriginalHandleRequestWithMiddleware,
+): OriginalHandleRequestWithoutMiddleware | OriginalHandleRequestWithMiddleware {
   return async function sentryInstrumentedHandleRequest(
     request: Request,
     responseStatusCode: number,
@@ -67,10 +85,39 @@ export function wrapSentryHandleRequest(originalHandle: OriginalHandleRequest):
     }
 
     try {
-      return await originalHandle(request, responseStatusCode, responseHeaders, routerContext, loadContext);
+      // Type guard to call the correct overload based on loadContext type
+      if (isRouterContextProvider(loadContext)) {
+        // loadContext is RouterContextProvider
+        return await (originalHandle as OriginalHandleRequestWithMiddleware)(
+          request,
+          responseStatusCode,
+          responseHeaders,
+          routerContext,
+          loadContext,
+        );
+      } else {
+        // loadContext is AppLoadContext
+        return await (originalHandle as OriginalHandleRequestWithoutMiddleware)(
+          request,
+          responseStatusCode,
+          responseHeaders,
+          routerContext,
+          loadContext,
+        );
+      }
     } finally {
       await flushIfServerless();
     }
+
+    /**
+     * Helper type guard to determine if the context is a RouterContextProvider.
+     *
+     * @param ctx - The context to check
+     * @returns True if the context is a RouterContextProvider
+     */
+    function isRouterContextProvider(ctx: AppLoadContext | RouterContextProvider): ctx is RouterContextProvider {
+      return typeof (ctx as RouterContextProvider)?.get === 'function';
+    }
   };
 }