Fix streaming server actions

packages/next/src/server/app-render/action-handler.ts

@@ -617,8 +617,9 @@ export async function handleAction({
           type: 'done',
           result: await generateFlight(req, ctx, requestStore, {
             actionResult: promise,
-            // We didn't execute an action, so no revalidations could have occurred. We can skip rendering the page.
-            skipFlight: true,
+            // We didn't execute an action, so no revalidations could have
+            // occurred. We can skip rendering the page.
+            skipPageRendering: true,
             temporaryReferences,
           }),
         }
@@ -729,16 +730,16 @@ export async function handleAction({
                 // Only warn if it's a server action, otherwise skip for other post requests
                 warnBadServerActionRequest()
 
-                const actionReturnedState =
-                  await executeActionAndPrepareForRender(
-                    action as () => Promise<unknown>,
-                    [],
-                    workStore,
-                    requestStore
-                  )
+                const { actionResult } = await executeActionAndPrepareForRender(
+                  action as () => Promise<unknown>,
+                  [],
+                  workStore,
+                  requestStore,
+                  actionWasForwarded
+                )
 
                 const formState = await decodeFormState(
-                  actionReturnedState,
+                  actionResult,
                   formData,
                   serverModuleMap
                 )
@@ -923,16 +924,16 @@ export async function handleAction({
                 // Only warn if it's a server action, otherwise skip for other post requests
                 warnBadServerActionRequest()
 
-                const actionReturnedState =
-                  await executeActionAndPrepareForRender(
-                    action as () => Promise<unknown>,
-                    [],
-                    workStore,
-                    requestStore
-                  )
+                const { actionResult } = await executeActionAndPrepareForRender(
+                  action as () => Promise<unknown>,
+                  [],
+                  workStore,
+                  requestStore,
+                  actionWasForwarded
+                )
 
                 const formState = await decodeFormState(
-                  actionReturnedState,
+                  actionResult,
                   formData,
                   serverModuleMap
                 )
@@ -1022,27 +1023,32 @@ export async function handleAction({
             actionId!
           ]
 
-        const returnVal = await executeActionAndPrepareForRender(
-          actionHandler,
-          boundActionArguments,
-          workStore,
-          requestStore
-        ).finally(() => {
-          addRevalidationHeader(res, { workStore, requestStore })
-        })
+        const { actionResult, skipPageRendering } =
+          await executeActionAndPrepareForRender(
+            actionHandler,
+            boundActionArguments,
+            workStore,
+            requestStore,
+            actionWasForwarded
+          ).finally(() => {
+            addRevalidationHeader(res, { workStore, requestStore })
+          })
 
         // For form actions, we need to continue rendering the page.
         if (isFetchAction) {
-          const actionResult = await generateFlight(req, ctx, requestStore, {
-            actionResult: Promise.resolve(returnVal),
-            // if the page was not revalidated, or if the action was forwarded from another worker, we can skip the rendering the flight tree
-            skipFlight: !workStore.pathWasRevalidated || actionWasForwarded,
-            temporaryReferences,
-          })
-
           return {
             type: 'done',
-            result: actionResult,
+            result: await generateFlight(req, ctx, requestStore, {
+              actionResult: Promise.resolve(actionResult),
+              skipPageRendering,
+              temporaryReferences,
+              // If we skip page rendering, we need to ensure pending
+              // revalidates are awaited before closing the response. Otherwise,
+              // this will be done after rendering the page.
+              waitUntil: skipPageRendering
+                ? executeRevalidates(workStore)
+                : undefined,
+            }),
           }
         } else {
           // TODO: this shouldn't be reachable, because all non-fetch codepaths return early.
@@ -1101,7 +1107,7 @@ export async function handleAction({
         return {
           type: 'done',
           result: await generateFlight(req, ctx, requestStore, {
-            skipFlight: false,
+            skipPageRendering: false,
             actionResult: promise,
             temporaryReferences,
           }),
@@ -1138,8 +1144,10 @@ export async function handleAction({
         type: 'done',
         result: await generateFlight(req, ctx, requestStore, {
           actionResult: promise,
-          // if the page was not revalidated, or if the action was forwarded from another worker, we can skip the rendering the flight tree
-          skipFlight: !workStore.pathWasRevalidated || actionWasForwarded,
+          // If the page was not revalidated, or if the action was forwarded
+          // from another worker, we can skip rendering the page.
+          skipPageRendering:
+            !workStore.pathWasRevalidated || actionWasForwarded,
           temporaryReferences,
         }),
       }
@@ -1156,29 +1164,45 @@ async function executeActionAndPrepareForRender<
   action: TFn,
   args: Parameters<TFn>,
   workStore: WorkStore,
-  requestStore: RequestStore
-): Promise<Awaited<ReturnType<TFn>>> {
+  requestStore: RequestStore,
+  actionWasForwarded: boolean
+): Promise<{
+  actionResult: Awaited<ReturnType<TFn>>
+  skipPageRendering: boolean
+}> {
   requestStore.phase = 'action'
+  let skipPageRendering = actionWasForwarded
+
   try {
-    return await workUnitAsyncStorage.run(requestStore, () =>
+    const actionResult = await workUnitAsyncStorage.run(requestStore, () =>
       action.apply(null, args)
     )
-  } finally {
-    requestStore.phase = 'render'
 
-    // When we switch to the render phase, cookies() will return
-    // `workUnitStore.cookies` instead of `workUnitStore.userspaceMutableCookies`.
-    // We want the render to see any cookie writes that we performed during the action,
-    // so we need to update the immutable cookies to reflect the changes.
-    synchronizeMutableCookies(requestStore)
+    // If the page was not revalidated, or if the action was forwarded from
+    // another worker, we can skip rendering the page.
+    skipPageRendering ||= !workStore.pathWasRevalidated
 
-    // The server action might have toggled draft mode, so we need to reflect
-    // that in the work store to be up-to-date for subsequent rendering.
-    workStore.isDraftMode = requestStore.draftMode.isEnabled
-
-    // If the action called revalidateTag/revalidatePath, then that might affect data used by the subsequent render,
-    // so we need to make sure all revalidations are applied before that
-    await executeRevalidates(workStore)
+    return { actionResult, skipPageRendering }
+  } finally {
+    if (!skipPageRendering) {
+      requestStore.phase = 'render'
+
+      // When we switch to the render phase, cookies() will return
+      // `workUnitStore.cookies` instead of
+      // `workUnitStore.userspaceMutableCookies`. We want the render to see any
+      // cookie writes that we performed during the action, so we need to update
+      // the immutable cookies to reflect the changes.
+      synchronizeMutableCookies(requestStore)
+
+      // The server action might have toggled draft mode, so we need to reflect
+      // that in the work store to be up-to-date for subsequent rendering.
+      workStore.isDraftMode = requestStore.draftMode.isEnabled
+
+      // If the action called revalidateTag/revalidatePath, then that might
+      // affect data used by the subsequent render, so we need to make sure all
+      // revalidations are applied before that.
+      await executeRevalidates(workStore)
+    }
   }
 }
 

packages/next/src/server/app-render/app-render.tsx

@@ -453,7 +453,7 @@ async function generateDynamicRSCPayload(
   ctx: AppRenderContext,
   options?: {
     actionResult?: ActionResult
-    skipFlight?: boolean
+    skipPageRendering?: boolean
     runtimePrefetchSentinel?: number
   }
 ): Promise<RSCPayload> {
@@ -485,7 +485,7 @@ async function generateDynamicRSCPayload(
 
   const serveStreamingMetadata = !!ctx.renderOpts.serveStreamingMetadata
 
-  if (!options?.skipFlight) {
+  if (!options?.skipPageRendering) {
     const preloadCallbacks: PreloadCallbacks = []
 
     const { Viewport, Metadata, MetadataOutlet } = createMetadataComponents({
@@ -588,10 +588,11 @@ async function generateDynamicFlightRenderResult(
   requestStore: RequestStore,
   options?: {
     actionResult: ActionResult
-    skipFlight: boolean
+    skipPageRendering: boolean
     componentTree?: CacheNodeSeedData
     preloadCallbacks?: PreloadCallbacks
     temporaryReferences?: WeakMap<any, string>
+    waitUntil?: Promise<unknown>
   }
 ): Promise<RenderResult> {
   const {
@@ -649,9 +650,11 @@ async function generateDynamicFlightRenderResult(
     }
   )
 
-  return new FlightRenderResult(flightReadableStream, {
-    fetchMetrics: workStore.fetchMetrics,
-  })
+  return new FlightRenderResult(
+    flightReadableStream,
+    { fetchMetrics: workStore.fetchMetrics },
+    options?.waitUntil
+  )
 }
 
 type RenderToReadableStreamServerOptions = NonNullable<

packages/next/src/server/app-render/flight-render-result.ts

@@ -7,8 +7,13 @@ import RenderResult, { type RenderResultMetadata } from '../render-result'
 export class FlightRenderResult extends RenderResult {
   constructor(
     response: string | ReadableStream<Uint8Array>,
-    metadata: RenderResultMetadata = {}
+    metadata: RenderResultMetadata = {},
+    waitUntil?: Promise<unknown>
   ) {
-    super(response, { contentType: RSC_CONTENT_TYPE_HEADER, metadata })
+    super(response, {
+      contentType: RSC_CONTENT_TYPE_HEADER,
+      metadata,
+      waitUntil,
+    })
   }
 }

packages/next/src/server/lib/patch-fetch.ts

@@ -220,25 +220,26 @@ async function createCachedDynamicResponse(
     .finally(handleUnlock)
 
   const pendingRevalidateKey = `cache-set-${cacheKey}`
-  workStore.pendingRevalidates ??= {}
+  const pendingRevalidates = (workStore.pendingRevalidates ??= {})
 
-  if (pendingRevalidateKey in workStore.pendingRevalidates) {
-    // there is already a pending revalidate entry that we need to await to
-    // avoid race conditions
-    await workStore.pendingRevalidates[pendingRevalidateKey]
+  let pendingRevalidatePromise = Promise.resolve()
+  if (pendingRevalidateKey in pendingRevalidates) {
+    // There is already a pending revalidate entry that we need to await to
+    // avoid race conditions.
+    pendingRevalidatePromise = pendingRevalidates[pendingRevalidateKey]
   }
 
-  workStore.pendingRevalidates[pendingRevalidateKey] = cacheSetPromise.finally(
-    () => {
+  pendingRevalidates[pendingRevalidateKey] = pendingRevalidatePromise
+    .then(() => cacheSetPromise)
+    .finally(() => {
       // If the pending revalidate is not present in the store, then we have
       // nothing to delete.
-      if (!workStore.pendingRevalidates?.[pendingRevalidateKey]) {
+      if (!pendingRevalidates?.[pendingRevalidateKey]) {
         return
       }
 
-      delete workStore.pendingRevalidates[pendingRevalidateKey]
-    }
-  )
+      delete pendingRevalidates[pendingRevalidateKey]
+    })
 
   return cloned2
 }

test/e2e/app-dir/actions-streaming/actions-streaming.test.ts

@@ -0,0 +1,43 @@
+import { nextTestSetup } from 'e2e-utils'
+import { retry, waitFor } from 'next-test-utils'
+
+describe('actions-streaming', () => {
+  const { next } = nextTestSetup({
+    files: __dirname,
+  })
+
+  describe('actions returning a ReadableStream', () => {
+    it('should properly stream the response without buffering', async () => {
+      const browser = await next.browser('/readable-stream')
+      await browser.elementById('stream-button').click()
+
+      expect(await browser.elementById('stream-button').text()).toBe(
+        'Streaming...'
+      )
+
+      // If we're streaming properly, we should see the first chunks arrive
+      // quickly.
+      expect(await browser.elementByCss('h3').text()).toMatch(
+        /Received \d+ chunks/
+      )
+      expect(await browser.elementById('chunks').text()).toInclude(
+        'Lorem ipsum dolor sit'
+      )
+
+      // Finally, wait for the response to finish streaming.
+      await waitFor(5000)
+      await retry(
+        async () => {
+          expect(await browser.elementByCss('h3').text()).toBe(
+            'Received 50 chunks'
+          )
+          expect(await browser.elementById('stream-button').text()).toBe(
+            'Start Stream'
+          )
+        },
+        10000,
+        1000
+      )
+    })
+  })
+})

test/e2e/app-dir/actions-streaming/app/layout.tsx

@@ -0,0 +1,8 @@
+import { ReactNode } from 'react'
+export default function Root({ children }: { children: ReactNode }) {
+  return (
+    <html>
+      <body>{children}</body>
+    </html>
+  )
+}

test/e2e/app-dir/actions-streaming/app/page.tsx

@@ -0,0 +1,9 @@
+import Link from 'next/link'
+
+export default function Page() {
+  return (
+    <p>
+      <Link href="/readable-stream">Readable Stream</Link>
+    </p>
+  )
+}

test/e2e/app-dir/actions-streaming/app/readable-stream/actions.ts

@@ -0,0 +1,7 @@
+'use server'
+
+export async function streamData(origin: string) {
+  const response = await fetch(new URL('/readable-stream/api', origin))
+
+  return response.body!
+}

test/e2e/app-dir/actions-streaming/app/readable-stream/api/route.ts

@@ -0,0 +1,20 @@
+import { setTimeout } from 'timers/promises'
+
+const loremIpsum =
+  'Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt.\n'
+
+export async function GET() {
+  const encoder = new TextEncoder()
+
+  const stream = new ReadableStream({
+    async start(controller) {
+      for (let i = 0; i < 50; i++) {
+        await setTimeout(100)
+        controller.enqueue(encoder.encode(loremIpsum))
+      }
+      controller.close()
+    },
+  })
+
+  return new Response(stream, { headers: { 'Content-Type': 'text/plain' } })
+}