Merge updated PR cloudflare#2710 head into static byte ranges

Author: james-elicx

packages/vinext/src/server/app-rsc-handler.ts

@@ -32,6 +32,7 @@ import {
   closeAfterResponse,
   closeAfterResponseWithBody,
   createRequestContext,
+  preserveFullyBufferedBodyMetadata,
   runWithRequestContext,
 } from "vinext/shims/unified-request-context";
 import { flattenErrorCauses } from "../utils/error-cause.js";
@@ -146,11 +147,14 @@ function applyMiddlewareContextToResponse(
   const headers = new Headers(response.headers);
   mergeMiddlewareResponseHeaders(headers, middlewareContext.headers);
 
-  return new Response(response.body, {
-    status: middlewareContext.status ?? response.status,
-    statusText: response.statusText,
-    headers,
-  });
+  return preserveFullyBufferedBodyMetadata(
+    response,
+    new Response(response.body, {
+      status: middlewareContext.status ?? response.status,
+      statusText: response.statusText,
+      headers,
+    }),
+  );
 }
 
 type DispatchMatchedPageOptions<TRoute> = {

packages/vinext/src/server/metadata-route-response.ts

@@ -10,6 +10,7 @@ import {
   type SitemapEntry,
 } from "./metadata-routes.js";
 import { notFoundResponse } from "./http-error-responses.js";
+import { markFullyBufferedBody } from "vinext/shims/unified-request-context";
 
 type AppPageParams = Record<string, string | string[]>;
 type MetadataRouteFunction = (props: Record<string, unknown>) => unknown;
@@ -223,12 +224,15 @@ async function handleGeneratedSitemap(
   if (!isSitemapEntries(result)) {
     throw new TypeError("Metadata sitemap routes must return an array.");
   }
-  return new Response(sitemapToXml(result), {
-    headers: {
-      "Content-Type": route.contentType,
-      "Cache-Control": metadataRouteCacheHeader(route),
-    },
-  });
+  // Body serialized to a string here — fully materialized, no producer left.
+  return markFullyBufferedBody(
+    new Response(sitemapToXml(result), {
+      headers: {
+        "Content-Type": route.contentType,
+        "Cache-Control": metadataRouteCacheHeader(route),
+      },
+    }),
+  );
 }
 
 function findGeneratedImageId(
@@ -327,12 +331,15 @@ async function callDynamicMetadataRoute(
     body = JSON.stringify(result);
   }
 
-  return new Response(body, {
-    headers: {
-      "Content-Type": route.contentType,
-      "Cache-Control": metadataRouteCacheHeader(route),
-    },
-  });
+  // Every branch above serialized `result` to a string — fully materialized.
+  return markFullyBufferedBody(
+    new Response(body, {
+      headers: {
+        "Content-Type": route.contentType,
+        "Cache-Control": metadataRouteCacheHeader(route),
+      },
+    }),
+  );
 }
 
 function serveStaticMetadataRoute(route: MetadataRuntimeRoute): Response {
@@ -348,12 +355,15 @@ function serveStaticMetadataRoute(route: MetadataRuntimeRoute): Response {
     for (let index = 0; index < binary.length; index++) {
       bytes[index] = binary.charCodeAt(index);
     }
-    return new Response(bytes, {
-      headers: {
-        "Content-Type": route.contentType,
-        "Cache-Control": metadataRouteCacheHeader(route),
-      },
-    });
+    // Static file bytes, fully in memory — no producer left.
+    return markFullyBufferedBody(
+      new Response(bytes, {
+        headers: {
+          "Content-Type": route.contentType,
+          "Cache-Control": metadataRouteCacheHeader(route),
+        },
+      }),
+    );
   } catch (error) {
     const reason = error instanceof Error && error.message ? `: ${error.message}` : "";
     throw new Error(

packages/vinext/src/shims/server.ts

@@ -23,6 +23,7 @@ import {
   getRequestContext,
   isInsideUnifiedScope,
   queueAfterCallback,
+  trackAfterPromise,
 } from "./unified-request-context.js";
 import { assertSafeNavigationUrl } from "./url-safety.js";
 import { hasBasePath, stripBasePath } from "../utils/base-path.js";
@@ -1243,9 +1244,12 @@ export function after<T>(task: Promise<T> | (() => T | Promise<T>)): void {
     if (task == null || typeof (task as PromiseLike<T>).then !== "function") {
       throw new TypeError("`after()`: Argument must be a promise or a function");
     }
-    const guarded = Promise.resolve(task).catch((error) => {
-      console.error("[vinext] after() task failed:", error);
-    });
+    const guarded = trackAfterPromise(
+      requestContext,
+      Promise.resolve(task).catch((error) => {
+        console.error("[vinext] after() task failed:", error);
+      }),
+    );
     getRequestExecutionContext()?.waitUntil(guarded);
     return;
   }

packages/vinext/src/shims/unified-request-context.ts

@@ -63,6 +63,7 @@ export type AfterRequestContext = {
   callbacks: Array<() => unknown>;
   responseClosed: boolean;
   pendingCallbacks: number;
+  pendingPromises: number;
   completion: Promise<void> | null;
   resolveCompletion: (() => void) | null;
 };
@@ -129,6 +130,7 @@ export function createRequestContext(opts?: Partial<UnifiedRequestContext>): Uni
       callbacks: [],
       responseClosed: false,
       pendingCallbacks: 0,
+      pendingPromises: 0,
       completion: null,
       resolveCompletion: null,
     },
@@ -194,6 +196,14 @@ export function queueAfterCallback(ctx: UnifiedRequestContext, callback: () => u
   }
 }
 
+/** Track promise-form after() work that can register a callback before settling. */
+export function trackAfterPromise<T>(ctx: UnifiedRequestContext, promise: Promise<T>): Promise<T> {
+  ctx.afterContext.pendingPromises += 1;
+  return promise.finally(() => {
+    ctx.afterContext.pendingPromises -= 1;
+  });
+}
+
 /** Bind a callback to every AsyncLocalStorage context active at registration. */
 export function bindRequestContextSnapshot<T>(
   ctx: UnifiedRequestContext,
@@ -226,16 +236,80 @@ export async function closeAfterResponse(ctx: UnifiedRequestContext): Promise<vo
   return state.completion ?? Promise.resolve();
 }
 
-/** Wrap a response so deferred callbacks start on stream completion or cancellation. */
+/**
+ * Whether this request has function-form `after()` work that still needs to
+ * observe the response body closing.
+ *
+ * Promise-form `after(promise)` is included because its continuation can
+ * register a function-form `after()` before the promise settles. Function-form
+ * work needs the body's close observed because its contract is "run once the
+ * response has been sent". `resolveCompletion` stays non-null for as long as
+ * any callback is queued or in flight, so checking it alongside the explicit
+ * counters keeps this correct on re-entry after callbacks have started.
+ */
+export function requiresResponseCloseTracking(ctx: UnifiedRequestContext): boolean {
+  const state = ctx.afterContext;
+  return (
+    state.callbacks.length > 0 ||
+    state.pendingCallbacks > 0 ||
+    state.pendingPromises > 0 ||
+    state.resolveCompletion !== null
+  );
+}
+
+type ResponseWithFullyBufferedBodyMetadata = Response & {
+  __vinextFullyBufferedBody?: boolean;
+};
+
+/**
+ * Mark a response whose body vinext constructed from a fully in-memory string
+ * or byte array, as opposed to a body handed back by user code, which could
+ * still be producing. With no producer left, no `after()` call can originate
+ * from this body — the one signal that makes it safe for
+ * `closeAfterResponseWithBody()` to skip close tracking.
+ *
+ * Not set for a metadata route's `result instanceof Response` passthrough (a
+ * user `icon.tsx`/`opengraph-image.tsx` can return a streaming
+ * `ImageResponse`) or any handler-returned `new Response(stream)` — those
+ * bodies can still be producing and must keep close tracking.
+ */
+export function markFullyBufferedBody(response: Response): Response {
+  (response as ResponseWithFullyBufferedBodyMetadata).__vinextFullyBufferedBody = true;
+  return response;
+}
+
+function isFullyBufferedBody(response: Response): boolean {
+  return (response as ResponseWithFullyBufferedBodyMetadata).__vinextFullyBufferedBody === true;
+}
+
+/** Preserve the internal buffered-body signal when response metadata is rebuilt. */
+export function preserveFullyBufferedBodyMetadata(source: Response, target: Response): Response {
+  return isFullyBufferedBody(source) ? markFullyBufferedBody(target) : target;
+}
+
+/**
+ * Wrap a response so deferred `after()` callbacks start on stream completion
+ * or cancellation. Skipped only when the body is marked fully buffered (see
+ * `markFullyBufferedBody`) and nothing is currently registered — that lets
+ * the runtime send it with an accurate `Content-Length` instead of chunked
+ * transfer encoding.
+ */
 export function closeAfterResponseWithBody(
   response: Response,
   ctx: UnifiedRequestContext,
 ): Response {
   if (!response.body) {
+    // Resolve the after-lifecycle now (a no-op if nothing is queued) so a
+    // callback registered after this call returns doesn't wait forever on a
+    // responseClosed flag nothing else will set.
     queueMicrotask(() => void closeAfterResponse(ctx));
     return response;
   }
 
+  if (isFullyBufferedBody(response) && !requiresResponseCloseTracking(ctx)) {
+    return response;
+  }
+
   const passthrough = new TransformStream<Uint8Array, Uint8Array>();
   void response.body.pipeTo(passthrough.writable).then(
     () => void closeAfterResponse(ctx),