[Flight] Add unstable_allowPartialStream option to Flight Client (#35731)

When using a partial prerender stream, i.e. a prerender that is
intentionally aborted before all I/O has resolved, consumers of
`createFromReadableStream` would need to keep the stream unclosed to
prevent React Flight from erroring on unresolved chunks. However, some
browsers (e.g. Chrome, Firefox) keep unclosed ReadableStreams with
pending reads as native GC roots, retaining the entire Flight response.

With this PR we're adding an `unstable_allowPartialStream` option, that
allows consumers to close the stream normally. The Flight Client's
`close()` function then transitions pending chunks to halted instead of
erroring them. Halted chunks keep Suspense fallbacks showing (i.e. they
never resolve), and their `.then()` is a no-op so no new listeners
accumulate. Inner stream chunks (ReadableStream/AsyncIterable) are
closed gracefully, and `getChunk()` returns halted chunks for new IDs
that are accessed after closing the response. Blocked chunks are left
alone because they may be waiting on client-side async operations like
module loading, or on forward references to chunks that appeared later
in the stream, both of which resolve independently of closing.

packages/react-client/src/ReactFlightClient.js

@@ -359,6 +359,7 @@ type Response = {
   _stringDecoder: StringDecoder,
   _closed: boolean,
   _closedReason: mixed,
+  _allowPartialStream: boolean,
   _tempRefs: void | TemporaryReferenceSet, // the set temporary references can be resolved from
   _timeOrigin: number, // Profiling-only
   _pendingInitialRender: null | TimeoutID, // Profiling-only,
@@ -1456,9 +1457,19 @@ function getChunk(response: Response, id: number): SomeChunk<any> {
   let chunk = chunks.get(id);
   if (!chunk) {
     if (response._closed) {
-      // We have already errored the response and we're not going to get
+      if (response._allowPartialStream) {
-      // anything more streaming in so this will immediately error.
+        // For partial streams, chunks accessed after close should be HALTED
-      chunk = createErrorChunk(response, response._closedReason);
+        // (never resolve).
+        chunk = createPendingChunk(response);
+        const haltedChunk: HaltedChunk<any> = (chunk: any);
+        haltedChunk.status = HALTED;
+        haltedChunk.value = null;
+        haltedChunk.reason = null;
+      } else {
+        // We have already errored the response and we're not going to get
+        // anything more streaming in so this will immediately error.
+        chunk = createErrorChunk(response, response._closedReason);
+      }
     } else {
       chunk = createPendingChunk(response);
     }
@@ -2655,6 +2666,7 @@ function ResponseInstance(
   encodeFormAction: void | EncodeFormActionCallback,
   nonce: void | string,
   temporaryReferences: void | TemporaryReferenceSet,
+  allowPartialStream: boolean,
   findSourceMapURL: void | FindSourceMapURLCallback, // DEV-only
   replayConsole: boolean, // DEV-only
   environmentName: void | string, // DEV-only
@@ -2674,6 +2686,7 @@ function ResponseInstance(
   this._fromJSON = (null: any);
   this._closed = false;
   this._closedReason = null;
+  this._allowPartialStream = allowPartialStream;
   this._tempRefs = temporaryReferences;
   if (enableProfilerTimer && enableComponentPerformanceTrack) {
     this._timeOrigin = 0;
@@ -2767,6 +2780,7 @@ export function createResponse(
   encodeFormAction: void | EncodeFormActionCallback,
   nonce: void | string,
   temporaryReferences: void | TemporaryReferenceSet,
+  allowPartialStream: boolean,
   findSourceMapURL: void | FindSourceMapURLCallback, // DEV-only
   replayConsole: boolean, // DEV-only
   environmentName: void | string, // DEV-only
@@ -2792,6 +2806,7 @@ export function createResponse(
       encodeFormAction,
       nonce,
       temporaryReferences,
+      allowPartialStream,
       findSourceMapURL,
       replayConsole,
       environmentName,
@@ -5243,11 +5258,45 @@ function createFromJSONCallback(response: Response) {
 }
 
 export function close(weakResponse: WeakResponse): void {
-  // In case there are any remaining unresolved chunks, they won't
+  // In case there are any remaining unresolved chunks, they won't be resolved
-  // be resolved now. So we need to issue an error to those.
+  // now. So we either error or halt them depending on whether partial streams
-  // Ideally we should be able to early bail out if we kept a
+  // are allowed.
-  // ref count of pending chunks.
+  // TODO: Ideally we should be able to bail out early if we kept a ref count of
-  reportGlobalError(weakResponse, new Error('Connection closed.'));
+  // pending chunks.
+  if (hasGCedResponse(weakResponse)) {
+    return;
+  }
+  const response = unwrapWeakResponse(weakResponse);
+  if (response._allowPartialStream) {
+    // For partial streams, we halt pending chunks instead of erroring them.
+    response._closed = true;
+    response._chunks.forEach(chunk => {
+      if (chunk.status === PENDING) {
+        // Clear listeners to release closures and transition to HALTED.
+        // Future .then() calls on HALTED chunks are no-ops.
+        releasePendingChunk(response, chunk);
+        const haltedChunk: HaltedChunk<any> = (chunk: any);
+        haltedChunk.status = HALTED;
+        haltedChunk.value = null;
+        haltedChunk.reason = null;
+      } else if (chunk.status === INITIALIZED && chunk.reason !== null) {
+        // Stream chunk - close gracefully instead of erroring.
+        chunk.reason.close('"$undefined"');
+      }
+    });
+    if (__DEV__) {
+      const debugChannel = response._debugChannel;
+      if (debugChannel !== undefined) {
+        closeDebugChannel(debugChannel);
+        response._debugChannel = undefined;
+        if (debugChannelRegistry !== null) {
+          debugChannelRegistry.unregister(response);
+        }
+      }
+    }
+  } else {
+    reportGlobalError(weakResponse, new Error('Connection closed.'));
+  }
 }
 
 function getCurrentOwnerInDEV(): null | ReactComponentInfo {

packages/react-markup/src/ReactMarkupServer.js

@@ -89,6 +89,7 @@ export function experimental_renderToHTML(
       noServerCallOrFormAction,
       undefined,
       undefined,
+      false,
       undefined,
       false,
       undefined,

packages/react-noop-renderer/src/ReactNoopFlightClient.js

@@ -71,6 +71,7 @@ function read<T>(source: Source, options: ReadOptions): Thenable<T> {
     undefined,
     undefined,
     undefined,
+    false,
     options !== undefined ? options.findSourceMapURL : undefined,
     true,
     undefined,

packages/react-server-dom-esm/src/client/ReactFlightDOMClientBrowser.js

@@ -49,6 +49,7 @@ export type Options = {
   callServer?: CallServerCallback,
   debugChannel?: {writable?: WritableStream, readable?: ReadableStream, ...},
   temporaryReferences?: TemporaryReferenceSet,
+  unstable_allowPartialStream?: boolean,
   findSourceMapURL?: FindSourceMapURLCallback,
   replayConsoleLogs?: boolean,
   environmentName?: string,
@@ -98,6 +99,9 @@ function createResponseFromOptions(options: void | Options) {
     options && options.temporaryReferences
       ? options.temporaryReferences
       : undefined,
+    options && options.unstable_allowPartialStream
+      ? options.unstable_allowPartialStream
+      : false,
     __DEV__ && options && options.findSourceMapURL
       ? options.findSourceMapURL
       : undefined,

packages/react-server-dom-esm/src/client/ReactFlightDOMClientNode.js

@@ -54,6 +54,7 @@ type EncodeFormActionCallback = <A>(
 export type Options = {
   nonce?: string,
   encodeFormAction?: EncodeFormActionCallback,
+  unstable_allowPartialStream?: boolean,
   findSourceMapURL?: FindSourceMapURLCallback,
   replayConsoleLogs?: boolean,
   environmentName?: string,
@@ -104,6 +105,9 @@ function createFromNodeStream<T>(
     options ? options.encodeFormAction : undefined,
     options && typeof options.nonce === 'string' ? options.nonce : undefined,
     undefined, // TODO: If encodeReply is supported, this should support temporaryReferences
+    options && options.unstable_allowPartialStream
+      ? options.unstable_allowPartialStream
+      : false,
     __DEV__ && options && options.findSourceMapURL
       ? options.findSourceMapURL
       : undefined,

packages/react-server-dom-parcel/src/client/ReactFlightDOMClientBrowser.js

@@ -124,6 +124,9 @@ function createResponseFromOptions(options: void | Options) {
     options && options.temporaryReferences
       ? options.temporaryReferences
       : undefined,
+    options && options.unstable_allowPartialStream
+      ? options.unstable_allowPartialStream
+      : false,
     __DEV__ ? findSourceMapURL : undefined,
     __DEV__ ? (options ? options.replayConsoleLogs !== false : true) : false, // defaults to true
     __DEV__ && options && options.environmentName
@@ -207,6 +210,7 @@ function startReadingFromStream(
 export type Options = {
   debugChannel?: {writable?: WritableStream, readable?: ReadableStream, ...},
   temporaryReferences?: TemporaryReferenceSet,
+  unstable_allowPartialStream?: boolean,
   replayConsoleLogs?: boolean,
   environmentName?: string,
   startTime?: number,

packages/react-server-dom-parcel/src/client/ReactFlightDOMClientEdge.js

@@ -77,6 +77,7 @@ export type Options = {
   nonce?: string,
   encodeFormAction?: EncodeFormActionCallback,
   temporaryReferences?: TemporaryReferenceSet,
+  unstable_allowPartialStream?: boolean,
   replayConsoleLogs?: boolean,
   environmentName?: string,
   startTime?: number,
@@ -104,6 +105,9 @@ function createResponseFromOptions(options?: Options) {
     options && options.temporaryReferences
       ? options.temporaryReferences
       : undefined,
+    options && options.unstable_allowPartialStream
+      ? options.unstable_allowPartialStream
+      : false,
     __DEV__ ? findSourceMapURL : undefined,
     __DEV__ && options ? options.replayConsoleLogs === true : false, // defaults to false
     __DEV__ && options && options.environmentName

packages/react-server-dom-parcel/src/client/ReactFlightDOMClientNode.js

@@ -50,6 +50,7 @@ type EncodeFormActionCallback = <A>(
 export type Options = {
   nonce?: string,
   encodeFormAction?: EncodeFormActionCallback,
+  unstable_allowPartialStream?: boolean,
   replayConsoleLogs?: boolean,
   environmentName?: string,
   startTime?: number,
@@ -97,6 +98,9 @@ export function createFromNodeStream<T>(
     options ? options.encodeFormAction : undefined,
     options && typeof options.nonce === 'string' ? options.nonce : undefined,
     undefined, // TODO: If encodeReply is supported, this should support temporaryReferences
+    options && options.unstable_allowPartialStream
+      ? options.unstable_allowPartialStream
+      : false,
     __DEV__ ? findSourceMapURL : undefined,
     __DEV__ && options ? options.replayConsoleLogs === true : false, // defaults to false
     __DEV__ && options && options.environmentName

packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientBrowser.js

@@ -48,6 +48,7 @@ export type Options = {
   callServer?: CallServerCallback,
   debugChannel?: {writable?: WritableStream, readable?: ReadableStream, ...},
   temporaryReferences?: TemporaryReferenceSet,
+  unstable_allowPartialStream?: boolean,
   findSourceMapURL?: FindSourceMapURLCallback,
   replayConsoleLogs?: boolean,
   environmentName?: string,
@@ -97,6 +98,9 @@ function createResponseFromOptions(options: void | Options) {
     options && options.temporaryReferences
       ? options.temporaryReferences
       : undefined,
+    options && options.unstable_allowPartialStream
+      ? options.unstable_allowPartialStream
+      : false,
     __DEV__ && options && options.findSourceMapURL
       ? options.findSourceMapURL
       : undefined,

packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientEdge.js

@@ -76,6 +76,7 @@ export type Options = {
   nonce?: string,
   encodeFormAction?: EncodeFormActionCallback,
   temporaryReferences?: TemporaryReferenceSet,
+  unstable_allowPartialStream?: boolean,
   findSourceMapURL?: FindSourceMapURLCallback,
   replayConsoleLogs?: boolean,
   environmentName?: string,
@@ -104,6 +105,9 @@ function createResponseFromOptions(options: Options) {
     options && options.temporaryReferences
       ? options.temporaryReferences
       : undefined,
+    options && options.unstable_allowPartialStream
+      ? options.unstable_allowPartialStream
+      : false,
     __DEV__ && options && options.findSourceMapURL
       ? options.findSourceMapURL
       : undefined,

packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientNode.js

@@ -57,6 +57,7 @@ type EncodeFormActionCallback = <A>(
 export type Options = {
   nonce?: string,
   encodeFormAction?: EncodeFormActionCallback,
+  unstable_allowPartialStream?: boolean,
   findSourceMapURL?: FindSourceMapURLCallback,
   replayConsoleLogs?: boolean,
   environmentName?: string,
@@ -106,6 +107,9 @@ function createFromNodeStream<T>(
     options ? options.encodeFormAction : undefined,
     options && typeof options.nonce === 'string' ? options.nonce : undefined,
     undefined, // TODO: If encodeReply is supported, this should support temporaryReferences
+    options && options.unstable_allowPartialStream
+      ? options.unstable_allowPartialStream
+      : false,
     __DEV__ && options && options.findSourceMapURL
       ? options.findSourceMapURL
       : undefined,

packages/react-server-dom-unbundled/src/client/ReactFlightDOMClientEdge.js

@@ -76,6 +76,7 @@ export type Options = {
   nonce?: string,
   encodeFormAction?: EncodeFormActionCallback,
   temporaryReferences?: TemporaryReferenceSet,
+  unstable_allowPartialStream?: boolean,
   findSourceMapURL?: FindSourceMapURLCallback,
   replayConsoleLogs?: boolean,
   environmentName?: string,
@@ -104,6 +105,9 @@ function createResponseFromOptions(options: Options) {
     options && options.temporaryReferences
       ? options.temporaryReferences
       : undefined,
+    options && options.unstable_allowPartialStream
+      ? options.unstable_allowPartialStream
+      : false,
     __DEV__ && options && options.findSourceMapURL
       ? options.findSourceMapURL
       : undefined,

packages/react-server-dom-unbundled/src/client/ReactFlightDOMClientNode.js

@@ -57,6 +57,7 @@ type EncodeFormActionCallback = <A>(
 export type Options = {
   nonce?: string,
   encodeFormAction?: EncodeFormActionCallback,
+  unstable_allowPartialStream?: boolean,
   findSourceMapURL?: FindSourceMapURLCallback,
   replayConsoleLogs?: boolean,
   environmentName?: string,
@@ -106,6 +107,9 @@ function createFromNodeStream<T>(
     options ? options.encodeFormAction : undefined,
     options && typeof options.nonce === 'string' ? options.nonce : undefined,
     undefined, // TODO: If encodeReply is supported, this should support temporaryReferences
+    options && options.unstable_allowPartialStream
+      ? options.unstable_allowPartialStream
+      : false,
     __DEV__ && options && options.findSourceMapURL
       ? options.findSourceMapURL
       : undefined,

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js

@@ -2504,6 +2504,171 @@ describe('ReactFlightDOMBrowser', () => {
     expect(container.innerHTML).toBe('');
   });
 
+  it('renders Suspense fallback for unresolved promises with unstable_allowPartialStream', async () => {
+    let resolveGreeting;
+    const greetingPromise = new Promise(resolve => {
+      resolveGreeting = resolve;
+    });
+
+    function App() {
+      return (
+        <Suspense fallback="loading...">
+          <Greeting />
+        </Suspense>
+      );
+    }
+
+    async function Greeting() {
+      const greeting = await greetingPromise;
+      return greeting;
+    }
+
+    const controller = new AbortController();
+    const {pendingResult} = await serverAct(async () => {
+      return {
+        pendingResult: ReactServerDOMStaticServer.prerender(
+          <App />,
+          webpackMap,
+          {
+            signal: controller.signal,
+          },
+        ),
+      };
+    });
+
+    controller.abort();
+    resolveGreeting('Hello, World!');
+    const {prelude} = await serverAct(() => pendingResult);
+
+    function ClientRoot({response}) {
+      return use(response);
+    }
+
+    const response = ReactServerDOMClient.createFromReadableStream(
+      passThrough(prelude),
+      {
+        unstable_allowPartialStream: true,
+      },
+    );
+    const container = document.createElement('div');
+    const errors = [];
+    const root = ReactDOMClient.createRoot(container, {
+      onUncaughtError(err) {
+        errors.push(err);
+      },
+    });
+
+    await act(() => {
+      root.render(<ClientRoot response={response} />);
+    });
+
+    // With `unstable_allowPartialStream`, we should see the fallback instead of a
+    // 'Connection closed.' error
+    expect(errors).toEqual([]);
+    expect(container.innerHTML).toBe('loading...');
+  });
+
+  it('renders client components that are blocked on chunks with unstable_allowPartialStream', async () => {
+    let resolveClientComponentChunk;
+
+    const ClientComponent = clientExports(
+      function ClientComponent({children}) {
+        return <div>{children}</div>;
+      },
+      '42',
+      '/test.js',
+      new Promise(resolve => (resolveClientComponentChunk = resolve)),
+    );
+
+    function App() {
+      return <ClientComponent>Hello, World!</ClientComponent>;
+    }
+
+    const controller = new AbortController();
+    const {pendingResult} = await serverAct(async () => {
+      return {
+        pendingResult: ReactServerDOMStaticServer.prerender(
+          <App />,
+          webpackMap,
+          {
+            signal: controller.signal,
+          },
+        ),
+      };
+    });
+
+    controller.abort();
+    const {prelude} = await serverAct(() => pendingResult);
+
+    function ClientRoot({response}) {
+      return use(response);
+    }
+
+    const response = ReactServerDOMClient.createFromReadableStream(
+      passThrough(prelude),
+      {
+        unstable_allowPartialStream: true,
+      },
+    );
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+
+    await act(() => {
+      root.render(<ClientRoot response={response} />);
+    });
+
+    expect(container.innerHTML).toBe('');
+
+    await act(() => {
+      resolveClientComponentChunk();
+    });
+
+    expect(container.innerHTML).toBe('<div>Hello, World!</div>');
+  });
+
+  it('closes inner ReadableStreams gracefully with unstable_allowPartialStream', async () => {
+    let streamController;
+    const innerStream = new ReadableStream({
+      start(c) {
+        streamController = c;
+      },
+    });
+
+    const abortController = new AbortController();
+    const {pendingResult} = await serverAct(async () => {
+      streamController.enqueue({hello: 'world'});
+      return {
+        pendingResult: ReactServerDOMStaticServer.prerender(
+          {stream: innerStream},
+          webpackMap,
+          {
+            signal: abortController.signal,
+          },
+        ),
+      };
+    });
+
+    abortController.abort();
+    const {prelude} = await serverAct(() => pendingResult);
+
+    const response = await ReactServerDOMClient.createFromReadableStream(
+      passThrough(prelude),
+      {
+        unstable_allowPartialStream: true,
+      },
+    );
+
+    // The inner stream should be readable up to what was enqueued.
+    const reader = response.stream.getReader();
+    const {value, done} = await reader.read();
+    expect(value).toEqual({hello: 'world'});
+    expect(done).toBe(false);
+
+    // The next read should signal the stream is done (closed, not errored).
+    const final = await reader.read();
+    expect(final.done).toBe(true);
+  });
+
   it('can dedupe references inside promises', async () => {
     const foo = {};
     const bar = {
@@ -2902,9 +3067,9 @@ describe('ReactFlightDOMBrowser', () => {
               [
                 "Object.<anonymous>",
                 "/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js",
-                2824,
+                2989,
                 19,
-                2808,
+                2973,
                 89,
               ],
             ],

packages/react-server-dom-webpack/src/client/ReactFlightDOMClientBrowser.js

@@ -48,6 +48,7 @@ export type Options = {
   callServer?: CallServerCallback,
   debugChannel?: {writable?: WritableStream, readable?: ReadableStream, ...},
   temporaryReferences?: TemporaryReferenceSet,
+  unstable_allowPartialStream?: boolean,
   findSourceMapURL?: FindSourceMapURLCallback,
   replayConsoleLogs?: boolean,
   environmentName?: string,
@@ -97,6 +98,9 @@ function createResponseFromOptions(options: void | Options) {
     options && options.temporaryReferences
       ? options.temporaryReferences
       : undefined,
+    options && options.unstable_allowPartialStream
+      ? options.unstable_allowPartialStream
+      : false,
     __DEV__ && options && options.findSourceMapURL
       ? options.findSourceMapURL
       : undefined,

packages/react-server-dom-webpack/src/client/ReactFlightDOMClientEdge.js

@@ -76,6 +76,7 @@ export type Options = {
   nonce?: string,
   encodeFormAction?: EncodeFormActionCallback,
   temporaryReferences?: TemporaryReferenceSet,
+  unstable_allowPartialStream?: boolean,
   findSourceMapURL?: FindSourceMapURLCallback,
   replayConsoleLogs?: boolean,
   environmentName?: string,
@@ -104,6 +105,9 @@ function createResponseFromOptions(options: Options) {
     options && options.temporaryReferences
       ? options.temporaryReferences
       : undefined,
+    options && options.unstable_allowPartialStream
+      ? options.unstable_allowPartialStream
+      : false,
     __DEV__ && options && options.findSourceMapURL
       ? options.findSourceMapURL
       : undefined,

packages/react-server-dom-webpack/src/client/ReactFlightDOMClientNode.js

@@ -57,6 +57,7 @@ type EncodeFormActionCallback = <A>(
 export type Options = {
   nonce?: string,
   encodeFormAction?: EncodeFormActionCallback,
+  unstable_allowPartialStream?: boolean,
   findSourceMapURL?: FindSourceMapURLCallback,
   replayConsoleLogs?: boolean,
   environmentName?: string,
@@ -106,6 +107,9 @@ function createFromNodeStream<T>(
     options ? options.encodeFormAction : undefined,
     options && typeof options.nonce === 'string' ? options.nonce : undefined,
     undefined, // TODO: If encodeReply is supported, this should support temporaryReferences
+    options && options.unstable_allowPartialStream
+      ? options.unstable_allowPartialStream
+      : false,
     __DEV__ && options && options.findSourceMapURL
       ? options.findSourceMapURL
       : undefined,