fix (ai/core): prevent onFinish from masking stream errors (#4318)

vercel · Jan 8, 2025 · e4ce80c · e4ce80c
1 parent 927936a
commit e4ce80c
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 12 deletions.
diff --git a/.changeset/sweet-carrots-shout.md b/.changeset/sweet-carrots-shout.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix (ai/core): prevent onFinish from masking stream errors
diff --git a/packages/ai/core/generate-text/__snapshots__/stream-text.test.ts.snap b/packages/ai/core/generate-text/__snapshots__/stream-text.test.ts.snap
@@ -1699,7 +1699,7 @@ exports[`streamText > options.maxSteps > 4 steps: initial, continue, continue, c
 ]
 `;
 
-exports[`streamText > options.onFinish > options.onFinish should send correct information 1`] = `
+exports[`streamText > options.onFinish > should send correct information 1`] = `
 {
   "experimental_providerMetadata": {
     "testProvider": {

diff --git a/packages/ai/core/generate-text/stream-text.test.ts b/packages/ai/core/generate-text/stream-text.test.ts
@@ -1640,7 +1640,7 @@ describe('streamText', () => {
   });
 
   describe('options.onFinish', () => {
-    it('options.onFinish should send correct information', async () => {
+    it('should send correct information', async () => {
       let result!: Parameters<
         Required<Parameters<typeof streamText>[0]>['onFinish']
       >[0];
@@ -1692,6 +1692,27 @@ describe('streamText', () => {
 
       expect(result).toMatchSnapshot();
     });
+
+    it('should not prevent error from being forwarded', async () => {
+      const result = streamText({
+        model: new MockLanguageModelV1({
+          doStream: async () => {
+            throw new Error('test error');
+          },
+        }),
+        prompt: 'test-input',
+        onFinish() {}, // just defined; do nothing
+      });
+
+      expect(
+        await convertAsyncIterableToArray(result.fullStream),
+      ).toStrictEqual([
+        {
+          type: 'error',
+          error: new Error('test error'),
+        },
+      ]);
+    });
   });
 
   describe('result.responseMessages', () => {

diff --git a/packages/ai/core/generate-text/stream-text.ts b/packages/ai/core/generate-text/stream-text.ts
@@ -637,18 +637,21 @@ class DefaultStreamTextResult<
 
       async flush(controller) {
         try {
+          if (recordedSteps.length === 0) {
+            return; // no steps recorded (e.g. in error scenario)
+          }
+
           // from last step (when there are errors there may be no last step)
           const lastStep = recordedSteps[recordedSteps.length - 1];
-          if (lastStep) {
-            self.warningsPromise.resolve(lastStep.warnings);
-            self.requestPromise.resolve(lastStep.request);
-            self.responsePromise.resolve(lastStep.response);
-            self.toolCallsPromise.resolve(lastStep.toolCalls);
-            self.toolResultsPromise.resolve(lastStep.toolResults);
-            self.providerMetadataPromise.resolve(
-              lastStep.experimental_providerMetadata,
-            );
-          }
+
+          self.warningsPromise.resolve(lastStep.warnings);
+          self.requestPromise.resolve(lastStep.request);
+          self.responsePromise.resolve(lastStep.response);
+          self.toolCallsPromise.resolve(lastStep.toolCalls);
+          self.toolResultsPromise.resolve(lastStep.toolResults);
+          self.providerMetadataPromise.resolve(
+            lastStep.experimental_providerMetadata,
+          );
 
           // derived:
           const finishReason = recordedFinishReason ?? 'unknown';