fix (core): send buffered text in smooth stream when stream parts change (#5531)

lgrammel · carlbrugger · web-flow · commit c45d100a8ec2 · 2025-04-03T14:03:54.000+02:00
Co-authored-by: Carl Brugger &lt;cebrugg@gmail.com&gt;
diff --git a/.changeset/smooth-mirrors-kneel.md b/.changeset/smooth-mirrors-kneel.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix (core): send buffered text in smooth stream when stream parts change
diff --git a/packages/ai/core/generate-text/smooth-stream.test.ts b/packages/ai/core/generate-text/smooth-stream.test.ts
@@ -173,6 +173,166 @@ describe('smoothStream', () => {
         },
       ]);
     });
+
+    it('should send remaining text buffer before tool call starts', async () => {
+      const stream = convertArrayToReadableStream([
+        { type: 'text-delta', textDelta: 'I will check the' },
+        { type: 'text-delta', textDelta: ' weather in Lon' },
+        { type: 'text-delta', textDelta: 'don.' },
+        { type: 'tool-call', name: 'weather', args: { city: 'London' } },
+        { type: 'step-finish' },
+        { type: 'finish' },
+      ]).pipeThrough(
+        smoothStream({
+          delayInMs: 10,
+          _internal: { delay },
+        })({ tools: {} }),
+      );
+
+      await consumeStream(stream);
+
+      expect(events).toMatchInlineSnapshot(`
+        [
+          "delay 10",
+          {
+            "textDelta": "I ",
+            "type": "text-delta",
+          },
+          "delay 10",
+          {
+            "textDelta": "will ",
+            "type": "text-delta",
+          },
+          "delay 10",
+          {
+            "textDelta": "check ",
+            "type": "text-delta",
+          },
+          "delay 10",
+          {
+            "textDelta": "the ",
+            "type": "text-delta",
+          },
+          "delay 10",
+          {
+            "textDelta": "weather ",
+            "type": "text-delta",
+          },
+          "delay 10",
+          {
+            "textDelta": "in ",
+            "type": "text-delta",
+          },
+          {
+            "textDelta": "London.",
+            "type": "text-delta",
+          },
+          {
+            "args": {
+              "city": "London",
+            },
+            "name": "weather",
+            "type": "tool-call",
+          },
+          {
+            "type": "step-finish",
+          },
+          {
+            "type": "finish",
+          },
+        ]
+      `);
+    });
+
+    it('should send remaining text buffer before tool call starts and tool call streaming is enabled', async () => {
+      const stream = convertArrayToReadableStream([
+        { type: 'text-delta', textDelta: 'I will check the' },
+        { type: 'text-delta', textDelta: ' weather in Lon' },
+        { type: 'text-delta', textDelta: 'don.' },
+        {
+          type: 'tool-call-streaming-start',
+          name: 'weather',
+          args: { city: 'London' },
+        },
+        { type: 'tool-call-delta', name: 'weather', args: { city: 'London' } },
+        { type: 'tool-call', name: 'weather', args: { city: 'London' } },
+        { type: 'step-finish' },
+        { type: 'finish' },
+      ]).pipeThrough(
+        smoothStream({
+          delayInMs: 10,
+          _internal: { delay },
+        })({ tools: {} }),
+      );
+
+      await consumeStream(stream);
+
+      expect(events).toMatchInlineSnapshot(`
+        [
+          "delay 10",
+          {
+            "textDelta": "I ",
+            "type": "text-delta",
+          },
+          "delay 10",
+          {
+            "textDelta": "will ",
+            "type": "text-delta",
+          },
+          "delay 10",
+          {
+            "textDelta": "check ",
+            "type": "text-delta",
+          },
+          "delay 10",
+          {
+            "textDelta": "the ",
+            "type": "text-delta",
+          },
+          "delay 10",
+          {
+            "textDelta": "weather ",
+            "type": "text-delta",
+          },
+          "delay 10",
+          {
+            "textDelta": "in ",
+            "type": "text-delta",
+          },
+          {
+            "textDelta": "London.",
+            "type": "text-delta",
+          },
+          {
+            "args": {
+              "city": "London",
+            },
+            "name": "weather",
+            "type": "tool-call-streaming-start",
+          },
+          {
+            "args": {
+              "city": "London",
+            },
+            "name": "weather",
+            "type": "tool-call-delta",
+          },
+          {
+            "args": {
+              "city": "London",
+            },
+            "name": "weather",
+            "type": "tool-call",
+          },
+          {
+            "type": "step-finish",
+          },
+          {
+            "type": "finish",
+          },
+        ]
+      `);
+    });
   });
 
   describe('line chunking', () => {
diff --git a/packages/ai/core/generate-text/smooth-stream.ts b/packages/ai/core/generate-text/smooth-stream.ts
@@ -44,9 +44,10 @@ export function smoothStream<TOOLS extends ToolSet>({
 
   return () => {
     let buffer = '';
+
     return new TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>({
       async transform(chunk, controller) {
-        if (chunk.type === 'step-finish') {
+        if (chunk.type !== 'text-delta') {
           if (buffer.length > 0) {
             controller.enqueue({ type: 'text-delta', textDelta: buffer });
             buffer = '';
@@ -56,11 +57,6 @@ export function smoothStream<TOOLS extends ToolSet>({
           return;
         }
 
-        if (chunk.type !== 'text-delta') {
-          controller.enqueue(chunk);
-          return;
-        }
-
         buffer += chunk.textDelta;
 
         let match;

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +'ai': patch
 +---
++
 +fix (core): send buffered text in smooth stream when stream parts change