[Flight] Add Debug Channel option for stateful connection to the backend in DEV

fixtures/flight/server/global.js

@@ -104,6 +104,9 @@ async function renderApp(req, res, next) {
   if (req.headers['cache-control']) {
     proxiedHeaders['Cache-Control'] = req.get('cache-control');
   }
+  if (req.get('rsc-request-id')) {
+    proxiedHeaders['rsc-request-id'] = req.get('rsc-request-id');
+  }
 
   const requestsPrerender = req.path === '/prerender';
 

fixtures/flight/server/region.js

@@ -50,7 +50,27 @@ const {readFile} = require('fs').promises;
 
 const React = require('react');
 
-async function renderApp(res, returnValue, formState, noCache) {
+const activeDebugChannels =
+  process.env.NODE_ENV === 'development' ? new Map() : null;
+
+function getDebugChannel(req) {
+  if (process.env.NODE_ENV !== 'development') {
+    return undefined;
+  }
+  const requestId = req.get('rsc-request-id');
+  if (!requestId) {
+    return undefined;
+  }
+  return activeDebugChannels.get(requestId);
+}
+
+async function renderApp(
+  res,
+  returnValue,
+  formState,
+  noCache,
+  promiseForDebugChannel
+) {
   const {renderToPipeableStream} = await import(
     'react-server-dom-webpack/server'
   );
@@ -101,7 +121,9 @@ async function renderApp(res, returnValue, formState, noCache) {
   );
   // For client-invoked server actions we refresh the tree and return a return value.
   const payload = {root, returnValue, formState};
-  const {pipe} = renderToPipeableStream(payload, moduleMap);
+  const {pipe} = renderToPipeableStream(payload, moduleMap, {
+    debugChannel: await promiseForDebugChannel,
+  });
   pipe(res);
 }
 
@@ -166,7 +188,7 @@ app.get('/', async function (req, res) {
   if ('prerender' in req.query) {
     await prerenderApp(res, null, null, noCache);
   } else {
-    await renderApp(res, null, null, noCache);
+    await renderApp(res, null, null, noCache, getDebugChannel(req));
   }
 });
 
@@ -204,7 +226,7 @@ app.post('/', bodyParser.text(), async function (req, res) {
       // We handle the error on the client
     }
     // Refresh the client and return the value
-    renderApp(res, result, null, noCache);
+    renderApp(res, result, null, noCache, getDebugChannel(req));
   } else {
     // This is the progressive enhancement case
     const UndiciRequest = require('undici').Request;
@@ -220,11 +242,11 @@ app.post('/', bodyParser.text(), async function (req, res) {
       // Wait for any mutations
       const result = await action();
       const formState = decodeFormState(result, formData);
-      renderApp(res, null, formState, noCache);
+      renderApp(res, null, formState, noCache, undefined);
     } catch (x) {
       const {setServerState} = await import('../src/ServerState.js');
       setServerState('Error: ' + x.message);
-      renderApp(res, null, null, noCache);
+      renderApp(res, null, null, noCache, undefined);
     }
   }
 });
@@ -324,7 +346,7 @@ if (process.env.NODE_ENV === 'development') {
   });
 }
 
-app.listen(3001, () => {
+const httpServer = app.listen(3001, () => {
   console.log('Regional Flight Server listening on port 3001...');
 });
 
@@ -346,3 +368,27 @@ app.on('error', function (error) {
       throw error;
   }
 });
+
+if (process.env.NODE_ENV === 'development') {
+  // Open a websocket server for Debug information
+  const WebSocket = require('ws');
+  const webSocketServer = new WebSocket.Server({noServer: true});
+
+  httpServer.on('upgrade', (request, socket, head) => {
+    const DEBUG_CHANNEL_PATH = '/debug-channel?';
+    if (request.url.startsWith(DEBUG_CHANNEL_PATH)) {
+      const requestId = request.url.slice(DEBUG_CHANNEL_PATH.length);
+      const promiseForWs = new Promise(resolve => {
+        webSocketServer.handleUpgrade(request, socket, head, ws => {
+          ws.on('close', () => {
+            activeDebugChannels.delete(requestId);
+          });
+          resolve(ws);
+        });
+      });
+      activeDebugChannels.set(requestId, promiseForWs);
+    } else {
+      socket.destroy();
+    }
+  });
+}

fixtures/flight/src/App.js

@@ -123,7 +123,6 @@ async function ServerComponent({noCache}) {
 export default async function App({prerender, noCache}) {
   const res = await fetch('http://localhost:3001/todos');
   const todos = await res.json();
-  console.log(res);
 
   const dedupedChild = <ServerComponent noCache={noCache} />;
   const message = getServerState();

fixtures/flight/src/index.js

@@ -42,17 +42,43 @@ function Shell({data}) {
 }
 
 async function hydrateApp() {
-  const {root, returnValue, formState} = await createFromFetch(
-    fetch('/', {
-      headers: {
-        Accept: 'text/x-component',
-      },
-    }),
-    {
-      callServer,
-      findSourceMapURL,
-    }
-  );
+  let response;
+  if (
+    process.env.NODE_ENV === 'development' &&
+    typeof WebSocketStream === 'function'
+  ) {
+    const requestId = crypto.randomUUID();
+    const wss = new WebSocketStream(
+      'ws://localhost:3001/debug-channel?' + requestId
+    );
+    const debugChannel = await wss.opened;
+    response = createFromFetch(
+      fetch('/', {
+        headers: {
+          Accept: 'text/x-component',
+          'rsc-request-id': requestId,
+        },
+      }),
+      {
+        callServer,
+        debugChannel,
+        findSourceMapURL,
+      }
+    );
+  } else {
+    response = createFromFetch(
+      fetch('/', {
+        headers: {
+          Accept: 'text/x-component',
+        },
+      }),
+      {
+        callServer,
+        findSourceMapURL,
+      }
+    );
+  }
+  const {root, returnValue, formState} = await response;
 
   ReactDOM.hydrateRoot(
     document,

packages/react-client/src/ReactFlightClient.js

@@ -328,6 +328,8 @@ export type FindSourceMapURLCallback = (
   environmentName: string,
 ) => null | string;
 
+export type DebugChannelCallback = (message: string) => void;
+
 export type Response = {
   _bundlerConfig: ServerConsumerModuleMap,
   _serverReferenceConfig: null | ServerManifest,
@@ -351,6 +353,7 @@ export type Response = {
   _debugRootStack?: null | Error, // DEV-only
   _debugRootTask?: null | ConsoleTask, // DEV-only
   _debugFindSourceMapURL?: void | FindSourceMapURLCallback, // DEV-only
+  _debugChannel?: void | DebugChannelCallback, // DEV-only
   _replayConsole: boolean, // DEV-only
   _rootEnvironmentName: string, // DEV-only, the requested environment name.
 };
@@ -687,6 +690,15 @@ export function reportGlobalError(response: Response, error: Error): void {
       triggerErrorOnChunk(chunk, error);
     }
   });
+  if (__DEV__) {
+    const debugChannel = response._debugChannel;
+    if (debugChannel !== undefined) {
+      // If we don't have any more ways of reading data, we don't have to send any
+      // more neither. So we close the writable side.
+      debugChannel('');
+      response._debugChannel = undefined;
+    }
+  }
   if (enableProfilerTimer && enableComponentPerformanceTrack) {
     markAllTracksInOrder();
     flushComponentPerformance(
@@ -1667,6 +1679,14 @@ function parseModelString(
       }
       case 'Y': {
         if (__DEV__) {
+          if (value.length > 2) {
+            const debugChannel = response._debugChannel;
+            if (debugChannel) {
+              const ref = value.slice(2);
+              debugChannel('R:' + ref); // Release this reference immediately
+            }
+          }
+
           // In DEV mode we encode omitted objects in logs as a getter that throws
           // so that when you try to access it on the client, you know why that
           // happened.
@@ -1730,9 +1750,10 @@ function ResponseInstance(
   encodeFormAction: void | EncodeFormActionCallback,
   nonce: void | string,
   temporaryReferences: void | TemporaryReferenceSet,
-  findSourceMapURL: void | FindSourceMapURLCallback,
-  replayConsole: boolean,
-  environmentName: void | string,
+  findSourceMapURL: void | FindSourceMapURLCallback, // DEV-only
+  replayConsole: boolean, // DEV-only
+  environmentName: void | string, // DEV-only
+  debugChannel: void | DebugChannelCallback, // DEV-only
 ) {
   const chunks: Map<number, SomeChunk<any>> = new Map();
   this._bundlerConfig = bundlerConfig;
@@ -1787,6 +1808,7 @@ function ResponseInstance(
       );
     }
     this._debugFindSourceMapURL = findSourceMapURL;
+    this._debugChannel = debugChannel;
     this._replayConsole = replayConsole;
     this._rootEnvironmentName = rootEnv;
   }
@@ -1802,9 +1824,10 @@ export function createResponse(
   encodeFormAction: void | EncodeFormActionCallback,
   nonce: void | string,
   temporaryReferences: void | TemporaryReferenceSet,
-  findSourceMapURL: void | FindSourceMapURLCallback,
-  replayConsole: boolean,
-  environmentName: void | string,
+  findSourceMapURL: void | FindSourceMapURLCallback, // DEV-only
+  replayConsole: boolean, // DEV-only
+  environmentName: void | string, // DEV-only
+  debugChannel: void | DebugChannelCallback, // DEV-only
 ): Response {
   // $FlowFixMe[invalid-constructor]: the shapes are exact here but Flow doesn't like constructors
   return new ResponseInstance(
@@ -1818,6 +1841,7 @@ export function createResponse(
     findSourceMapURL,
     replayConsole,
     environmentName,
+    debugChannel,
   );
 }