Simplify state in both stream managers

Author: ammario

src/services/streamManager.ts

@@ -300,8 +300,6 @@ export class StreamManager extends EventEmitter {
         historySequence,
       } as StreamStartEvent);
 
-      let currentTextBuffer = "";
-
       // Use fullStream to capture all events including tool calls
       const toolCalls = new Map<
         string,
@@ -323,31 +321,18 @@ export class StreamManager extends EventEmitter {
 
         switch (part.type) {
           case "text-delta":
-            currentTextBuffer += part.text;
             this.emit("stream-delta", {
               type: "stream-delta",
               workspaceId: workspaceId as string,
               messageId: streamInfo.messageId,
               delta: part.text,
             } as StreamDeltaEvent);
 
-            // Update parts array for partial message
-            if (
-              streamInfo.parts.length > 0 &&
-              streamInfo.parts[streamInfo.parts.length - 1].type === "text"
-            ) {
-              // Update existing text part
-              streamInfo.parts[streamInfo.parts.length - 1] = {
-                type: "text",
-                text: currentTextBuffer,
-              };
-            } else {
-              // Add new text part
-              streamInfo.parts.push({
-                type: "text",
-                text: currentTextBuffer,
-              });
-            }
+            // Append each delta as a new part (merging happens at display time)
+            streamInfo.parts.push({
+              type: "text",
+              text: part.text,
+            });
 
             // Schedule partial write (throttled, fire-and-forget to not block stream)
             void this.schedulePartialWrite(workspaceId, streamInfo);
@@ -356,21 +341,11 @@ export class StreamManager extends EventEmitter {
           case "reasoning-delta": {
             const delta = (part as ReasoningDeltaPart).text ?? "";
 
-            // Check if last part is reasoning (consistent with text-delta handling)
-            const lastPart = streamInfo.parts[streamInfo.parts.length - 1];
-            if (lastPart?.type === "reasoning") {
-              // Update existing reasoning part
-              streamInfo.parts[streamInfo.parts.length - 1] = {
-                type: "reasoning",
-                text: lastPart.text + delta,
-              };
-            } else {
-              // Push new reasoning part
-              streamInfo.parts.push({
-                type: "reasoning",
-                text: delta,
-              });
-            }
+            // Append each delta as a new part (merging happens at display time)
+            streamInfo.parts.push({
+              type: "reasoning",
+              text: delta,
+            });
 
             this.emit("reasoning-delta", {
               type: "reasoning-delta",
@@ -393,10 +368,6 @@ export class StreamManager extends EventEmitter {
           }
 
           case "tool-call": {
-            // Reset text buffer to separate text segments (text is already in parts from text-delta handler)
-            // We don't push here because text-delta already maintains parts array
-            currentTextBuffer = "";
-
             // Tool call started - store in map for later lookup
             toolCalls.set(part.toolCallId, {
               toolCallId: part.toolCallId,

src/utils/StreamingMessageAggregator.test.ts

@@ -171,8 +171,8 @@ describe("StreamingMessageAggregator", () => {
     expect(messages).toHaveLength(1);
 
     const message = messages[0];
-    // Should have 3 parts: text, tool, text
-    expect(message.parts).toHaveLength(3);
+    // Should have 4 parts: text, tool, text, text (deltas not merged during streaming)
+    expect(message.parts).toHaveLength(4);
 
     // First text part (before tool)
     expect(message.parts[0].type).toBe("text");
@@ -186,10 +186,12 @@ describe("StreamingMessageAggregator", () => {
     expect(toolPart.toolName).toBe("searchFiles");
     expect(toolPart.state).toBe("output-available");
 
-    // Second text part (after tool)
+    // Second and third text parts (after tool) - separate deltas not yet merged
     expect(message.parts[2].type).toBe("text");
-    if (message.parts[2].type === "text") {
-      expect(message.parts[2].text).toBe("I found the following results: file1.ts and file2.ts");
+    expect(message.parts[3].type).toBe("text");
+    if (message.parts[2].type === "text" && message.parts[3].type === "text") {
+      expect(message.parts[2].text).toBe("I found the following results: ");
+      expect(message.parts[3].text).toBe("file1.ts and file2.ts");
     }
 
     // Test DisplayedMessages split
@@ -348,9 +350,18 @@ describe("StreamingMessageAggregator", () => {
     const messages = aggregator.getAllMessages();
     expect(messages).toHaveLength(1);
 
+    // Raw parts are separate deltas (2 parts: "Hello, " and "world!")
+    expect(messages[0].parts).toHaveLength(2);
     const firstPart = messages[0].parts[0];
     if (firstPart.type === "text") {
-      expect(firstPart.text).toBe("Hello, world!");
+      expect(firstPart.text).toBe("Hello, ");
+    }
+
+    // DisplayedMessages should merge them
+    const displayedMessages = aggregator.getDisplayedMessages();
+    expect(displayedMessages).toHaveLength(1);
+    if (displayedMessages[0].type === "assistant") {
+      expect(displayedMessages[0].content).toBe("Hello, world!");
     }
   });
 

src/utils/StreamingMessageAggregator.ts

@@ -111,15 +111,13 @@ export class StreamingMessageAggregator {
 
     this.activeStreams.set(context.streamingId, context);
 
-    // Create initial streaming message
+    // Create initial streaming message with empty parts (deltas will append)
     const streamingMessage = createCmuxMessage(data.messageId, "assistant", "", {
       historySequence: data.historySequence,
       streamingId: context.streamingId,
       timestamp: Date.now(),
       model: data.model,
     });
-    // Start with empty text part (streaming status inferred from activeStreams)
-    streamingMessage.parts[0] = { type: "text", text: "" };
 
     this.messages.set(data.messageId, streamingMessage);
     this.incrementDisplayVersion();
@@ -129,21 +127,11 @@ export class StreamingMessageAggregator {
     const message = this.messages.get(data.messageId);
     if (!message) return;
 
-    // Check if last part is text (consistent with text handling in streamManager)
-    const lastPart = message.parts[message.parts.length - 1];
-    if (lastPart?.type === "text") {
-      // Update existing text part
-      message.parts[message.parts.length - 1] = {
-        type: "text",
-        text: lastPart.text + data.delta,
-      };
-    } else {
-      // Push new text part
-      message.parts.push({
-        type: "text",
-        text: data.delta,
-      });
-    }
+    // Append each delta as a new part (merging happens at display time)
+    message.parts.push({
+      type: "text",
+      text: data.delta,
+    });
     this.incrementDisplayVersion();
   }
 
@@ -279,21 +267,11 @@ export class StreamingMessageAggregator {
     const message = this.messages.get(data.messageId);
     if (!message) return;
 
-    // Check if last part is reasoning (consistent with text handling)
-    const lastPart = message.parts[message.parts.length - 1];
-    if (lastPart?.type === "reasoning") {
-      // Update existing reasoning part
-      message.parts[message.parts.length - 1] = {
-        type: "reasoning",
-        text: lastPart.text + data.delta,
-      };
-    } else {
-      // Push new reasoning part
-      message.parts.push({
-        type: "reasoning",
-        text: data.delta,
-      });
-    }
+    // Append each delta as a new part (merging happens at display time)
+    message.parts.push({
+      type: "reasoning",
+      text: data.delta,
+    });
     this.incrementDisplayVersion();
   }
 
@@ -376,11 +354,37 @@ export class StreamingMessageAggregator {
         // Check if this message has an active stream (for inferring streaming status)
         const hasActiveStream = this.getActiveStreams().some((s) => s.messageId === message.id);
 
+        // Merge adjacent parts of same type (text with text, reasoning with reasoning)
+        // This is where all merging happens - streaming just appends raw deltas
+        const mergedParts: typeof message.parts = [];
+        for (let i = 0; i < message.parts.length; i++) {
+          const part = message.parts[i];
+          const lastMerged = mergedParts[mergedParts.length - 1];
+
+          // Try to merge with last part if same type
+          if (lastMerged?.type === "text" && part.type === "text") {
+            // Merge text parts
+            mergedParts[mergedParts.length - 1] = {
+              type: "text",
+              text: lastMerged.text + part.text,
+            };
+          } else if (lastMerged?.type === "reasoning" && part.type === "reasoning") {
+            // Merge reasoning parts
+            mergedParts[mergedParts.length - 1] = {
+              type: "reasoning",
+              text: lastMerged.text + part.text,
+            };
+          } else {
+            // Different type or tool part - add new part
+            mergedParts.push(part);
+          }
+        }
+
         // Find the last part that will produce a DisplayedMessage
         // (reasoning, text parts with content, OR tool parts)
         let lastPartIndex = -1;
-        for (let i = message.parts.length - 1; i >= 0; i--) {
-          const part = message.parts[i];
+        for (let i = mergedParts.length - 1; i >= 0; i--) {
+          const part = mergedParts[i];
           if (
             part.type === "reasoning" ||
             (part.type === "text" && part.text) ||
@@ -391,7 +395,7 @@ export class StreamingMessageAggregator {
           }
         }
 
-        message.parts.forEach((part, partIndex) => {
+        mergedParts.forEach((part, partIndex) => {
           const isLastPart = partIndex === lastPartIndex;
           // Part is streaming if: active stream exists AND this is the last part
           const isStreaming = hasActiveStream && isLastPart;