fix: unstable_io() returns hanging promise during prerendering (#979)

* fix: unstable_io() returns hanging promise during prerendering (#972)

During prerendering (static export, --prerender-all, TPR), unstable_io()
must return a hanging promise to suspend React past the IO boundary.
Previously it always returned a resolved promise, matching only the
browser/client implementation.

Changes:
- Add makeHangingPromise utility (never-resolving promise, rejects on abort)
- Define typed WorkUnitStore discriminated union in workUnitAsyncStorage
- unstable_io() branches on work unit store type per Next.js's io.ts
- Set prerender work unit store in RSC entry when VINEXT_PRERENDER=1
- Add tests for prerender, request, cache, and abort scenarios

* fix: add knip entries for prerender-work-unit and work-unit types (#972)

* fix(prerender): address PR review feedback for unstable_io work unit store

- AbortController now properly aborted in .finally() to prevent listener leaks

- Add optional route field to PrerenderStore for better error messages

- Pass pathname as route in app-rsc-entry handler

- Use actual route instead of hardcoded 'unknown' in unstable_io()

- Add explicit return in default branch of cache.ts for clarity

- Remove premature knip entry for make-hanging-promise.ts

- Replace V8-internal .status check with Promise.race pattern in tests

* fix(app-rsc-entry): place route option in correct location

Move { route: __pathname } to __runWithPrerenderWorkUnit() call
instead of incorrectly appending to _handleRequest closing brace.
Fixes parse error introduced in 035a54f.

Refs #979

* fix: address prerender hanging promise review feedback

- Suppress unhandled rejection when signal already aborted in makeHangingPromise
- Add lazy pathname extraction via getter function in runWithPrerenderWorkUnit
- Add listener array cleanup after abort in makeHangingPromise
- Add test for unhandled rejection suppression on early-return path
- Update entry template snapshots for lazy route getter

* fix: resolve merge conflicts with main for prerender work unit

- Take main's refactored app-rsc-entry.ts as base (extracted runtime primitives)
- Apply prerender-work-unit wrapping around handler function
- Update tests to match refactored code structure (delegated helpers)
- Update entry-templates snapshots for new generated code structure

* fix: resolve merge conflicts with main for test files

- Keep main's tests/app-router.test.ts (already has correct refactored tests)
- Update entry-templates snapshots to include prerender-work-unit generated code

Author: Divkix

knip.ts

@@ -57,6 +57,12 @@ export default {
         "src/shims/internal/api-utils.ts",
         "src/shims/internal/app-router-context.ts",
         "src/shims/internal/utils.ts",
+        // Typed WorkUnitStore exports consumed by cache.ts via AsyncLocalStorage
+        // generic — knip cannot trace type-only dependencies through ALS.
+        "src/shims/internal/work-unit-async-storage.ts",
+        // Imported via template string in app-rsc-entry.ts (generated code),
+        // so knip cannot trace the import statically.
+        "src/server/prerender-work-unit-setup.ts",
       ],
       project: ["src/**/*.{ts,tsx}"],
     },

packages/vinext/src/entries/app-rsc-entry.ts

@@ -74,6 +74,10 @@ const appPrerenderEndpointsPath = resolveEntryPath(
   "../server/app-prerender-endpoints.js",
   import.meta.url,
 );
+const prerenderWorkUnitSetupPath = resolveEntryPath(
+  "../server/prerender-work-unit-setup.js",
+  import.meta.url,
+);
 const rscStreamHintsPath = resolveEntryPath("../server/rsc-stream-hints.js", import.meta.url);
 const isrCachePath = resolveEntryPath("../server/isr-cache.js", import.meta.url);
 const rootParamsShimPath = resolveEntryPath("../shims/root-params.js", import.meta.url);
@@ -285,6 +289,7 @@ import {
 } from ${JSON.stringify(isrCachePath)};
 // Import server-only state module to register ALS-backed accessors.
 import "vinext/navigation-state";
+import { runWithPrerenderWorkUnit as __runWithPrerenderWorkUnit } from ${JSON.stringify(prerenderWorkUnitSetupPath)};
 import { runWithRequestContext as _runWithUnifiedCtx, createRequestContext as _createUnifiedCtx } from "vinext/unified-request-context";
 import { reportRequestError as _reportRequestError } from "vinext/instrumentation";
 import { flattenErrorCauses as __flattenErrorCauses } from ${JSON.stringify(errorCausePath)};
@@ -734,7 +739,8 @@ export default async function handler(request, ctx) {
     executionContext: ctx ?? _getRequestExecutionContext() ?? null,
     unstableCacheRevalidation: "background",
   });
-  return _runWithUnifiedCtx(__uCtx, async () => {
+  return _runWithUnifiedCtx(__uCtx, () =>
+    __runWithPrerenderWorkUnit(async () => {
     _ensureFetchPatch();
     const __reqCtx = requestContextFromRequest(request);
     // Per-request container for middleware state. Passed into
@@ -775,7 +781,8 @@ export default async function handler(request, ctx) {
       }
     }
     return response;
-  });
+    }, { route: () => new URL(request.url).pathname })
+  );
 }
 
 async function _handleRequest(request, __reqCtx, _mwCtx) {

packages/vinext/src/server/prerender-work-unit-setup.ts

@@ -0,0 +1,34 @@
+/**
+ * Sets up the work unit async storage for prerendering.
+ *
+ * When VINEXT_PRERENDER=1, wraps execution in a workUnitAsyncStorage.run()
+ * with a PrerenderStore so that dynamic APIs (e.g., unstable_io()) can
+ * detect the prerender context and return hanging promises.
+ *
+ * Used by: app-rsc-entry.ts handler template.
+ *
+ * TODO: If future dynamic APIs need request-scoped stores for normal (non-prerender)
+ * requests, add a `{ type: "request" }` store during normal request handling.
+ */
+import { workUnitAsyncStorage } from "../shims/internal/work-unit-async-storage.js";
+
+export function runWithPrerenderWorkUnit<T>(
+  fn: () => Promise<T>,
+  options?: { route?: string | (() => string) },
+): Promise<T> {
+  if (process.env.VINEXT_PRERENDER === "1") {
+    const controller = new AbortController();
+    const route = typeof options?.route === "function" ? options.route() : options?.route;
+    return workUnitAsyncStorage
+      .run(
+        {
+          type: "prerender",
+          renderSignal: controller.signal,
+          route,
+        },
+        fn,
+      )
+      .finally(() => controller.abort());
+  }
+  return fn();
+}

packages/vinext/src/shims/cache.ts

@@ -25,6 +25,8 @@ import {
   getRequestContext,
   runWithUnifiedStateMutation,
 } from "./unified-request-context.js";
+import { workUnitAsyncStorage } from "./internal/work-unit-async-storage.js";
+import { makeHangingPromise } from "./internal/make-hanging-promise.js";
 
 // ---------------------------------------------------------------------------
 // Lazy accessor for cache context — avoids circular imports with cache-runtime.
@@ -447,17 +449,58 @@ const _resolvedIOPromise: Promise<void> = Promise.resolve(undefined);
 (_resolvedIOPromise as unknown as Record<string, unknown>).value = undefined;
 
 /**
- * Marks an IO boundary in server components by returning a resolved promise.
+ * Marks an IO boundary in server components by returning a resolved promise
+ * during requests and a hanging promise during prerendering.
  *
  * See: https://github.com/vercel/next.js/pull/92521
  * Guard removed: https://github.com/vercel/next.js/pull/92923
  *
- * In Next.js, `unstable_io()` during prerendering contexts returns a hanging
- * promise to prevent execution past the IO boundary. vinext does support
- * prerendering (static export, --prerender-all, TPR), but the hanging IO
- * boundary behavior is not yet implemented, so this always resolves immediately.
+ * Ported from Next.js: packages/next/src/server/request/io.ts
+ * https://github.com/vercel/next.js/blob/canary/packages/next/src/server/request/io.ts
+ *
+ * Behavior by work unit type:
+ * - request → resolve immediately (no delay needed for dynamic SSR)
+ * - prerender / prerender-client / prerender-runtime → hang (prevent
+ *   execution past IO boundary during static generation)
+ * - cache / private-cache / unstable-cache → resolve immediately
+ *   (caches capture IO results at fill time)
+ * - generate-static-params → resolve immediately (build time, no prerender to stall)
+ * - prerender-legacy → resolve immediately (no cache components)
+ *
+ * When no work unit store is present (e.g. client-side, standalone script),
+ * resolves immediately — matching the browser/client implementation.
  */
 export function unstable_io(): Promise<void> {
+  const workUnitStore = workUnitAsyncStorage.getStore();
+
+  if (workUnitStore) {
+    switch (workUnitStore.type) {
+      case "request":
+        return _resolvedIOPromise;
+      case "prerender":
+      case "prerender-client":
+      case "prerender-runtime":
+        // Prevent execution past the IO boundary during prerendering.
+        // The hanging promise suspends React's render indefinitely until
+        // the prerender is aborted or completed.
+        return makeHangingPromise(
+          workUnitStore.renderSignal,
+          /* route */ workUnitStore.route ?? "unknown",
+          "`unstable_io()`",
+        );
+      case "cache":
+      case "private-cache":
+      case "unstable-cache":
+      case "generate-static-params":
+      case "prerender-legacy":
+        return _resolvedIOPromise;
+      default:
+        workUnitStore satisfies never;
+        return _resolvedIOPromise;
+    }
+  }
+
+  // No work store — outside rendering context (client, standalone script).
   return _resolvedIOPromise;
 }
 

packages/vinext/src/shims/internal/make-hanging-promise.ts

@@ -0,0 +1,64 @@
+/**
+ * makeHangingPromise — returns a promise that never resolves during prerendering.
+ *
+ * When prerendering, `unstable_io()` must return a hanging promise to prevent
+ * React from executing past the IO boundary. The promise never resolves—it only
+ * rejects if the render signal is aborted (e.g., due to a dynamic error or
+ * cache-fill completion).
+ *
+ * Ported from Next.js: packages/next/src/server/dynamic-rendering-utils.ts
+ * https://github.com/vercel/next.js/blob/canary/packages/next/src/server/dynamic-rendering-utils.ts
+ */
+
+class HangingPromiseRejectionError extends Error {
+  constructor(route: string, expression: string) {
+    super(
+      `Route ${route} used ${expression} during prerendering but the render was aborted. ` +
+        `This is expected when prerendering is cut short (e.g. due to a dynamic access).`,
+    );
+    this.name = "HangingPromiseRejectionError";
+  }
+}
+
+const abortListenersBySignal = new WeakMap<AbortSignal, (() => void)[]>();
+
+function suppressUnhandledRejection(): void {
+  // intentionally empty — suppresses "unhandled rejection" warnings
+}
+
+export function makeHangingPromise<T>(
+  signal: AbortSignal,
+  route: string,
+  expression: string,
+): Promise<T> {
+  if (signal.aborted) {
+    const rejected = Promise.reject(new HangingPromiseRejectionError(route, expression));
+    rejected.catch(suppressUnhandledRejection);
+    return rejected;
+  }
+  const hangingPromise = new Promise<T>((_, reject) => {
+    const boundRejection = reject.bind(null, new HangingPromiseRejectionError(route, expression));
+    const currentListeners = abortListenersBySignal.get(signal);
+    if (currentListeners) {
+      currentListeners.push(boundRejection);
+    } else {
+      const listeners = [boundRejection];
+      abortListenersBySignal.set(signal, listeners);
+      signal.addEventListener(
+        "abort",
+        () => {
+          for (let i = 0; i < listeners.length; i++) {
+            listeners[i]();
+          }
+          listeners.length = 0;
+        },
+        { once: true },
+      );
+    }
+  });
+  // Suppress unhandled rejection — the promise is expected to be used with
+  // React's use() which handles rejections. If never awaited, the abort
+  // rejection is a no-op cleanup.
+  hangingPromise.catch(suppressUnhandledRejection);
+  return hangingPromise;
+}

packages/vinext/src/shims/internal/work-unit-async-storage.ts

@@ -2,13 +2,55 @@
  * Shim for next/dist/server/app-render/work-unit-async-storage.external
  * and next/dist/client/components/request-async-storage.external
  *
- * Used by: @sentry/nextjs (runtime resolve for request context injection).
- * Provides a minimal AsyncLocalStorage-like export that Sentry can detect
- * and use without crashing.
+ * Tracks the current rendering context type so that dynamic APIs
+ * (unstable_io, headers, cookies, etc.) can branch on whether they're
+ * inside a request, prerender, cache scope, or other context.
+ *
+ * Used by: @sentry/nextjs (runtime resolve for request context injection),
+ * unstable_io() for hanging-promise behavior during prerendering.
  */
 import { AsyncLocalStorage } from "node:async_hooks";
 
-export const workUnitAsyncStorage = new AsyncLocalStorage<unknown>();
+// ── WorkUnitStore discriminated union ───────────────────────────────────
+
+export type RequestStore = {
+  readonly type: "request";
+};
+
+export type PrerenderStore = {
+  readonly type: "prerender" | "prerender-client" | "prerender-runtime";
+  /** AbortSignal that fires when the prerender is cancelled or completed. */
+  readonly renderSignal: AbortSignal;
+  /** Optional route identifier for debugging and error messages. */
+  readonly route?: string;
+};
+
+export type CacheStore = {
+  readonly type: "cache" | "private-cache" | "unstable-cache";
+};
+
+export type GenerateStaticParamsStore = {
+  readonly type: "generate-static-params";
+};
+
+export type PrerenderLegacyStore = {
+  readonly type: "prerender-legacy";
+};
+
+/**
+ * Discriminated union of all known work unit types.
+ * Matches Next.js's WorkUnitStore: packages/next/src/server/app-render/work-unit-async-storage.external.ts
+ */
+export type WorkUnitStore =
+  | RequestStore
+  | PrerenderStore
+  | CacheStore
+  | GenerateStaticParamsStore
+  | PrerenderLegacyStore;
+
+export type WorkUnitAsyncStorage = AsyncLocalStorage<WorkUnitStore>;
+
+export const workUnitAsyncStorage: WorkUnitAsyncStorage = new AsyncLocalStorage();
 
 // Legacy name (Next 13.x–14.x)
 export const requestAsyncStorage = workUnitAsyncStorage;