fix (core): consume stream on abort (#5492)

Author: iteratetograceness

.changeset/thin-numbers-shave.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix (core): improve error handling in streamText's consumeStream method

content/docs/07-reference/01-ai-sdk-core/02-stream-text.mdx

@@ -2274,13 +2274,26 @@ To see `streamText` in action, check out [these examples](#examples).
     },
     {
       name: 'consumeStream',
-      type: '() => Promise<void>',
+      type: '(options?: ConsumeStreamOptions) => Promise<void>',
       description:
-        'Consumes the stream without processing the parts. This is useful to force the stream to finish.',
+        'Consumes the stream without processing the parts. This is useful to force the stream to finish. If an error occurs, it is passed to the optional `onError` callback.',
+      properties: [
+        {
+          type: 'ConsumeStreamOptions',
+          parameters: [
+            {
+              name: 'onError',
+              type: '(error: unknown) => void',
+              isOptional: true,
+              description: 'The error callback.',
+            },
+          ],
+        },
+      ],
     },
     {
       name: 'pipeDataStreamToResponse',
-      type: '(response: ServerResponse, options: PipeDataStreamToResponseOptions } => void',
+      type: '(response: ServerResponse, options: PipeDataStreamToResponseOptions) => void',
       description:
         'Writes stream data output to a Node.js response-like object. It sets a `Content-Type` header to `text/plain; charset=utf-8` and writes each stream data part as a separate chunk.',
       properties: [

examples/next-openai/app/api/use-chat-resilient-persistence/route.ts

@@ -1,6 +1,6 @@
 import { openai } from '@ai-sdk/openai';
-import { appendResponseMessages, createIdGenerator, streamText } from 'ai';
 import { saveChat } from '@util/chat-store';
+import { appendResponseMessages, createIdGenerator, streamText } from 'ai';
 
 export async function POST(req: Request) {
   const { messages, id } = await req.json();
@@ -26,7 +26,12 @@ export async function POST(req: Request) {
 
   // consume the stream to ensure it runs to completion and triggers onFinish
   // even when the client response is aborted (e.g. when the browser tab is closed).
-  result.consumeStream(); // no await
+  // no await
+  result.consumeStream({
+    onError: error => {
+      console.log('Error during background stream consumption: ', error); // optional error callback
+    },
+  });
 
   return result.toDataStreamResponse();
 }

examples/next-openai/app/use-chat-resilient-persistence/[id]/chat.tsx

@@ -1,19 +1,20 @@
 'use client';
 
-import { createIdGenerator } from 'ai';
 import { Message, useChat } from '@ai-sdk/react';
+import { createIdGenerator } from 'ai';
 
 export default function Chat({
   id,
   initialMessages,
 }: { id?: string | undefined; initialMessages?: Message[] } = {}) {
-  const { input, status, handleInputChange, handleSubmit, messages } = useChat({
-    api: '/api/use-chat-resilient-persistence',
-    id, // use the provided chatId
-    initialMessages, // initial messages if provided
-    sendExtraMessageFields: true, // send id and createdAt for each message
-    generateId: createIdGenerator({ prefix: 'msgc', size: 16 }), // id format for client-side messages
-  });
+  const { input, status, handleInputChange, handleSubmit, messages, stop } =
+    useChat({
+      api: '/api/use-chat-resilient-persistence',
+      id, // use the provided chatId
+      initialMessages, // initial messages if provided
+      sendExtraMessageFields: true, // send id and createdAt for each message
+      generateId: createIdGenerator({ prefix: 'msgc', size: 16 }), // id format for client-side messages
+    });
 
   return (
     <div className="flex flex-col w-full max-w-md py-24 mx-auto stretch">
@@ -32,6 +33,15 @@ export default function Chat({
           onChange={handleInputChange}
           disabled={status !== 'ready'}
         />
+        {status === 'streaming' && (
+          <button
+            className="fixed bottom-0 w-full max-w-md p-2 mb-8 border border-gray-300 rounded shadow-xl"
+            type="submit"
+            onClick={stop}
+          >
+            Stop
+          </button>
+        )}
       </form>
     </div>
   );

packages/ai/core/generate-text/stream-text-result.ts

@@ -61,6 +61,10 @@ export type DataStreamOptions = {
   experimental_sendStart?: boolean;
 };
 
+export type ConsumeStreamOptions = {
+  onError?: (error: unknown) => void;
+};
+
 /**
 A result object for accessing different stream types and additional information.
  */
@@ -203,8 +207,10 @@ Consumes the stream without processing the parts.
 This is useful to force the stream to finish.
 It effectively removes the backpressure and allows the stream to finish,
 triggering the `onFinish` callback and the promise resolution.
+
+If an error occurs, it is passed to the optional `onError` callback.
   */
-  consumeStream(): Promise<void>;
+  consumeStream(options?: ConsumeStreamOptions): Promise<void>;
 
   /**
   Converts the result to a data stream.

packages/ai/core/generate-text/stream-text.test.ts

@@ -1542,6 +1542,92 @@ describe('streamText', () => {
     });
   });
 
+  describe('result.consumeStream', () => {
+    it('should ignore AbortError during stream consumption', async () => {
+      const result = streamText({
+        model: createTestModel({
+          stream: new ReadableStream({
+            start(controller) {
+              controller.enqueue({ type: 'text-delta', textDelta: 'Hello' });
+              queueMicrotask(() => {
+                controller.error(
+                  Object.assign(new Error('Stream aborted'), {
+                    name: 'AbortError',
+                  }),
+                );
+              });
+            },
+          }),
+        }),
+        prompt: 'test-input',
+      });
+
+      await expect(result.consumeStream()).resolves.not.toThrow();
+    });
+
+    it('should ignore ResponseAborted error during stream consumption', async () => {
+      const result = streamText({
+        model: createTestModel({
+          stream: new ReadableStream({
+            start(controller) {
+              controller.enqueue({ type: 'text-delta', textDelta: 'Hello' });
+              queueMicrotask(() => {
+                controller.error(
+                  Object.assign(new Error('Response aborted'), {
+                    name: 'ResponseAborted',
+                  }),
+                );
+              });
+            },
+          }),
+        }),
+        prompt: 'test-input',
+      });
+
+      await expect(result.consumeStream()).resolves.not.toThrow();
+    });
+
+    it('should ignore any errors during stream consumption', async () => {
+      const result = streamText({
+        model: createTestModel({
+          stream: new ReadableStream({
+            start(controller) {
+              controller.enqueue({ type: 'text-delta', textDelta: 'Hello' });
+              queueMicrotask(() => {
+                controller.error(Object.assign(new Error('Some error')));
+              });
+            },
+          }),
+        }),
+        prompt: 'test-input',
+      });
+
+      await expect(result.consumeStream()).resolves.not.toThrow();
+    });
+
+    it('should call the onError callback with the error', async () => {
+      const onErrorCallback = vi.fn();
+      const result = streamText({
+        model: createTestModel({
+          stream: new ReadableStream({
+            start(controller) {
+              controller.enqueue({ type: 'text-delta', textDelta: 'Hello' });
+              queueMicrotask(() => {
+                controller.error(Object.assign(new Error('Some error')));
+              });
+            },
+          }),
+        }),
+        prompt: 'test-input',
+      });
+
+      await expect(
+        result.consumeStream({ onError: onErrorCallback }),
+      ).resolves.not.toThrow();
+      expect(onErrorCallback).toHaveBeenCalledWith(new Error('Some error'));
+    });
+  });
+
   describe('multiple stream consumption', () => {
     it('should support text stream, ai stream, full stream on single result object', async () => {
       const result = streamText({

packages/ai/core/generate-text/stream-text.ts

@@ -8,6 +8,7 @@ import { InvalidStreamPartError } from '../../errors/invalid-stream-part-error';
 import { NoOutputSpecifiedError } from '../../errors/no-output-specified-error';
 import { StreamData } from '../../streams/stream-data';
 import { asArray } from '../../util/as-array';
+import { consumeStream } from '../../util/consume-stream';
 import { DelayedPromise } from '../../util/delayed-promise';
 import { DataStreamWriter } from '../data-stream/data-stream-writer';
 import { CallSettings } from '../prompt/call-settings';
@@ -53,6 +54,7 @@ import {
 } from './run-tools-transformation';
 import { ResponseMessage, StepResult } from './step-result';
 import {
+  ConsumeStreamOptions,
   DataStreamOptions,
   StreamTextResult,
   TextStreamPart,
@@ -1591,10 +1593,14 @@ However, the LLM results are expected to be small enough to not cause issues.
     );
   }
 
-  async consumeStream(): Promise<void> {
-    const stream = this.fullStream;
-    for await (const part of stream) {
-      // no op
+  async consumeStream(options?: ConsumeStreamOptions): Promise<void> {
+    try {
+      await consumeStream({
+        stream: this.fullStream,
+        onError: options?.onError,
+      });
+    } catch (error) {
+      options?.onError?.(error);
     }
   }
 

packages/ai/util/consume-stream.ts

@@ -8,10 +8,22 @@
  * @param {ReadableStream} stream - The ReadableStream to be consumed.
  * @returns {Promise<void>} A promise that resolves when the stream is fully consumed.
  */
-export async function consumeStream(stream: ReadableStream): Promise<void> {
+export async function consumeStream({
+  stream,
+  onError,
+}: {
+  stream: ReadableStream;
+  onError?: (error: unknown) => void;
+}): Promise<void> {
   const reader = stream.getReader();
-  while (true) {
-    const { done } = await reader.read();
-    if (done) break;
+  try {
+    while (true) {
+      const { done } = await reader.read();
+      if (done) break;
+    }
+  } catch (error) {
+    onError?.(error);
+  } finally {
+    reader.releaseLock();
   }
 }