From 720de7f81a21e4be4d9dcc0c1040c123a7bf2dc6 Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Fri, 20 Dec 2024 09:01:12 -0800
Subject: [PATCH] (cherrypick) [Flight Reply] Reject any new Chunks not yet
 discovered at the time of reportGlobalError (#31840) (#31872)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

We might have already resolved models that are not pending and so are
not rejected by aborting the stream. When those later get parsed they
might discover new chunks which end up as pending. These should be
errored since they will never be able to resolve later.

This avoids infinitely hanging the stream.

This same fix needs to be ported to ReactFlightClient that has the same
issue.

Co-authored-by: Sebastian Markbåge <sebastian@calyptus.eu>
---
 .../react-server/src/ReactFlightReplyServer.js | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/packages/react-server/src/ReactFlightReplyServer.js b/packages/react-server/src/ReactFlightReplyServer.js
index 291da78707..e2056d7df2 100644
--- a/packages/react-server/src/ReactFlightReplyServer.js
+++ b/packages/react-server/src/ReactFlightReplyServer.js
@@ -135,6 +135,8 @@ export type Response = {
   _formData: FormData,
   _chunks: Map<number, SomeChunk<any>>,
   _fromJSON: (key: string, value: JSONValue) => any,
+  _closed: boolean,
+  _closedReason: mixed,
 };
 
 export function getRoot<T>(response: Response): Thenable<T> {
@@ -198,6 +200,14 @@ function createResolvedModelChunk<T>(
   return new Chunk(RESOLVED_MODEL, value, null, response);
 }
 
+function createErroredChunk<T>(
+  response: Response,
+  reason: mixed,
+): ErroredChunk<T> {
+  // $FlowFixMe[invalid-constructor] Flow doesn't support functions as constructors
+  return new Chunk(ERRORED, null, reason, response);
+}
+
 function resolveModelChunk<T>(chunk: SomeChunk<T>, value: string): void {
   if (chunk.status !== PENDING) {
     // We already resolved. We didn't expect to see this.
@@ -297,6 +307,8 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
 // Report that any missing chunks in the model is now going to throw this
 // error upon read. Also notify any pending promises.
 export function reportGlobalError(response: Response, error: Error): void {
+  response._closed = true;
+  response._closedReason = error;
   response._chunks.forEach(chunk => {
     // If this chunk was already resolved or errored, it won't
     // trigger an error but if it wasn't then we need to
@@ -318,6 +330,10 @@ function getChunk(response: Response, id: number): SomeChunk<any> {
     if (backingEntry != null) {
       // We assume that this is a string entry for now.
       chunk = createResolvedModelChunk(response, (backingEntry: any));
+    } else if (response._closed) {
+      // We have already errored the response and we're not going to get
+      // anything more streaming in so this will immediately error.
+      chunk = createErroredChunk(response, response._closedReason);
     } else {
       // We're still waiting on this entry to stream in.
       chunk = createPendingChunk(response);
@@ -519,6 +535,8 @@ export function createResponse(
       }
       return value;
     },
+    _closed: false,
+    _closedReason: null,
   };
   return response;
 }