ai 6.0.292 → 6.0.294

package/src/util/write-to-server-response.ts

@@ -2,6 +2,10 @@ import type { ServerResponse } from 'node:http';
 
 /**
  * Writes the content of a stream to a server response.
+ *
+ * When the client disconnects before the stream has been fully written
+ * (premature `close`), the stream is cancelled so that upstream resources
+ * can be released, and no further chunks are written.
  */
 export function writeToServerResponse({
   response,
@@ -24,22 +28,57 @@ export function writeToServerResponse({
   }
 
   const reader = stream.getReader();
+
+  // Detect client disconnects. `close` also fires after a regular `end()`;
+  // `writableFinished` distinguishes the two cases.
+  let clientDisconnected = false;
+  let onDisconnect: (() => void) | undefined;
+  const handleClose = () => {
+    if (response.writableFinished) {
+      return;
+    }
+
+    clientDisconnected = true;
+    onDisconnect?.();
+
+    // cancelling the reader resolves a pending read with `done: true`:
+    reader.cancel(new Error('Client disconnected.')).catch(() => {});
+  };
+  response.once('close', handleClose);
+
+  // The response may have been destroyed before the close listener was
+  // registered. Cancel immediately instead of waiting indefinitely for the
+  // next stream chunk.
+  if (response.destroyed) {
+    handleClose();
+  }
+
   const read = async () => {
     try {
       while (true) {
         const { done, value } = await reader.read();
-        if (done) break;
+        if (done || clientDisconnected || response.destroyed) break;
 
         // Respect backpressure: if write() returns false, wait for 'drain' event
         const canContinue = response.write(value);
         if (!canContinue) {
           await new Promise<void>(resolve => {
+            // don't wait for `drain` on a disconnected response:
+            onDisconnect = resolve;
             response.once('drain', resolve);
           });
+          onDisconnect = undefined;
         }
       }
     } finally {
-      response.end();
+      response.off('close', handleClose);
+
+      if (clientDisconnected || response.destroyed) {
+        // release the stream; ending a destroyed response is not possible:
+        reader.cancel().catch(() => {});
+      } else {
+        response.end();
+      }
     }
   };