From 6c85dd37c1e9f18465ce638a755f606629dc659f Mon Sep 17 00:00:00 2001
From: Nathan Nguyen <146415969+NathanDrake2406@users.noreply.github.com>
Date: Fri, 3 Apr 2026 14:46:15 +1100
Subject: [PATCH] fix: clean up request context on stream error in
 deferUntilStreamConsumed
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

deferUntilStreamConsumed's pull() handler had no error path — when the
upstream RSC/SSR stream errored mid-consumption, the onFlush callback
(clearRequestContext) was never invoked, leaking per-request state.
The cancel hook handled client disconnect and the flush hook handled
normal completion, but stream errors were uncovered. Adds an error
handler that calls once() before propagating the error.
---
 packages/vinext/src/server/app-page-stream.ts | 20 ++++---
 tests/app-page-stream.test.ts                 | 55 +++++++++++++++++++
 2 files changed, 68 insertions(+), 7 deletions(-)

diff --git a/packages/vinext/src/server/app-page-stream.ts b/packages/vinext/src/server/app-page-stream.ts
index 7f2f78365..1f7443325 100644
--- a/packages/vinext/src/server/app-page-stream.ts
+++ b/packages/vinext/src/server/app-page-stream.ts
@@ -106,13 +106,19 @@ export function deferUntilStreamConsumed(
   const reader = piped.getReader();
   return new ReadableStream<Uint8Array>({
     pull(controller) {
-      return reader.read().then(({ done, value }) => {
-        if (done) {
-          controller.close();
-        } else {
-          controller.enqueue(value);
-        }
-      });
+      return reader.read().then(
+        ({ done, value }) => {
+          if (done) {
+            controller.close();
+          } else {
+            controller.enqueue(value);
+          }
+        },
+        (error) => {
+          once();
+          controller.error(error);
+        },
+      );
     },
     cancel(reason) {
       // Stream cancelled before fully consumed (e.g. client disconnected).
diff --git a/tests/app-page-stream.test.ts b/tests/app-page-stream.test.ts
index 705b61e4c..2fe5d6989 100644
--- a/tests/app-page-stream.test.ts
+++ b/tests/app-page-stream.test.ts
@@ -2,6 +2,7 @@ import { describe, expect, it, vi } from "vite-plus/test";
 import {
   createAppPageFontData,
   createAppPageRscErrorTracker,
+  deferUntilStreamConsumed,
   renderAppPageHtmlResponse,
   renderAppPageHtmlStream,
   renderAppPageHtmlStreamWithRecovery,
@@ -100,6 +101,60 @@ describe("app page stream helpers", () => {
     expect(contextCleared).toHaveLength(1);
   });
 
+  it("calls onFlush when the upstream stream errors mid-consumption", async () => {
+    const onFlush = vi.fn();
+    const streamError = new Error("component threw during streaming");
+
+    // Emit one chunk, then error on the next pull — simulates a component
+    // throwing partway through RSC/SSR streaming.
+    let pullCount = 0;
+    const source = new ReadableStream<Uint8Array>({
+      pull(controller) {
+        pullCount++;
+        if (pullCount === 1) {
+          controller.enqueue(new TextEncoder().encode("partial"));
+        } else {
+          controller.error(streamError);
+        }
+      },
+    });
+
+    const wrapped = deferUntilStreamConsumed(source, onFlush);
+    const reader = wrapped.getReader();
+
+    // First read succeeds with the enqueued chunk.
+    const { value } = await reader.read();
+    expect(new TextDecoder().decode(value)).toBe("partial");
+
+    // Second read should surface the upstream error.
+    await expect(reader.read()).rejects.toThrow("component threw during streaming");
+
+    // onFlush must have been called despite the error — this is the bug fix.
+    expect(onFlush).toHaveBeenCalledTimes(1);
+  });
+
+  it("calls onFlush only once when the stream errors then is cancelled", async () => {
+    const onFlush = vi.fn();
+    const streamError = new Error("stream error");
+
+    // Error on the very first pull — simulates immediate failure.
+    const source = new ReadableStream<Uint8Array>({
+      pull(controller) {
+        controller.error(streamError);
+      },
+    });
+
+    const wrapped = deferUntilStreamConsumed(source, onFlush);
+    const reader = wrapped.getReader();
+
+    // Reading the errored stream triggers the error handler.
+    await expect(reader.read()).rejects.toThrow("stream error");
+
+    // The idempotent once() guard prevents double invocation — even if
+    // some code path triggered cleanup again, onFlush fires exactly once.
+    expect(onFlush).toHaveBeenCalledTimes(1);
+  });
+
   it("builds an HTML response, including link headers, and defers clearing request context until after body is consumed", async () => {
     const clearRequestContext = vi.fn();