fix: improve runtime context logic

Author: zllkjc

Diff for: packages/runtime/plugin-runtime/src/core/browser/index.tsx

@@ -1,7 +1,11 @@
 import cookieTool from 'cookie';
 import type React from 'react';
 import { getGlobalAppInit, getGlobalInternalRuntimeContext } from '../context';
-import { type RuntimeContext, getInitialContext } from '../context/runtime';
+import {
+  type RuntimeContext,
+  getInitialContext,
+  makeRequestContext,
+} from '../context/runtime';
 import { createLoaderManager } from '../loader/loaderManager';
 import { wrapRuntimeContextProvider } from '../react/wrapper';
 import type { SSRContainer } from '../types';
@@ -33,12 +37,12 @@ const getQuery = () =>
       return res;
     }, {});
 
-function getSSRData(): SSRContainer & ExtraSSRContainer {
+function getSSRData(): SSRContainer {
   const ssrData = window._SSR_DATA;
 
   const ssrRequest = ssrData?.context?.request;
 
-  const finalSSRData: SSRContainer & ExtraSSRContainer = {
+  const finalSSRData: SSRContainer = {
     ...(ssrData || {
       renderLevel: 0,
       mode: 'string',
@@ -117,8 +121,10 @@ export async function render(
       }),
       // garfish plugin params
       _internalRouterBaseName: App.props.basename,
-      ...{ ssrContext: ssrData.context },
+      ssrContext: ssrData.context,
     });
+    const requestContext = makeRequestContext(context);
+    context.context = requestContext;
 
     context.initialData = ssrData.data?.initialData;
     const initialData = await runBeforeRender(context);

Diff for: packages/runtime/plugin-runtime/src/core/compatible.tsx

@@ -10,6 +10,7 @@ import {
   type RuntimeContext,
   RuntimeReactContext,
   type TRuntimeContext,
+  makeRequestContext,
 } from './context/runtime';
 import { createLoaderManager } from './loader/loaderManager';
 import { type Plugin, registerPlugin } from './plugin';
@@ -146,8 +147,10 @@ export const bootstrap: BootStrap = async (
         loaderManager: createLoaderManager(initialLoadersState, {
           skipStatic: true,
         }),
-        ...(ssrData ? { ssrContext: ssrData?.context } : {}),
+        ssrContext: ssrData?.context,
       });
+      const requestContext = makeRequestContext(context);
+      context.context = requestContext;
 
       context.initialData = ssrData?.data?.initialData;
       const initialData = await runBeforeRender(context);
@@ -228,23 +231,10 @@ export const bootstrap: BootStrap = async (
 export const useRuntimeContext = () => {
   const context = useContext(RuntimeReactContext);
 
-  const baseSSRContext = context.ssrContext;
-  const tSSRContext = baseSSRContext
-    ? {
-        isBrowser: context.isBrowser,
-        request: baseSSRContext.request || ({} as TSSRContext['request']),
-        response: baseSSRContext.response || ({} as TSSRContext['response']),
-        logger: baseSSRContext.logger || ({} as TSSRContext['logger']),
-        getInitData: () => {
-          return Object.freeze(context.initialData);
-        },
-      }
-    : ({} as TSSRContext);
-
   // TODO: Here we should not provide all the RuntimeReactContext to the user
   const pickedContext: TRuntimeContext = {
     ...context,
-    context: tSSRContext,
+    context: context.context || ({} as TSSRContext),
     request: context.ssrContext?.request,
     response: context.ssrContext?.response,
   };

Diff for: packages/runtime/plugin-runtime/src/core/context/runtime.ts

@@ -16,12 +16,14 @@ interface BaseRuntimeContext {
   store?: Store;
   routeManifest: RouteManifest;
   routerContext?: StaticHandlerContext;
+  context?: TSSRContext;
   /**
    * private
    */
   unstable_getBlockNavState?: () => boolean;
 }
 
+// It's the RuntimeContext passed internally
 export interface RuntimeContext extends BaseRuntimeContext {
   [key: string]: any;
 }
@@ -30,6 +32,7 @@ export const RuntimeReactContext = createContext<RuntimeContext>({} as any);
 
 export const ServerRouterContext = createContext({} as any);
 
+// TODO: We should export this context to user as RuntimeContext, use in `init` function
 export interface TRuntimeContext extends Partial<BaseRuntimeContext> {
   initialData?: Record<string, unknown>;
   isBrowser: boolean;
@@ -53,3 +56,17 @@ export const getInitialContext = (
     routeManifest ||
     (typeof window !== 'undefined' && (window as any)[ROUTE_MANIFEST]),
 });
+
+export const makeRequestContext = (context: RuntimeContext) => {
+  const baseSSRContext = context.ssrContext;
+  const requestContext = baseSSRContext
+    ? {
+        isBrowser: context.isBrowser,
+        request: baseSSRContext.request || ({} as TSSRContext['request']),
+        response: baseSSRContext.response || ({} as TSSRContext['response']),
+        logger: baseSSRContext.logger || ({} as TSSRContext['logger']),
+      }
+    : ({} as TSSRContext);
+
+  return requestContext;
+};

Diff for: packages/runtime/plugin-runtime/src/core/server/requestHandler.ts

@@ -16,7 +16,7 @@ import {
   getGlobalInternalRuntimeContext,
   getGlobalRSCRoot,
 } from '../context';
-import { getInitialContext } from '../context/runtime';
+import { getInitialContext, makeRequestContext } from '../context/runtime';
 import { createLoaderManager } from '../loader/loaderManager';
 import { createRoot } from '../react';
 import type { SSRServerContext } from '../types';
@@ -133,6 +133,7 @@ function createSSRContext(
       params,
       headers: headersData,
       host,
+      referer: headers.get('referer')!,
       raw: request,
     },
     response: {
@@ -207,6 +208,8 @@ export const createRequestHandler: CreateRequestHandler = async (
             },
           ),
         });
+        const requestContext = makeRequestContext(context);
+        context.context = requestContext;
 
         // Handle redirects from React Router with an HTTP redirect
         const getRedirectResponse = (result: any) => {

Diff for: packages/runtime/plugin-runtime/src/core/server/stream/afterTemplate.ts

@@ -1,3 +1,4 @@
+import type { IncomingHttpHeaders } from 'http';
 import { serializeJson } from '@modern-js/runtime-utils/node';
 import type { HeadersData } from '@modern-js/runtime-utils/universal/request';
 import { type RenderLevel, SSR_DATA_JSON_ID } from '../../constants';
@@ -99,7 +100,7 @@ function createReplaceSSRData(options: {
         pathname: request.pathname,
         host: request.host,
         url: request.url,
-        headers,
+        headers: headers as IncomingHttpHeaders,
       },
     },
     mode: 'stream',

Diff for: packages/runtime/plugin-runtime/src/core/server/string/ssrData.ts

@@ -1,3 +1,4 @@
+import type { IncomingHttpHeaders } from 'http';
 import { serializeJson } from '@modern-js/runtime-utils/node';
 import type { StaticHandlerContext } from '@modern-js/runtime-utils/remix-router';
 import type { HeadersData } from '@modern-js/runtime-utils/universal/request';
@@ -68,7 +69,7 @@ export class SSRDataCollector implements Collector {
           pathname: request.pathname,
           host: request.host,
           url: request.url,
-          headers,
+          headers: headers as IncomingHttpHeaders,
         },
         reporter: {
           sessionId: reporter?.sessionId,

Diff for: packages/runtime/plugin-runtime/src/core/types.ts

@@ -1,3 +1,4 @@
+import type { IncomingHttpHeaders } from 'node:http';
 import type { OnError, OnTiming } from '@modern-js/app-tools';
 import type { BaseSSRServerContext } from '@modern-js/types';
 import type { RenderLevel } from './constants';
@@ -31,14 +32,7 @@ export interface SSRContainer {
   mode: SSRMode;
   data?: SSRData; // string ssr data
   context?: {
-    request: {
-      params: Record<string, any>;
-      query: Record<string, string>;
-      pathname: string;
-      host: string;
-      url: string;
-      headers?: Record<string, string | undefined>;
-    };
+    request: BaseSSRServerContext['request'];
     reporter?: {
       sessionId?: string;
     };
@@ -62,9 +56,6 @@ export type SSRServerContext = Pick<
 > & {
   request: BaseSSRServerContext['request'] & {
     baseUrl: string;
-    userAgent: string;
-    cookie: string;
-    cookieMap: Record<string, string>;
     raw: Request;
   };
   htmlModifiers: BuildHtmlCb[];
@@ -76,21 +67,19 @@ export type SSRServerContext = Pick<
 
 /* 通过 useRuntimeContext 获取的 SSRContext */
 interface TSSRBaseContext {
-  request: BaseSSRServerContext['request'] & {
-    userAgent: string;
-    cookie: string;
-  };
+  request: BaseSSRServerContext['request'];
   [propName: string]: any;
 }
 
-interface ServerContext extends TSSRBaseContext {
+export interface ServerContext extends TSSRBaseContext {
   isBrowser: false;
   response: BaseSSRServerContext['response'];
   logger: BaseSSRServerContext['logger'];
 }
 
-interface ClientContext extends TSSRBaseContext {
+export interface ClientContext extends TSSRBaseContext {
   isBrowser: true;
 }
 
+// TODO: rename it, maybe requestContext or renderContext
 export declare type TSSRContext = ServerContext | ClientContext;

Diff for: packages/toolkit/types/server/context.d.ts

@@ -76,6 +76,10 @@ export type BaseSSRServerContext<T extends 'node' | 'worker' = 'node'> = {
     headers: IncomingHttpHeaders;
     host: string;
     url: string;
+    referer?: string;
+    userAgent?: string;
+    cookie?: string;
+    cookieMap?: Record<string, string>;
     [propsName: string]: any;
   };
   response: BaseResponseLike;