Consume the RSC stream twice in the Flight fixture

Author: sebmarkbage

fixtures/flight/server/global.js

@@ -37,6 +37,7 @@ const React = require('react');
 
 const {renderToPipeableStream} = require('react-dom/server');
 const {createFromNodeStream} = require('react-server-dom-webpack/client');
+const {PassThrough} = require('stream');
 
 const app = express();
 
@@ -146,34 +147,33 @@ app.all('/', async function (req, res, next) {
       // so we start by consuming the RSC payload. This needs a module
       // map that reverse engineers the client-side path to the SSR path.
 
-      // This is a bad hack to set the form state after SSR has started. It works
-      // because we block the root component until we have the form state and
-      // any form that reads it necessarily will come later. It also only works
-      // because the formstate type is an object which may change in the future
-      const lazyFormState = [];
-
-      let cachedResult = null;
-      async function getRootAndFormState() {
-        const {root, formState} = await createFromNodeStream(
-          rscResponse,
-          ssrManifest
-        );
-        // We shouldn't be assuming formState is an object type but at the moment
-        // we have no way of setting the form state from within the render
-        Object.assign(lazyFormState, formState);
-        return root;
-      }
+      // We need to get the formState before we start rendering but we also
+      // need to run the Flight client inside the render to get all the preloads.
+      // The API is ambivalent about what's the right one so we need two for now.
+
+      // Tee the response into two streams so that we can do both.
+      const rscResponse1 = new PassThrough();
+      const rscResponse2 = new PassThrough();
+
+      rscResponse.pipe(rscResponse1);
+      rscResponse.pipe(rscResponse2);
+
+      const {formState} = await createFromNodeStream(rscResponse1, ssrManifest);
+      rscResponse1.end();
+
+      let cachedResult;
       let Root = () => {
         if (!cachedResult) {
-          cachedResult = getRootAndFormState();
+          // Read this stream inside the render.
+          cachedResult = createFromNodeStream(rscResponse2, ssrManifest);
         }
-        return React.use(cachedResult);
+        return React.use(cachedResult).root;
       };
       // Render it into HTML by resolving the client components
       res.set('Content-type', 'text/html');
       const {pipe} = renderToPipeableStream(React.createElement(Root), {
         bootstrapScripts: mainJSChunks,
-        formState: lazyFormState,
+        formState: formState,
       });
       pipe(res);
     } catch (e) {