Modify smoothStream transform to support tool call streams

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

@@ -43,33 +43,114 @@ export function smoothStream<TOOLS extends ToolSet>({
   }
 
   return () => {
-    let buffer = '';
+    // Separate buffers for text and tool call args
+    let textBuffer = '';
+    let toolCallBuffer: { id: string; name: string; buffer: string } | null =
+      null;
+
     return new TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>({
       async transform(chunk, controller) {
-        if (chunk.type === 'step-finish') {
-          if (buffer.length > 0) {
-            controller.enqueue({ type: 'text-delta', textDelta: buffer });
-            buffer = '';
+        // Handle special cases that should flush buffers and pass through
+        if (
+          chunk.type === 'step-finish' ||
+          chunk.type === 'tool-result' ||
+          chunk.type === 'tool-call'
+        ) {
+          // Flush text buffer if any
+          if (textBuffer.length > 0) {
+            controller.enqueue({ type: 'text-delta', textDelta: textBuffer });
+            textBuffer = '';
+          }
+
+          // For tool-call states (call/result), clear the buffer without sending an extra delta
+          // This prevents an extra partial-call from appearing between call and result states
+          if (chunk.type === 'tool-call') {
+            // Just clear the buffer without sending anything
+            toolCallBuffer = null;
+          }
+          // For other cases, flush the buffer normally
+          else if (toolCallBuffer && toolCallBuffer.buffer.length > 0) {
+            controller.enqueue({
+              type: 'tool-call-delta',
+              toolCallId: toolCallBuffer.id,
+              toolName: toolCallBuffer.name,
+              argsTextDelta: toolCallBuffer.buffer,
+            });
+            toolCallBuffer = null;
           }
 
           controller.enqueue(chunk);
           return;
         }
 
-        if (chunk.type !== 'text-delta') {
+        // Pass through any non-streaming chunks
+        if (chunk.type !== 'text-delta' && chunk.type !== 'tool-call-delta') {
           controller.enqueue(chunk);
           return;
         }
 
-        buffer += chunk.textDelta;
+        // Handle text-delta chunks
+        if (chunk.type === 'text-delta') {
+          textBuffer += chunk.textDelta;
+
+          let match;
+          while ((match = chunkingRegexp.exec(textBuffer)) != null) {
+            const chunkText = match[0];
+            controller.enqueue({ type: 'text-delta', textDelta: chunkText });
+            textBuffer = textBuffer.slice(chunkText.length);
+
+            await delay(delayInMs);
+          }
+          return;
+        }
+
+        // Handle tool-call-delta chunks
+        if (chunk.type === 'tool-call-delta') {
+          // Initialize tool call buffer if needed
+          if (!toolCallBuffer || toolCallBuffer.id !== chunk.toolCallId) {
+            toolCallBuffer = {
+              id: chunk.toolCallId,
+              name: chunk.toolName,
+              buffer: '',
+            };
+          }
+
+          // Add new content to buffer
+          toolCallBuffer.buffer += chunk.argsTextDelta;
 
-        let match;
-        while ((match = chunkingRegexp.exec(buffer)) != null) {
-          const chunk = match[0];
-          controller.enqueue({ type: 'text-delta', textDelta: chunk });
-          buffer = buffer.slice(chunk.length);
+          // Process buffer chunks
+          let match;
+          while ((match = chunkingRegexp.exec(toolCallBuffer.buffer)) != null) {
+            const chunkText = match[0];
+            controller.enqueue({
+              type: 'tool-call-delta',
+              toolCallId: toolCallBuffer.id,
+              toolName: toolCallBuffer.name,
+              argsTextDelta: chunkText,
+            });
+            toolCallBuffer.buffer = toolCallBuffer.buffer.slice(
+              chunkText.length,
+            );
+
+            await delay(delayInMs);
+          }
+          return;
+        }
+      },
+      flush(controller) {
+        // Flush any remaining content in the text buffer
+        if (textBuffer.length > 0) {
+          controller.enqueue({ type: 'text-delta', textDelta: textBuffer });
+        }
 
-          await delay(delayInMs);
+        // Flush any remaining content in the tool call buffer
+        if (toolCallBuffer && toolCallBuffer.buffer.length > 0) {
+          controller.enqueue({
+            type: 'tool-call-delta',
+            toolCallId: toolCallBuffer.id,
+            toolName: toolCallBuffer.name,
+            argsTextDelta: toolCallBuffer.buffer,
+          });
         }
       },
     });