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) {