fix: ensure background work is finished when response has 3xx or 5xx status code

Author: pieh

src/run/handlers/server.ts

@@ -113,19 +113,7 @@ export default async (request: Request) => {
     setVaryHeaders(response.headers, request, nextConfig)
     setCacheStatusHeader(response.headers)
 
-    // Temporary workaround for an issue where sending a response with an empty
-    // body causes an unhandled error. This doesn't catch everything, but redirects are the
-    // most common case of sending empty bodies. We can't check it directly because these are streams.
-    // The side effect is that responses which do contain data will not be streamed to the client,
-    // but that's fine for redirects.
-    // TODO: Remove once a fix has been rolled out.
-    if ((response.status > 300 && response.status < 400) || response.status >= 500) {
-      const body = await response.text()
-      return new Response(body || null, response)
-    }
-
-    const keepOpenUntilNextFullyRendered = new TransformStream({
-      async flush() {
       // it's important to keep the stream open until the next handler has finished
       await nextHandlerPromise
 
@@ -138,9 +126,18 @@ export default async (request: Request) => {
       // are resolved. If `context.waitUntil` is available, `requestContext.backgroundWorkPromise` will be empty
       // resolved promised and so awaiting it is no-op
       await requestContext.backgroundWorkPromise
+    }
+
+    const keepOpenUntilNextFullyRendered = new TransformStream({
+      async flush() {
+        await waitForBackgroundWork()
       },
     })
 
+    if (!response.body) {
+      await waitForBackgroundWork()
+    }
+
     return new Response(response.body?.pipeThrough(keepOpenUntilNextFullyRendered), response)
   })
 }