[Flight] Allow aborting encodeReply

packages/react-client/src/ReactFlightReplyClient.js

@@ -185,7 +185,7 @@ export function processReply(
   temporaryReferences: void | TemporaryReferenceSet,
   resolve: (string | FormData) => void,
   reject: (error: mixed) => void,
-): void {
+): (reason: mixed) => void {
   let nextPartId = 1;
   let pendingParts = 0;
   let formData: null | FormData = null;
@@ -841,6 +841,19 @@ export function processReply(
     return JSON.stringify(model, resolveToJSON);
   }
 
+  function abort(reason: mixed): void {
+    if (pendingParts > 0) {
+      pendingParts = 0; // Don't resolve again later.
+      // Resolve with what we have so far, which may have holes at this point.
+      // They'll error when the stream completes on the server.
+      if (formData === null) {
+        resolve(json);
+      } else {
+        resolve(formData);
+      }
+    }
+  }
+
   const json = serializeModel(root, 0);
 
   if (formData === null) {
@@ -854,6 +867,8 @@ export function processReply(
       resolve(formData);
     }
   }
+
+  return abort;
 }
 
 const boundCache: WeakMap<

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

@@ -121,12 +121,12 @@ function createFromFetch<T>(
 
 function encodeReply(
   value: ReactServerValue,
-  options?: {temporaryReferences?: TemporaryReferenceSet},
+  options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
 ): Promise<
   string | URLSearchParams | FormData,
 > /* We don't use URLSearchParams yet but maybe */ {
   return new Promise((resolve, reject) => {
-    processReply(
+    const abort = processReply(
       value,
       '',
       options && options.temporaryReferences
@@ -135,6 +135,18 @@ function encodeReply(
       resolve,
       reject,
     );
+    if (options && options.signal) {
+      const signal = options.signal;
+      if (signal.aborted) {
+        abort((signal: any).reason);
+      } else {
+        const listener = () => {
+          abort((signal: any).reason);
+          signal.removeEventListener('abort', listener);
+        };
+        signal.addEventListener('abort', listener);
+      }
+    }
   });
 }
 

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

@@ -120,12 +120,12 @@ function createFromFetch<T>(
 
 function encodeReply(
   value: ReactServerValue,
-  options?: {temporaryReferences?: TemporaryReferenceSet},
+  options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
 ): Promise<
   string | URLSearchParams | FormData,
 > /* We don't use URLSearchParams yet but maybe */ {
   return new Promise((resolve, reject) => {
-    processReply(
+    const abort = processReply(
       value,
       '',
       options && options.temporaryReferences
@@ -134,6 +134,18 @@ function encodeReply(
       resolve,
       reject,
     );
+    if (options && options.signal) {
+      const signal = options.signal;
+      if (signal.aborted) {
+        abort((signal: any).reason);
+      } else {
+        const listener = () => {
+          abort((signal: any).reason);
+          signal.removeEventListener('abort', listener);
+        };
+        signal.addEventListener('abort', listener);
+      }
+    }
   });
 }
 

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

@@ -149,12 +149,12 @@ function createFromFetch<T>(
 
 function encodeReply(
   value: ReactServerValue,
-  options?: {temporaryReferences?: TemporaryReferenceSet},
+  options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
 ): Promise<
   string | URLSearchParams | FormData,
 > /* We don't use URLSearchParams yet but maybe */ {
   return new Promise((resolve, reject) => {
-    processReply(
+    const abort = processReply(
       value,
       '',
       options && options.temporaryReferences
@@ -163,6 +163,18 @@ function encodeReply(
       resolve,
       reject,
     );
+    if (options && options.signal) {
+      const signal = options.signal;
+      if (signal.aborted) {
+        abort((signal: any).reason);
+      } else {
+        const listener = () => {
+          abort((signal: any).reason);
+          signal.removeEventListener('abort', listener);
+        };
+        signal.addEventListener('abort', listener);
+      }
+    }
   });
 }
 

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

@@ -618,4 +618,20 @@ describe('ReactFlightDOMReply', () => {
     const root = await ReactServerDOMServer.decodeReply(body, webpackServerMap);
     expect(root.prop.obj).toBe(root.prop);
   });
+
+  it('can abort an unresolved model and get the partial result', async () => {
+    const promise = new Promise(r => {});
+    const controller = new AbortController();
+    const bodyPromise = ReactServerDOMClient.encodeReply(
+      {promise: promise, hello: 'world'},
+      {signal: controller.signal},
+    );
+    controller.abort();
+
+    const result = await ReactServerDOMServer.decodeReply(await bodyPromise);
+    expect(result.hello).toBe('world');
+    // TODO: await result.promise should reject at this point because the stream
+    // has closed but that's a bug in both ReactFlightReplyServer and ReactFlightClient.
+    // It just halts in this case.
+  });
 });

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

@@ -120,12 +120,12 @@ function createFromFetch<T>(
 
 function encodeReply(
   value: ReactServerValue,
-  options?: {temporaryReferences?: TemporaryReferenceSet},
+  options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
 ): Promise<
   string | URLSearchParams | FormData,
 > /* We don't use URLSearchParams yet but maybe */ {
   return new Promise((resolve, reject) => {
-    processReply(
+    const abort = processReply(
       value,
       '',
       options && options.temporaryReferences
@@ -134,6 +134,18 @@ function encodeReply(
       resolve,
       reject,
     );
+    if (options && options.signal) {
+      const signal = options.signal;
+      if (signal.aborted) {
+        abort((signal: any).reason);
+      } else {
+        const listener = () => {
+          abort((signal: any).reason);
+          signal.removeEventListener('abort', listener);
+        };
+        signal.addEventListener('abort', listener);
+      }
+    }
   });
 }
 

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

@@ -149,12 +149,12 @@ function createFromFetch<T>(
 
 function encodeReply(
   value: ReactServerValue,
-  options?: {temporaryReferences?: TemporaryReferenceSet},
+  options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
 ): Promise<
   string | URLSearchParams | FormData,
 > /* We don't use URLSearchParams yet but maybe */ {
   return new Promise((resolve, reject) => {
-    processReply(
+    const abort = processReply(
       value,
       '',
       options && options.temporaryReferences
@@ -163,6 +163,18 @@ function encodeReply(
       resolve,
       reject,
     );
+    if (options && options.signal) {
+      const signal = options.signal;
+      if (signal.aborted) {
+        abort((signal: any).reason);
+      } else {
+        const listener = () => {
+          abort((signal: any).reason);
+          signal.removeEventListener('abort', listener);
+        };
+        signal.addEventListener('abort', listener);
+      }
+    }
   });
 }