feat: add SuspenseWrappers to Runtime (#7131)

* feat: add SuspenseWrappers to Runtime

* chore(runtime): format code

* fix(runtime): Apply suggestions from code review

Co-authored-by: Copilot <[email protected]>

* fix(runtime): remove redundant composeSuspenseWrappers function

---------

Co-authored-by: Copilot <[email protected]>

Author: riopop

.changeset/ten-paws-sort.md

@@ -0,0 +1,5 @@
+---
+'@ice/runtime': patch
+---
+
+feat: add SuspenseWrappers to Runtime

packages/runtime/src/Suspense.tsx

@@ -135,7 +135,8 @@ export function withSuspense(Component) {
   return (props: SuspenseProps) => {
     const { fallback, id, ...componentProps } = props;
 
-    const [suspenseState, updateSuspenseData] = React.useState({
+
+    const [suspenseState, updateSuspenseData] = React.useState<SuspenseState>({
       id: id,
       data: null,
       done: false,
@@ -156,24 +157,47 @@ export function withSuspense(Component) {
       updateSuspenseData(newState);
     }
 
-    return (
-      <React.Suspense fallback={fallback || null}>
+
+    // Get SuspenseWrappers from app context
+    const { SuspenseWrappers = [] } = useAppContext();
+
+    // Compose SuspenseWrappers
+    const composeSuspenseWrappers = React.useCallback(
+      (children: React.ReactNode) => {
+        if (!SuspenseWrappers.length) return children;
+
+        return SuspenseWrappers.reduce((WrappedComponent, wrapperConfig) => {
+          const { Wrapper } = wrapperConfig;
+          return <Wrapper id={id}>{WrappedComponent}</Wrapper>;
+        }, children);
+      },
+      [SuspenseWrappers, id],
+    );
+
+    const wrappedComponent = (
+      <>
         <InlineScript
           id={`suspense-parse-start-${id}`}
           script={`(${DISPATCH_SUSPENSE_EVENT_STRING})('ice-suspense-parse-start','${id}');`}
         />
-        <SuspenseContext.Provider value={suspenseState}>
-          <Component {...componentProps} />
-          <InlineScript
-            id={`suspense-parse-data-${id}`}
-            script={`(${DISPATCH_SUSPENSE_EVENT_STRING})('ice-suspense-parse-data','${id}');`}
-          />
-          <Data id={id} />
-        </SuspenseContext.Provider>
+        <Component {...componentProps} />
+        <InlineScript
+          id={`suspense-parse-data-${id}`}
+          script={`(${DISPATCH_SUSPENSE_EVENT_STRING})('ice-suspense-parse-data','${id}');`}
+        />
+        <Data id={id} />
         <InlineScript
           id={`suspense-parse-end-${id}`}
           script={`(${DISPATCH_SUSPENSE_EVENT_STRING})('ice-suspense-parse-end','${id}');`}
         />
+      </>
+    );
+
+    return (
+      <React.Suspense fallback={fallback || null}>
+        <SuspenseContext.Provider value={suspenseState}>
+          {composeSuspenseWrappers(wrappedComponent)}
+        </SuspenseContext.Provider>
       </React.Suspense>
     );
   };

packages/runtime/src/runtime.tsx

@@ -12,7 +12,9 @@ import type {
   SetAppRouter,
   AddProvider,
   AddWrapper,
+  AddSuspenseWrapper,
   RouteWrapperConfig,
+  SuspenseWrapperConfig,
   SetRender,
   AppRouterProps,
   ComponentWithChildren,
@@ -33,6 +35,8 @@ class Runtime {
 
   private RouteWrappers: RouteWrapperConfig[];
 
+  private SuspenseWrappers: SuspenseWrapperConfig[];
+
   private render: Renderer;
 
   private responseHandlers: ResponseHandler[];
@@ -46,6 +50,7 @@ class Runtime {
       return root;
     };
     this.RouteWrappers = [];
+    this.SuspenseWrappers = [];
     this.runtimeOptions = runtimeOptions;
     this.responseHandlers = [];
     this.getAppRouter = this.getAppRouter.bind(this);
@@ -55,6 +60,7 @@ class Runtime {
     return {
       ...this.appContext,
       RouteWrappers: this.RouteWrappers,
+      SuspenseWrappers: this.SuspenseWrappers,
     };
   };
 
@@ -72,6 +78,8 @@ class Runtime {
 
   public getWrappers = () => this.RouteWrappers;
 
+  public getSuspenseWrappers = () => this.SuspenseWrappers;
+
   public loadModule(module: RuntimePlugin | StaticRuntimePlugin | CommonJsRuntime) {
     let runtimeAPI: RuntimeAPI = {
       addProvider: this.addProvider,
@@ -80,6 +88,7 @@ class Runtime {
       getAppRouter: this.getAppRouter,
       setRender: this.setRender,
       addWrapper: this.addWrapper,
+      addSuspenseWrapper: this.addSuspenseWrapper,
       appContext: this.appContext,
       setAppRouter: this.setAppRouter,
       useData: process.env.ICE_CORE_ROUTER === 'true' ? useData : useSingleRouterData,
@@ -122,6 +131,12 @@ class Runtime {
     });
   };
 
+  private addSuspenseWrapper: AddSuspenseWrapper = (Wrapper) => {
+    this.SuspenseWrappers.push({
+      Wrapper,
+    });
+  };
+
   public setAppRouter: SetAppRouter = (AppRouter) => {
     this.AppRouter = AppRouter;
   };

packages/runtime/src/types.ts

@@ -117,6 +117,7 @@ export interface AppContext {
   loaderData?: LoadersData;
   routeModules?: RouteModules;
   RouteWrappers?: RouteWrapperConfig[];
+  SuspenseWrappers?: SuspenseWrapperConfig[];
   routePath?: string;
   matches?: RouteMatch[];
   routes?: RouteItem[];
@@ -187,16 +188,26 @@ export interface RouteWrapperConfig {
 
 export type AppProvider = ComponentWithChildren<any>;
 export type RouteWrapper = ComponentType<any>;
+
+export type SuspenseWrapper = ComponentWithChildren<{
+  id: string;
+}>;
+
 export type ResponseHandler = (
   req: IncomingMessage,
   res: ServerResponse,
 ) => any | Promise<any>;
 
+export interface SuspenseWrapperConfig {
+  Wrapper: SuspenseWrapper;
+}
+
 export type SetAppRouter = <T>(AppRouter: ComponentType<T>) => void;
 export type GetAppRouter = () => AppProvider;
 export type AddProvider = (Provider: AppProvider) => void;
 export type SetRender = (render: Renderer) => void;
 export type AddWrapper = (wrapper: RouteWrapper, forLayout?: boolean) => void;
+export type AddSuspenseWrapper = (wrapper: SuspenseWrapper) => void;
 export type AddResponseHandler = (handler: ResponseHandler) => void;
 export type GetResponseHandlers = () => ResponseHandler[];
 
@@ -227,6 +238,7 @@ export interface RuntimeAPI {
   getResponseHandlers: GetResponseHandlers;
   setRender: SetRender;
   addWrapper: AddWrapper;
+  addSuspenseWrapper: AddSuspenseWrapper;
   appContext: AppContext;
   useData: UseData;
   useConfig: UseConfig;