umijs · chenshuai2144 · Jun 21, 2023 · Jun 16, 2023 · Jun 16, 2023 · Jun 16, 2023
diff --git a/examples/ssr-demo/src/pages/index.tsx b/examples/ssr-demo/src/pages/index.tsx
@@ -1,5 +1,10 @@
 import React from 'react';
-import { Link, useClientLoaderData, useServerLoaderData } from 'umi';
+import {
+  Link,
+  useClientLoaderData,
+  useServerInsertedHTML,
+  useServerLoaderData,
+} from 'umi';
 import Button from '../components/Button';
 // @ts-ignore
 import bigImage from './big_image.jpg';
@@ -14,6 +19,11 @@ import umiLogo from './umi.png';
 export default function HomePage() {
   const clientLoaderData = useClientLoaderData();
   const serverLoaderData = useServerLoaderData();
+
+  useServerInsertedHTML(() => {
+    return <div>inserted html</div>;
+  });
+
   return (
     <div>
       <h1 className="title">Hello~</h1>

diff --git a/packages/preset-umi/src/features/ssr/ssr.ts b/packages/preset-umi/src/features/ssr/ssr.ts
@@ -69,6 +69,30 @@ export default (api: IApi) => {
       content: `
       import * as React from 'react';
 export { React };
+`,
+    });
+
+    api.writeTmpFile({
+      noPluginDir: true,
+      path: 'core/serverInsertedHTMLContext.ts',
+      content: `
+// Use React.createContext to avoid errors from the RSC checks because
+// it can't be imported directly in Server Components:
+import React from 'react'
+
+export type ServerInsertedHTMLHook = (callbacks: () => React.ReactNode) => void;
+// More info: https://github.com/vercel/next.js/pull/40686
+export const ServerInsertedHTMLContext =
+  React.createContext<ServerInsertedHTMLHook | null>(null as any);
+
+// copy form https://github.com/vercel/next.js/blob/fa076a3a69c9ccf63c9d1e53e7b681aa6dc23db7/packages/next/src/shared/lib/server-inserted-html.tsx#L13
+export function useServerInsertedHTML(callback: () => React.ReactNode): void {
+  const addInsertedServerHTMLCallback = React.useContext(ServerInsertedHTMLContext);
+  // Should have no effects on client where there's no flush effects provider
+  if (addInsertedServerHTMLCallback) {
+    addInsertedServerHTMLCallback(callback);
+  }
+}
 `,
     });
   });

diff --git a/packages/preset-umi/src/features/tmpFiles/tmpFiles.ts b/packages/preset-umi/src/features/tmpFiles/tmpFiles.ts
@@ -617,6 +617,11 @@ if (process.env.NODE_ENV === 'development') {
           exports.push(`export { TestBrowser } from './testBrowser';`);
         }
       }
+      if (api.config.ssr && api.appData.framework === 'react') {
+        exports.push(
+          `export { useServerInsertedHTML } from './core/serverInsertedHTMLContext';`,
+        );
+      }
       // plugins
       exports.push('// plugins');
       const allPlugins = readdirSync(api.paths.absTmpPath).filter((file) =>

diff --git a/packages/preset-umi/templates/server.tpl b/packages/preset-umi/templates/server.tpl
@@ -6,11 +6,18 @@ import { PluginManager } from '{{{ umiPluginPath }}}';
 import createRequestHandler, { createMarkupGenerator } from '{{{ umiServerPath }}}';
 
 let helmetContext;
+let ServerInsertedHTMLContext;
 
 try {
   helmetContext = require('./core/helmetContext').context;
 } catch { /* means `helmet: false`, do noting */ }
 
+
+try {
+  ServerInsertedHTMLContext = require('./core/serverInsertedHTMLContext').ServerInsertedHTMLContext;
+} catch { /* means `helmet: false`, do noting */ }
+
+
 const routesWithServerLoader = {
 {{#routesWithServerLoader}}
   '{{{ id }}}': () => import('{{{ path }}}'),
@@ -50,6 +57,7 @@ const createOpts = {
   getClientRootComponent,
   helmetContext,
   createHistory,
+  ServerInsertedHTMLContext,
 };
 const requestHandler = createRequestHandler(createOpts);
 

diff --git a/packages/renderer-react/src/server.tsx b/packages/renderer-react/src/server.tsx
@@ -63,6 +63,7 @@ export async function getClientRootComponent(opts: {
 function Html({ children, loaderData, manifest }: any) {
   // TODO: 处理 head 标签，比如 favicon.ico 的一致性
   // TODO: root 支持配置
+
   return (
     <html lang="en">
       <head>
@@ -78,6 +79,7 @@ function Html({ children, loaderData, manifest }: any) {
             __html: `<b>Enable JavaScript to run this app.</b>`,
           }}
         />
+
         <div id="root">{children}</div>
         <script
           dangerouslySetInnerHTML={{

diff --git a/packages/server/src/ssr.ts b/packages/server/src/ssr.ts
@@ -1,5 +1,5 @@
-import type { ReactElement } from 'react';
-import { renderToPipeableStream } from 'react-dom/server';
+import React, { ReactElement } from 'react';
+import * as ReactDomServer from 'react-dom/server';
 import { matchRoutes } from 'react-router-dom';
 import { Writable } from 'stream';
 import type { IRoutesById } from './types';
@@ -8,6 +8,8 @@ interface RouteLoaders {
   [key: string]: () => Promise<any>;
 }
 
+export type ServerInsertedHTMLHook = (callbacks: () => React.ReactNode) => void;
+
 interface CreateRequestHandlerOptions {
   routesWithServerLoader: RouteLoaders;
   PluginManager: any;
@@ -20,8 +22,11 @@ interface CreateRequestHandlerOptions {
   getClientRootComponent: (PluginManager: any) => any;
   createHistory: (opts: any) => any;
   helmetContext?: any;
+  ServerInsertedHTMLContext: React.Context<ServerInsertedHTMLHook | null>;
 }
 
+const serverInsertedHTMLCallbacks: Set<() => React.ReactNode> = new Set();
+
 function createJSXGenerator(opts: CreateRequestHandlerOptions) {
   return async (url: string) => {
     const {
@@ -81,19 +86,61 @@ function createJSXGenerator(opts: CreateRequestHandlerOptions) {
       loaderData,
     };
 
+    const element = (await opts.getClientRootComponent(
+      context,
+    )) as ReactElement;
+
+    const JSXProvider = (props: any) => {
+      const addInsertedHtml = React.useCallback(
+        (handler: () => React.ReactNode) => {
+          serverInsertedHTMLCallbacks.add(handler);
+        },
+        [],
+      );
+
+      return React.createElement(opts.ServerInsertedHTMLContext.Provider, {
+        children: props.children,
+        value: addInsertedHtml,
+      });
+    };
+
     return {
-      element: (await opts.getClientRootComponent(context)) as ReactElement,
+      element: React.createElement(JSXProvider, { children: element }),
       manifest,
     };
   };
 }
 
+const getGenerateStaticHTML = () => {
+  return (
+    ReactDomServer.renderToString(
+      React.createElement(React.Fragment, {
+        children: Array.from(serverInsertedHTMLCallbacks).map((callback) =>
+          callback(),
+        ),
+      }),
+    ) || ''
+  );
+};
+
 export function createMarkupGenerator(opts: CreateRequestHandlerOptions) {
   const jsxGeneratorDeferrer = createJSXGenerator(opts);
+  const JSXProvider = (props: any) => {
+    const addInsertedHtml = React.useCallback(
+      (handler: () => React.ReactNode) => {
+        serverInsertedHTMLCallbacks.add(handler);
+      },
+      [],
+    );
+
+    return React.createElement(opts.ServerInsertedHTMLContext.Provider, {
+      children: props.children,
+      value: addInsertedHtml,
+    });
+  };
 
   return async (url: string) => {
     const jsx = await jsxGeneratorDeferrer(url);
-
     if (jsx) {
       return new Promise(async (resolve, reject) => {
         let chunks: Buffer[] = [];
@@ -103,9 +150,9 @@ export function createMarkupGenerator(opts: CreateRequestHandlerOptions) {
           chunks.push(Buffer.from(chunk));
           next();
         };
-        writable.on('finish', () => {
+        writable.on('finish', async () => {
           let html = Buffer.concat(chunks).toString('utf8');
-
+          html += await getGenerateStaticHTML();
           // append helmet tags to head
           if (opts.helmetContext) {
             html = html.replace(
@@ -128,12 +175,15 @@ export function createMarkupGenerator(opts: CreateRequestHandlerOptions) {
 
         // why not use `renderToStaticMarkup` or `renderToString`?
         // they will return empty root by unknown reason (maybe umi has suspense logic?)
-        const stream = renderToPipeableStream(jsx.element, {
-          onShellReady() {
-            stream.pipe(writable);
+        const stream = ReactDomServer.renderToPipeableStream(
+          React.createElement(JSXProvider, { children: jsx.element }),
+          {
+            onShellReady() {
+              stream.pipe(writable);
+            },
+            onError: reject,
           },
-          onError: reject,
-        });
+        );
       });
     }
 
@@ -161,11 +211,22 @@ export default function createRequestHandler(
 
     if (!jsx) return next();
 
-    const stream = renderToPipeableStream(jsx.element, {
+    const writable = new Writable();
+
+    writable._write = (chunk, _encoding, next) => {
+      res.write(chunk);
+      next();
+    };
+
+    writable.on('finish', async () => {
+      res.write(await getGenerateStaticHTML());
+      res.end();
+    });
+
+    const stream = await ReactDomServer.renderToPipeableStream(jsx.element, {
       bootstrapScripts: [jsx.manifest.assets['umi.js'] || '/umi.js'],
       onShellReady() {
-        res.setHeader('Content-type', 'text/html');
-        stream.pipe(res);
+        stream.pipe(writable);
       },
       onError(x: any) {
         console.error(x);