embed rsc payload inside the html page

AbanoubGhadban · AbanoubGhadban · commit 3b956375ccd8 · 2025-03-27T12:04:34.000+02:00
diff --git a/lib/react_on_rails/helper.rb b/lib/react_on_rails/helper.rb
@@ -362,6 +362,10 @@ def json_safe_and_pretty(hash_or_string)
     def rails_context(server_side: true)
       # ALERT: Keep in sync with node_package/src/types/index.ts for the properties of RailsContext
       @rails_context ||= begin
+        rsc_url = if ReactOnRails::Utils.react_on_rails_pro?
+                    ReactOnRailsPro.configuration.rsc_payload_generation_url_path
+                  end
+
         result = {
           componentRegistryTimeout: ReactOnRails.configuration.component_registry_timeout,
           railsEnv: Rails.env,
@@ -371,7 +375,8 @@ def rails_context(server_side: true)
           i18nDefaultLocale: I18n.default_locale,
           rorVersion: ReactOnRails::VERSION,
           # TODO: v13 just use the version if existing
-          rorPro: ReactOnRails::Utils.react_on_rails_pro?
+          rorPro: ReactOnRails::Utils.react_on_rails_pro?,
+          rscPayloadGenerationUrl: rsc_url
         }
         if ReactOnRails::Utils.react_on_rails_pro?
           result[:rorProVersion] = ReactOnRails::Utils.react_on_rails_pro_version
diff --git a/node_package/src/RSCClientRoot.ts b/node_package/src/RSCClientRoot.ts
@@ -1,5 +1,7 @@
 'use client';
 
+/* eslint-disable no-underscore-dangle */
+
 import * as React from 'react';
 import * as ReactDOMClient from 'react-dom/client';
 import { createFromReadableStream } from 'react-on-rails-rsc/client';
@@ -13,6 +15,12 @@ if (typeof use !== 'function') {
   throw new Error('React.use is not defined. Please ensure you are using React 19 to use server components.');
 }
 
+declare global {
+  interface Window {
+    __FLIGHT_DATA: unknown[];
+  }
+}
+
 export type RSCClientRootProps = {
   componentName: string;
   rscPayloadGenerationUrlPath: string;
@@ -35,6 +43,45 @@ const fetchRSC = ({ componentName, rscPayloadGenerationUrlPath, componentProps }
   return createFromFetch(fetch(`/${strippedUrlPath}/${componentName}?props=${propsString}`));
 };
 
+const createRSCStreamFromPage = () => {
+  let streamController: ReadableStreamController<unknown> | undefined;
+  const stream = new ReadableStream({
+    start(controller) {
+      if (typeof window === 'undefined') {
+        return;
+      }
+      const handleChunk = (chunk: unknown) => {
+        controller.enqueue(chunk);
+      };
+      if (!window.__FLIGHT_DATA) {
+        window.__FLIGHT_DATA = [];
+      }
+      window.__FLIGHT_DATA.forEach(handleChunk);
+      window.__FLIGHT_DATA.push = (...chunks: unknown[]) => {
+        chunks.forEach(handleChunk);
+        return chunks.length;
+      };
+      streamController = controller;
+    }
+  });
+
+  if (typeof document !== 'undefined' && document.readyState === 'loading') {
+    document.addEventListener('DOMContentLoaded', () => {
+      streamController?.close();
+    });
+  } else {
+    streamController?.close();
+  }
+
+  return stream;
+}
+
+const createFromRSCStream = () => {
+  const stream = createRSCStreamFromPage();
+  const transformedStream = transformRSCStreamAndReplayConsoleLogs(stream);
+  return createFromReadableStream<React.ReactNode>(transformedStream);
+}
+
 /**
  * RSCClientRoot is a React component that handles client-side rendering of React Server Components (RSC).
  * It manages the fetching, caching, and rendering of RSC payloads from the server.
@@ -53,7 +100,6 @@ const RSCClientRoot: RenderFunction = async (
   _railsContext?: RailsContext,
   domNodeId?: string,
 ) => {
-  const root = await fetchRSC({ componentName, rscPayloadGenerationUrlPath, componentProps });
   if (!domNodeId) {
     throw new Error('RSCClientRoot: No domNodeId provided');
   }
@@ -62,8 +108,10 @@ const RSCClientRoot: RenderFunction = async (
     throw new Error(`RSCClientRoot: No DOM node found for id: ${domNodeId}`);
   }
   if (domNode.innerHTML) {
+    const root = await createFromRSCStream();
     ReactDOMClient.hydrateRoot(domNode, root);
   } else {
+    const root = await fetchRSC({ componentName, rscPayloadGenerationUrlPath, componentProps })
     ReactDOMClient.createRoot(domNode).render(root);
   }
   // Added only to satisfy the return type of RenderFunction
diff --git a/node_package/src/RSCPayloadContainer.ts b/node_package/src/RSCPayloadContainer.ts
@@ -0,0 +1,100 @@
+import * as React from 'react';
+
+type StreamChunk = {
+  chunk: string;
+  isLastChunk: boolean;
+}
+
+type RSCPayloadContainerProps = {
+  RSCPayloadStream: NodeJS.ReadableStream;
+}
+
+type RSCPayloadContainerInnerProps = {
+  chunkIndex: number;
+  getChunkPromise: (chunkIndex: number) => Promise<StreamChunk>;
+}
+
+function escapeScript(script: string) {
+  return script
+    .replace(/<!--/g, '<\\!--')
+    .replace(/<\/(script)/gi, '</\\$1');
+}
+
+const RSCPayloadContainer = (
+  { chunkIndex, getChunkPromise }: RSCPayloadContainerInnerProps,
+): React.ReactNode => {
+  const chunkPromise = getChunkPromise(chunkIndex);
+  const chunk = React.use(chunkPromise);
+
+  const scriptElement = React.createElement('script', {
+    dangerouslySetInnerHTML: {
+      __html: escapeScript(`(self.__FLIGHT_DATA||=[]).push(${chunk.chunk})`),
+    },
+    key: `script-${chunkIndex}`,
+  });
+
+  if (chunk.isLastChunk) {
+    return scriptElement;
+  }
+
+  return React.createElement(React.Fragment, null, [
+    scriptElement,
+    React.createElement(
+      React.Suspense,
+      { fallback: null, key: `suspense-${chunkIndex}` },
+      React.createElement(
+        RSCPayloadContainer,
+        { chunkIndex: chunkIndex + 1, getChunkPromise },
+      ),
+    ),
+  ]);
+}
+
+export default function RSCPayloadContainerWrapper({ RSCPayloadStream }: RSCPayloadContainerProps) {
+  const [chunkPromises] = React.useState<Promise<StreamChunk>[]>(() => {
+    const promises: Promise<StreamChunk>[] = [];
+    let resolveCurrentPromise: (streamChunk: StreamChunk) => void = () => {};
+    let rejectCurrentPromise: (error: Error) => void = () => {};
+    const decoder = new TextDecoder();
+
+    const createNewPromise = () => {
+      const promise = new Promise<StreamChunk>((resolve, reject) => {
+        resolveCurrentPromise = resolve;
+        rejectCurrentPromise = reject;
+      });
+
+      promises.push(promise);
+    };
+
+    createNewPromise();
+    RSCPayloadStream.on('data', (streamChunk) => {
+      resolveCurrentPromise({ chunk: decoder.decode(streamChunk), isLastChunk: false });
+      createNewPromise();
+    });
+
+    RSCPayloadStream.on('error', (error) => {
+      rejectCurrentPromise(error);
+      createNewPromise();
+    });
+
+    RSCPayloadStream.on('end', () => {
+      resolveCurrentPromise({ chunk: '', isLastChunk: true });
+    });
+
+    return promises;
+  });
+
+  const getChunkPromise = React.useCallback((chunkIndex: number) => {
+    if (chunkIndex > chunkPromises.length) {
+      throw new Error('React on Rails Error: RSC Chunk index out of bounds');
+    }
+
+    return chunkPromises[chunkIndex];
+  }, [chunkPromises]);
+
+  return React.createElement(
+    React.Suspense,
+    { fallback: null },
+    React.createElement(RSCPayloadContainer, { chunkIndex: 0, getChunkPromise }),
+  );
+}
diff --git a/node_package/src/RSCServerRoot.ts b/node_package/src/RSCServerRoot.ts
@@ -3,6 +3,8 @@ import { createFromNodeStream } from 'react-on-rails-rsc/client.node';
 import type { RenderFunction, RailsContext } from './types';
 import transformRSCStream from './transformRSCNodeStreamAndReplayConsoleLogs';
 import loadJsonFile from './loadJsonFile';
+import RSCPayloadContainer from './RSCPayloadContainer';
+import { PassThrough } from 'stream';
 
 declare global {
   function generateRSCPayload(
@@ -82,9 +84,28 @@ const RSCServerRoot: RenderFunction = async ({ componentName, componentProps }:
     componentProps,
     railsContext.serverSideRSCPayloadParameters
   );
-  const serverComponentElement = createFromReactOnRailsNodeStream(rscPayloadStream, ssrManifest);
 
-  return () => use(serverComponentElement);
+  // Tee the stream to pass it to the server component and the payload container
+  const rscPayloadStream1 = new PassThrough();
+  rscPayloadStream.pipe(rscPayloadStream1);
+  const rscPayloadStream2 = new PassThrough();
+  rscPayloadStream.pipe(rscPayloadStream2);
+  const serverComponentElement = createFromReactOnRailsNodeStream(rscPayloadStream1, ssrManifest);
+
+  return () => React.createElement(
+    React.Fragment,
+    null, [
+      React.createElement(
+        React.Fragment,
+        { key: 'serverComponentElement' },
+        use(serverComponentElement)
+      ),
+      React.createElement(
+        RSCPayloadContainer,
+        { RSCPayloadStream: rscPayloadStream2, key: 'rscPayloadContainer' },
+      ),
+    ]
+  );
 };
 
 export default RSCServerRoot;
diff --git a/node_package/src/buildConsoleReplay.ts b/node_package/src/buildConsoleReplay.ts
@@ -51,11 +51,12 @@ export function consoleReplay(
 }
 
 export default function buildConsoleReplay(
-  customConsoleHistory: (typeof console)['history'] | undefined = undefined,
+  customConsoleHistory: typeof console['history'] | undefined = undefined,
   numberOfMessagesToSkip: number = 0,
 ): string {
-  return RenderUtils.wrapInScriptTags(
-    'consoleReplayLog',
-    consoleReplay(customConsoleHistory, numberOfMessagesToSkip),
-  );
+  const consoleReplayJS = consoleReplay(customConsoleHistory, numberOfMessagesToSkip);
+  if (consoleReplayJS.length === 0) {
+    return '';
+  }
+  return RenderUtils.wrapInScriptTags('consoleReplayLog', consoleReplayJS);
 }
diff --git a/node_package/src/streamServerRenderedReactComponent.ts b/node_package/src/streamServerRenderedReactComponent.ts
@@ -87,16 +87,41 @@ export const transformRenderStreamChunksToResultObject = (renderState: StreamRen
   const consoleHistory = console.history;
   let previouslyReplayedConsoleMessages = 0;
 
+  let consoleReplayTimeoutId: NodeJS.Timeout;
+  const buildConsoleReplayChunk = () => {
+    const consoleReplayScript = buildConsoleReplay(consoleHistory, previouslyReplayedConsoleMessages);
+    previouslyReplayedConsoleMessages = consoleHistory?.length || 0;
+    if (consoleReplayScript.length === 0) {
+      return null;
+    }
+    const consoleReplayJsonChunk = JSON.stringify(createResultObject('', consoleReplayScript, renderState));
+    return consoleReplayJsonChunk;
+  }
+
   const transformStream = new PassThrough({
     transform(chunk, _, callback) {
       // eslint-disable-next-line @typescript-eslint/no-unsafe-call,@typescript-eslint/no-unsafe-member-access
       const htmlChunk = chunk.toString() as string;
-      const consoleReplayScript = buildConsoleReplay(consoleHistory, previouslyReplayedConsoleMessages);
-      previouslyReplayedConsoleMessages = consoleHistory?.length || 0;
+      const jsonChunk = JSON.stringify(createResultObject(htmlChunk, '', renderState));
+      this.push(`${jsonChunk}\n`);
+
+      clearTimeout(consoleReplayTimeoutId);
+      consoleReplayTimeoutId = setTimeout(() => {
+        const consoleReplayChunk = buildConsoleReplayChunk();
+        if (consoleReplayChunk) {
+          this.push(`${consoleReplayChunk}\n`);
+        }
+      }, 0);
 
-      const jsonChunk = JSON.stringify(createResultObject(htmlChunk, consoleReplayScript, renderState));
+      callback();
+    },
 
-      this.push(`${jsonChunk}\n`);
+    flush(callback) {
+      clearTimeout(consoleReplayTimeoutId);
+      const consoleReplayChunk = buildConsoleReplayChunk();
+      if (consoleReplayChunk) {
+        this.push(`${consoleReplayChunk}\n`);
+      }
       callback();
     },
   });
diff --git a/node_package/src/transformRSCStreamAndReplayConsoleLogs.ts b/node_package/src/transformRSCStreamAndReplayConsoleLogs.ts
@@ -1,6 +1,6 @@
 import { RenderResult } from './types';
 
-export default function transformRSCStreamAndReplayConsoleLogs(stream: ReadableStream<Uint8Array>) {
+export default function transformRSCStreamAndReplayConsoleLogs(stream: ReadableStream<Uint8Array | RenderResult>) {
   return new ReadableStream({
     async start(controller) {
       const reader = stream.getReader();
@@ -9,35 +9,41 @@ export default function transformRSCStreamAndReplayConsoleLogs(stream: ReadableS
 
       let lastIncompleteChunk = '';
       let { value, done } = await reader.read();
-      while (!done) {
-        const decodedValue = lastIncompleteChunk + decoder.decode(value);
-        const chunks = decodedValue.split('\n');
-        lastIncompleteChunk = chunks.pop() ?? '';
+      
+      const handleJsonChunk = (chunk: RenderResult) => {
+        const { html, consoleReplayScript = '' } = chunk;
+        controller.enqueue(encoder.encode(html ?? ''));
 
-        const jsonChunks = chunks
-          .filter((line) => line.trim() !== '')
-          .map((line) => {
-            try {
-              return JSON.parse(line) as RenderResult;
-            } catch (error) {
-              console.error('Error parsing JSON:', line, error);
-              throw error;
-            }
-          });
+        const replayConsoleCode = consoleReplayScript.trim().replace(/^<script.*>/, '').replace(/<\/script>$/, '');
+        if (replayConsoleCode?.trim() !== '') {
+          const scriptElement = document.createElement('script');
+          scriptElement.textContent = replayConsoleCode;
+          document.body.appendChild(scriptElement);
+        }
+      };
+      
+      while (!done) {
+        if (ArrayBuffer.isView(value)) {
+          const decodedValue = lastIncompleteChunk + decoder.decode(value);
+          const chunks = decodedValue.split('\n');
+          lastIncompleteChunk = chunks.pop() ?? '';
 
-        for (const jsonChunk of jsonChunks) {
-          const { html, consoleReplayScript = '' } = jsonChunk;
-          controller.enqueue(encoder.encode(html ?? ''));
+          const jsonChunks = chunks
+            .filter(line => line.trim() !== '')
+            .map((line) => {
+              try {
+                return JSON.parse(line) as RenderResult;
+              } catch (error) {
+                console.error('Error parsing JSON:', line, error);
+                throw error;
+              }
+            });
 
-          const replayConsoleCode = consoleReplayScript
-            .trim()
-            .replace(/^<script.*>/, '')
-            .replace(/<\/script>$/, '');
-          if (replayConsoleCode?.trim() !== '') {
-            const scriptElement = document.createElement('script');
-            scriptElement.textContent = replayConsoleCode;
-            document.body.appendChild(scriptElement);
+          for (const jsonChunk of jsonChunks) {
+            handleJsonChunk(jsonChunk);
           }
+        } else if (value) {
+          handleJsonChunk(value);
         }
 
         // eslint-disable-next-line no-await-in-loop
diff --git a/node_package/src/types/index.ts b/node_package/src/types/index.ts
@@ -31,9 +31,9 @@ export type RailsContext = {
   pathname: string;
   search: string | null;
   httpAcceptLanguage: string;
+  rscPayloadGenerationUrl: string;
 } & ({
   serverSide: false;
-  rscPayloadGenerationUrl: string;
 } | {
   serverSide: true;
   // These parameters are passed from React on Rails Pro to the node renderer.
