refactor: update RailsContext types to include server component capabilities and streamline related functions

Author: AbanoubGhadban

eslint.config.ts

@@ -162,6 +162,8 @@ const config = tsEslint.config([
       parserOptions: {
         projectService: {
           allowDefaultProject: ['eslint.config.ts', 'knip.ts', 'node_package/tests/*.test.{ts,tsx}'],
+          // Needed because `import * as ... from` instead of `import ... from` doesn't work in this file
+          // for some imports.
           defaultProject: 'tsconfig.eslint.json',
         },
       },

node_package/src/RSCPayloadGenerator.ts

@@ -1,17 +1,21 @@
 import { PassThrough } from 'stream';
-import { RailsContext, RSCPayloadStreamInfo, RSCPayloadCallback } from './types/index.ts';
+import {
+  RailsContextWithServerComponentCapabilities,
+  RSCPayloadStreamInfo,
+  RSCPayloadCallback,
+} from './types/index.ts';
 
 declare global {
   function generateRSCPayload(
     componentName: string,
     props: unknown,
-    railsContext: RailsContext,
+    railsContext: RailsContextWithServerComponentCapabilities,
   ): Promise<NodeJS.ReadableStream>;
 }
 
-const mapRailsContextToRSCPayloadStreams = new Map<RailsContext, RSCPayloadStreamInfo[]>();
+const mapRailsContextToRSCPayloadStreams = new Map<string, RSCPayloadStreamInfo[]>();
 
-const rscPayloadCallbacks = new Map<RailsContext, Array<RSCPayloadCallback>>();
+const rscPayloadCallbacks = new Map<string, Array<RSCPayloadCallback>>();
 
 /**
  * Registers a callback to be executed when RSC payloads are generated.
@@ -27,13 +31,17 @@ const rscPayloadCallbacks = new Map<RailsContext, Array<RSCPayloadCallback>>();
  * @param railsContext - Context for the current request
  * @param callback - Function to call when an RSC payload is generated
  */
-export const onRSCPayloadGenerated = (railsContext: RailsContext, callback: RSCPayloadCallback) => {
-  const callbacks = rscPayloadCallbacks.get(railsContext) || [];
+export const onRSCPayloadGenerated = (
+  railsContext: RailsContextWithServerComponentCapabilities,
+  callback: RSCPayloadCallback,
+) => {
+  const { renderRequestId } = railsContext.componentSpecificMetadata;
+  const callbacks = rscPayloadCallbacks.get(renderRequestId) || [];
   callbacks.push(callback);
-  rscPayloadCallbacks.set(railsContext, callbacks);
+  rscPayloadCallbacks.set(renderRequestId, callbacks);
 
   // Call callback for any existing streams for this context
-  const existingStreams = mapRailsContextToRSCPayloadStreams.get(railsContext) || [];
+  const existingStreams = mapRailsContextToRSCPayloadStreams.get(renderRequestId) || [];
   existingStreams.forEach((streamInfo) => callback(streamInfo));
 };
 
@@ -56,7 +64,7 @@ export const onRSCPayloadGenerated = (railsContext: RailsContext, callback: RSCP
 export const getRSCPayloadStream = async (
   componentName: string,
   props: unknown,
-  railsContext: RailsContext,
+  railsContext: RailsContextWithServerComponentCapabilities,
 ): Promise<NodeJS.ReadableStream> => {
   if (typeof generateRSCPayload !== 'function') {
     throw new Error(
@@ -66,8 +74,9 @@ export const getRSCPayloadStream = async (
     );
   }
 
+  const { renderRequestId } = railsContext.componentSpecificMetadata;
   const stream = await generateRSCPayload(componentName, props, railsContext);
-  const streams = mapRailsContextToRSCPayloadStreams.get(railsContext) ?? [];
+  const streams = mapRailsContextToRSCPayloadStreams.get(renderRequestId) ?? [];
   const stream1 = new PassThrough();
   stream.pipe(stream1);
   const stream2 = new PassThrough();
@@ -79,25 +88,29 @@ export const getRSCPayloadStream = async (
     stream: stream2,
   };
   streams.push(streamInfo);
-  mapRailsContextToRSCPayloadStreams.set(railsContext, streams);
+  mapRailsContextToRSCPayloadStreams.set(renderRequestId, streams);
 
   // Notify callbacks about the new stream in a sync manner to maintain proper hydration timing
   // as described in the comment above onRSCPayloadGenerated
-  const callbacks = rscPayloadCallbacks.get(railsContext) || [];
+  const callbacks = rscPayloadCallbacks.get(renderRequestId) || [];
   callbacks.forEach((callback) => callback(streamInfo));
 
   return stream1;
 };
 
 export const getRSCPayloadStreams = (
-  railsContext: RailsContext,
+  railsContext: RailsContextWithServerComponentCapabilities,
 ): {
   componentName: string;
   props: unknown;
   stream: NodeJS.ReadableStream;
-}[] => mapRailsContextToRSCPayloadStreams.get(railsContext) ?? [];
+}[] => {
+  const { renderRequestId } = railsContext.componentSpecificMetadata;
+  return mapRailsContextToRSCPayloadStreams.get(renderRequestId) ?? [];
+};
 
-export const clearRSCPayloadStreams = (railsContext: RailsContext) => {
-  mapRailsContextToRSCPayloadStreams.delete(railsContext);
-  rscPayloadCallbacks.delete(railsContext);
+export const clearRSCPayloadStreams = (railsContext: RailsContextWithServerComponentCapabilities) => {
+  const { renderRequestId } = railsContext.componentSpecificMetadata;
+  mapRailsContextToRSCPayloadStreams.delete(renderRequestId);
+  rscPayloadCallbacks.delete(renderRequestId);
 };

node_package/src/ReactOnRailsRSC.ts

@@ -4,7 +4,7 @@ import { PassThrough, Readable } from 'stream';
 
 import {
   RSCRenderParams,
-  RailsContextWithComponentSpecificMetadata,
+  assertRailsContextWithServerComponentCapabilities,
   StreamRenderState,
   StreamableComponentResult,
 } from './types/index.ts';
@@ -34,11 +34,10 @@ const streamRenderRSCComponent = (
   options: RSCRenderParams,
 ): Readable => {
   const { throwJsErrors } = options;
-  if (!options.railsContext?.serverSide || !options.railsContext.reactClientManifestFileName) {
-    throw new Error('Rails context is not available');
-  }
+  const { railsContext } = options;
+  assertRailsContextWithServerComponentCapabilities(railsContext);
 
-  const { reactClientManifestFileName } = options.railsContext;
+  const { reactClientManifestFileName } = railsContext;
   const renderState: StreamRenderState = {
     result: null,
     hasErrors: false,
@@ -78,9 +77,7 @@ const streamRenderRSCComponent = (
     });
 
   readableStream.on('end', () => {
-    if (options.railsContext?.componentSpecificMetadata) {
-      notifySSREnd(options.railsContext as RailsContextWithComponentSpecificMetadata);
-    }
+    notifySSREnd(railsContext);
   });
   return readableStream;
 };

node_package/src/getReactServerComponent.server.ts

@@ -2,7 +2,7 @@ import { BundleManifest } from 'react-on-rails-rsc';
 import { buildClientRenderer } from 'react-on-rails-rsc/client.node';
 import transformRSCStream from './transformRSCNodeStream.ts';
 import loadJsonFile from './loadJsonFile.ts';
-import { RailsContext } from './types/index.ts';
+import { assertRailsContextWithServerComponentCapabilities, RailsContext } from './types/index.ts';
 
 type RSCServerRootProps = {
   componentName: string;
@@ -58,18 +58,7 @@ const getReactServerComponent = async ({
   componentProps,
   railsContext,
 }: RSCServerRootProps) => {
-  if (
-    !railsContext?.serverSide ||
-    !railsContext?.reactClientManifestFileName ||
-    !railsContext?.reactServerClientManifestFileName
-  ) {
-    throw new Error(
-      'serverClientManifestFileName and reactServerClientManifestFileName are required. ' +
-        'Please ensure that React Server Component webpack configurations are properly set ' +
-        'as stated in the React Server Component tutorial. ' +
-        'Make sure to use "stream_react_component" instead of "react_component" to SSR a server component.',
-    );
-  }
+  assertRailsContextWithServerComponentCapabilities(railsContext);
 
   if (typeof ReactOnRails.getRSCPayloadStream !== 'function') {
     throw new Error(

node_package/src/injectRSCPayload.ts

@@ -1,6 +1,6 @@
 import { PipeableStream } from 'react-dom/server';
 import { PassThrough, Transform } from 'stream';
-import { RailsContext } from './types/index.ts';
+import { RailsContextWithServerComponentCapabilities } from './types/index.ts';
 
 // In JavaScript, when an escape sequence with a backslash (\) is followed by a character
 // that isn't a recognized escape character, the backslash is ignored, and the character
@@ -50,7 +50,7 @@ function writeChunk(chunk: string, transform: Transform, cacheKey: string) {
  */
 export default function injectRSCPayload(
   pipeableHtmlStream: NodeJS.ReadableStream | PipeableStream,
-  railsContext: RailsContext,
+  railsContext: RailsContextWithServerComponentCapabilities,
 ) {
   const htmlStream = new PassThrough();
   pipeableHtmlStream.pipe(htmlStream);

node_package/src/postSSRHooks.ts

@@ -1,18 +1,18 @@
-import { RailsContextWithComponentSpecificMetadata } from './types/index.ts';
+import { RailsContextWithServerComponentCapabilities } from './types/index.ts';
 
 type PostSSRHook = () => void;
 const postSSRHooks = new Map<string, PostSSRHook[]>();
 
 export const addPostSSRHook = (
-  railsContext: RailsContextWithComponentSpecificMetadata,
+  railsContext: RailsContextWithServerComponentCapabilities,
   hook: PostSSRHook,
 ) => {
   const hooks = postSSRHooks.get(railsContext.componentSpecificMetadata.renderRequestId) || [];
   hooks.push(hook);
   postSSRHooks.set(railsContext.componentSpecificMetadata.renderRequestId, hooks);
 };
 
-export const notifySSREnd = (railsContext: RailsContextWithComponentSpecificMetadata) => {
+export const notifySSREnd = (railsContext: RailsContextWithServerComponentCapabilities) => {
   const hooks = postSSRHooks.get(railsContext.componentSpecificMetadata.renderRequestId);
   if (hooks) {
     hooks.forEach((hook) => hook());

node_package/src/streamServerRenderedReactComponent.ts

@@ -8,8 +8,8 @@ import buildConsoleReplay from './buildConsoleReplay.ts';
 import handleError from './handleError.ts';
 import { renderToPipeableStream, PipeableStream } from './ReactDOMServer.cts';
 import { createResultObject, convertToError, validateComponent } from './serverRenderUtils.ts';
-import type {
-  RailsContextWithComponentSpecificMetadata,
+import {
+  assertRailsContextWithServerComponentCapabilities,
   RenderParams,
   StreamRenderState,
   StreamableComponentResult,
@@ -138,7 +138,7 @@ const streamRenderReactComponent = (
   reactRenderingResult: StreamableComponentResult,
   options: RenderParams,
 ) => {
-  const { name: componentName, throwJsErrors, domNodeId } = options;
+  const { name: componentName, throwJsErrors, domNodeId, railsContext } = options;
   const renderState: StreamRenderState = {
     result: null,
     hasErrors: false,
@@ -163,10 +163,7 @@ const streamRenderReactComponent = (
     endStream();
   };
 
-  const { railsContext } = options;
-  if (!railsContext) {
-    throw new Error('railsContext is required to stream a React component');
-  }
+  assertRailsContextWithServerComponentCapabilities(railsContext);
 
   Promise.resolve(reactRenderingResult)
     .then((reactRenderedElement) => {
@@ -195,7 +192,7 @@ const streamRenderReactComponent = (
         },
         onAllReady() {
           if (railsContext.componentSpecificMetadata?.renderRequestId) {
-            notifySSREnd(railsContext as RailsContextWithComponentSpecificMetadata);
+            notifySSREnd(railsContext);
           }
         },
         identifierPrefix: domNodeId,

node_package/src/types/index.ts

@@ -62,12 +62,38 @@ export type RailsContext = {
     }
 );
 
-export type RailsContextWithComponentSpecificMetadata = RailsContext & {
+export type RailsContextWithServerComponentCapabilities = RailsContext & {
+  serverSide: true;
+  serverSideRSCPayloadParameters: unknown;
+  reactClientManifestFileName: string;
+  reactServerClientManifestFileName: string;
   componentSpecificMetadata: {
     renderRequestId: string;
   };
 };
 
+export const assertRailsContextWithServerComponentCapabilities: (
+  context: RailsContext | undefined,
+) => asserts context is RailsContextWithServerComponentCapabilities = (
+  context: RailsContext | undefined,
+): asserts context is RailsContextWithServerComponentCapabilities => {
+  if (
+    !context ||
+    !('serverSideRSCPayloadParameters' in context) ||
+    !('reactClientManifestFileName' in context) ||
+    !('reactServerClientManifestFileName' in context) ||
+    !('componentSpecificMetadata' in context)
+  ) {
+    throw new Error(
+      'Rails context does not have server side RSC payload parameters.\n\n' +
+        'Please ensure:\n' +
+        '1. You are using a compatible version of react_on_rails_pro\n' +
+        '2. Server components support is enabled by setting:\n' +
+        '   ReactOnRailsPro.configuration.enable_rsc_support = true',
+    );
+  }
+};
+
 // not strictly what we want, see https://github.com/microsoft/TypeScript/issues/17867#issuecomment-323164375
 type AuthenticityHeaders = Record<string, string> & {
   'X-CSRF-Token': string | null;
@@ -183,7 +209,8 @@ export interface RenderParams extends Params {
   renderingReturnsPromises: boolean;
 }
 
-export interface RSCRenderParams extends RenderParams {
+export interface RSCRenderParams extends Omit<RenderParams, 'railsContext'> {
+  railsContext: RailsContextWithServerComponentCapabilities;
   reactClientManifestFileName: string;
 }
 
@@ -301,7 +328,7 @@ export interface ReactOnRails {
    * Adds a post SSR hook to be called after the SSR has completed.
    * @param hook - The hook to be called after the SSR has completed.
    */
-  addPostSSRHook(railsContext: RailsContextWithComponentSpecificMetadata, hook: () => void): void;
+  addPostSSRHook(railsContext: RailsContextWithServerComponentCapabilities, hook: () => void): void;
 }
 
 export type RSCPayloadStreamInfo = {
@@ -432,7 +459,7 @@ export interface ReactOnRailsInternal extends ReactOnRails {
   getRSCPayloadStream?: (
     componentName: string,
     props: unknown,
-    railsContext: RailsContext,
+    railsContext: RailsContextWithServerComponentCapabilities,
   ) => Promise<NodeJS.ReadableStream>;
 
   /**
@@ -441,7 +468,7 @@ export interface ReactOnRailsInternal extends ReactOnRails {
     * @param railsContext - The Rails context of the current rendering request.
     * @returns An array of objects, each containing the component name and its corresponding NodeJS.ReadableStream.
     */
-  getRSCPayloadStreams?: (railsContext: RailsContext) => {
+  getRSCPayloadStreams?: (railsContext: RailsContextWithServerComponentCapabilities) => {
     componentName: string;
     props: unknown;
     stream: NodeJS.ReadableStream;
@@ -452,13 +479,16 @@ export interface ReactOnRailsInternal extends ReactOnRails {
    * @param railsContext - The Rails context of the current rendering request.
    * @param callback - The callback to be called when an RSC payload stream is generated.
    */
-  onRSCPayloadGenerated?: (railsContext: RailsContext, callback: RSCPayloadCallback) => void;
+  onRSCPayloadGenerated?: (
+    railsContext: RailsContextWithServerComponentCapabilities,
+    callback: RSCPayloadCallback,
+  ) => void;
 
   /**
    * Clears all RSC payload streams generated for the rendering request of the given Rails context.
    * @param railsContext - The Rails context of the current rendering request.
    */
-  clearRSCPayloadStreams?: (railsContext: RailsContext) => void;
+  clearRSCPayloadStreams?: (railsContext: RailsContextWithServerComponentCapabilities) => void;
 }
 
 export type RenderStateHtml = FinalHtmlResult | Promise<FinalHtmlResult>;