feat (ai/core): mask data stream error messages. (#2617)

Author: lgrammel

.changeset/tricky-trains-yawn.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat (ai/core): mask data stream error messages with streamText

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

@@ -769,11 +769,6 @@ for await (const textPart of textStream) {
         },
       ],
     },
-    {
-      name: 'toDataStream',
-      type: '(callbacks?: AIStreamCallbacksAndOptions) => data stream',
-      description: 'Converts the result to an data stream.',
-    },
     {
       name: 'pipeDataStreamToResponse',
       type: '(response: ServerResponse, init?: { headers?: Record<string, string>; status?: number } => void',
@@ -826,6 +821,13 @@ for await (const textPart of textStream) {
               optional: true,
               description: 'The stream data object.',
             },
+            {
+              name: 'getErrorMessage',
+              type: '(error: unknown) => string',
+              description:
+                'A function to get the error message from the error object. By default, all errors are masked as "" for safety reasons.',
+              optional: true,
+            },
           ],
         },
       ],

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

@@ -1054,6 +1054,48 @@ describe('result.toDataStreamResponse', () => {
       '0:"world!"\n',
     ]);
   });
+
+  it('should mask error messages by default', async () => {
+    const result = await streamText({
+      model: new MockLanguageModelV1({
+        doStream: async () => ({
+          stream: convertArrayToReadableStream([
+            { type: 'error', error: 'error' },
+          ]),
+          rawCall: { rawPrompt: 'prompt', rawSettings: {} },
+        }),
+      }),
+      prompt: 'test-input',
+    });
+
+    const response = result.toDataStreamResponse();
+
+    assert.deepStrictEqual(await convertResponseStreamToArray(response), [
+      '3:""\n',
+    ]);
+  });
+
+  it('should support custom error messages', async () => {
+    const result = await streamText({
+      model: new MockLanguageModelV1({
+        doStream: async () => ({
+          stream: convertArrayToReadableStream([
+            { type: 'error', error: 'error' },
+          ]),
+          rawCall: { rawPrompt: 'prompt', rawSettings: {} },
+        }),
+      }),
+      prompt: 'test-input',
+    });
+
+    const response = result.toDataStreamResponse({
+      getErrorMessage: error => `custom error message: ${error}`,
+    });
+
+    assert.deepStrictEqual(await convertResponseStreamToArray(response), [
+      '3:"custom error message: error"\n',
+    ]);
+  });
 });
 
 describe('result.toTextStreamResponse', () => {

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

@@ -513,6 +513,16 @@ However, the LLM results are expected to be small enough to not cause issues.
   }
 
   toAIStream(callbacks: AIStreamCallbacksAndOptions = {}) {
+    return this.toDataStream({ callbacks });
+  }
+
+  private toDataStream({
+    callbacks = {},
+    getErrorMessage = () => '', // mask error messages for safety by default
+  }: {
+    callbacks?: AIStreamCallbacksAndOptions;
+    getErrorMessage?: (error: unknown) => string;
+  } = {}) {
     let aggregatedResponse = '';
 
     const callbackTransformer = new TransformStream<
@@ -588,7 +598,7 @@ However, the LLM results are expected to be small enough to not cause issues.
             break;
           case 'error':
             controller.enqueue(
-              formatStreamPart('error', JSON.stringify(chunk.error)),
+              formatStreamPart('error', getErrorMessage(chunk.error)),
             );
             break;
           case 'finish':
@@ -632,7 +642,7 @@ However, the LLM results are expected to be small enough to not cause issues.
       ...init?.headers,
     });
 
-    const reader = this.toAIStream().getReader();
+    const reader = this.toDataStream().getReader();
 
     const read = async () => {
       try {
@@ -688,7 +698,13 @@ However, the LLM results are expected to be small enough to not cause issues.
   }
 
   toDataStreamResponse(
-    options?: ResponseInit | { init?: ResponseInit; data?: StreamData },
+    options?:
+      | ResponseInit
+      | {
+          init?: ResponseInit;
+          data?: StreamData;
+          getErrorMessage?: (error: unknown) => string;
+        },
   ): Response {
     const init: ResponseInit | undefined =
       options == null
@@ -709,9 +725,16 @@ However, the LLM results are expected to be small enough to not cause issues.
         ? options.data
         : undefined;
 
+    const getErrorMessage: ((error: unknown) => string) | undefined =
+      options == null
+        ? undefined
+        : 'getErrorMessage' in options
+        ? options.getErrorMessage
+        : undefined;
+
     const stream = data
-      ? mergeStreams(data.stream, this.toAIStream())
-      : this.toAIStream();
+      ? mergeStreams(data.stream, this.toDataStream({ getErrorMessage }))
+      : this.toDataStream({ getErrorMessage });
 
     return new Response(stream, {
       status: init?.status ?? 200,