🤖 refactor: simplify stream lifecycle plumbing

Author: ammar-agent

src/browser/utils/messages/StreamingMessageAggregator.ts

@@ -548,18 +548,6 @@ export class StreamingMessageAggregator {
     this.setPendingStreamStartTime(null);
 
     this.connectingStreams.set(data.messageId, { startTime: Date.now(), model: data.model });
-
-    // Create a placeholder assistant message (kept invisible until parts arrive)
-    // so that out-of-order deltas (if they ever occur) have somewhere to attach.
-    if (!this.messages.has(data.messageId)) {
-      const connectingMessage = createMuxMessage(data.messageId, "assistant", "", {
-        historySequence: data.historySequence,
-        timestamp: Date.now(),
-        model: data.model,
-      });
-      this.messages.set(data.messageId, connectingMessage);
-    }
-
     this.invalidateCache();
   }
 

src/common/utils/streamLifecycle.ts

@@ -0,0 +1,42 @@
+// Stream lifecycle events are emitted during an in-flight assistant response.
+//
+// Keeping the event list centralized makes it harder to accidentally forget to forward/buffer a
+// newly introduced lifecycle event.
+
+export const STREAM_LIFECYCLE_EVENTS = [
+  "stream-pending",
+  "stream-start",
+  "stream-delta",
+  "stream-abort",
+  "stream-end",
+] as const;
+
+export type StreamLifecycleEventName = (typeof STREAM_LIFECYCLE_EVENTS)[number];
+
+// Events that can be forwarded 1:1 from StreamManager -> AIService.
+// (`stream-abort` needs additional bookkeeping in AIService.)
+export const STREAM_LIFECYCLE_EVENTS_DIRECT_FORWARD = [
+  "stream-pending",
+  "stream-start",
+  "stream-delta",
+  "stream-end",
+] as const satisfies readonly StreamLifecycleEventName[];
+
+// Events that can be forwarded 1:1 from AIService -> AgentSession -> renderer.
+// (`stream-end` has additional session-side behavior.)
+export const STREAM_LIFECYCLE_EVENTS_SIMPLE_FORWARD = [
+  "stream-pending",
+  "stream-start",
+  "stream-delta",
+  "stream-abort",
+] as const satisfies readonly StreamLifecycleEventName[];
+
+export function forwardStreamLifecycleEvents(params: {
+  events: readonly StreamLifecycleEventName[];
+  listen: (event: StreamLifecycleEventName, handler: (payload: unknown) => void) => void;
+  emit: (event: StreamLifecycleEventName, payload: unknown) => void;
+}): void {
+  for (const event of params.events) {
+    params.listen(event, (payload) => params.emit(event, payload));
+  }
+}

src/node/services/agentSession.ts

@@ -35,6 +35,10 @@ import { AttachmentService } from "./attachmentService";
 import type { PostCompactionAttachment, PostCompactionExclusions } from "@/common/types/attachment";
 import { TURNS_BETWEEN_ATTACHMENTS } from "@/common/constants/attachments";
 import { extractEditedFileDiffs } from "@/common/utils/messages/extractEditedFiles";
+import {
+  forwardStreamLifecycleEvents,
+  STREAM_LIFECYCLE_EVENTS_SIMPLE_FORWARD,
+} from "@/common/utils/streamLifecycle";
 import { isValidModelFormat } from "@/common/utils/ai/models";
 
 /**
@@ -619,9 +623,15 @@ export class AgentSession {
       this.aiService.on(event, wrapped as never);
     };
 
-    forward("stream-pending", (payload) => this.emitChatEvent(payload));
-    forward("stream-start", (payload) => this.emitChatEvent(payload));
-    forward("stream-delta", (payload) => this.emitChatEvent(payload));
+    forwardStreamLifecycleEvents({
+      events: STREAM_LIFECYCLE_EVENTS_SIMPLE_FORWARD,
+      listen: (event, handler) => {
+        forward(event, handler);
+      },
+      emit: (_event, payload) => {
+        this.emitChatEvent(payload as WorkspaceChatMessage);
+      },
+    });
     forward("tool-call-start", (payload) => this.emitChatEvent(payload));
     forward("tool-call-delta", (payload) => this.emitChatEvent(payload));
     forward("tool-call-end", (payload) => {
@@ -644,7 +654,6 @@ export class AgentSession {
     forward("reasoning-delta", (payload) => this.emitChatEvent(payload));
     forward("reasoning-end", (payload) => this.emitChatEvent(payload));
     forward("usage-delta", (payload) => this.emitChatEvent(payload));
-    forward("stream-abort", (payload) => this.emitChatEvent(payload));
 
     forward("stream-end", async (payload) => {
       const handled = await this.compactionHandler.handleCompletion(payload as StreamEndEvent);

src/node/services/aiService.ts

@@ -58,6 +58,10 @@ import { applyToolPolicy, type ToolPolicy } from "@/common/utils/tools/toolPolic
 import { MockScenarioPlayer } from "./mock/mockScenarioPlayer";
 import { EnvHttpProxyAgent, type Dispatcher } from "undici";
 import { getPlanFilePath } from "@/common/utils/planStorage";
+import {
+  forwardStreamLifecycleEvents,
+  STREAM_LIFECYCLE_EVENTS_DIRECT_FORWARD,
+} from "@/common/utils/streamLifecycle";
 import { getPlanModeInstruction } from "@/common/utils/ui/modeUtils";
 import type { UIMode } from "@/common/types/mode";
 import { MUX_APP_ATTRIBUTION_TITLE, MUX_APP_ATTRIBUTION_URL } from "@/constants/appAttribution";
@@ -348,10 +352,15 @@ export class AIService extends EventEmitter {
    * Forward all stream events from StreamManager to AIService consumers
    */
   private setupStreamEventForwarding(): void {
-    this.streamManager.on("stream-pending", (data) => this.emit("stream-pending", data));
-    this.streamManager.on("stream-start", (data) => this.emit("stream-start", data));
-    this.streamManager.on("stream-delta", (data) => this.emit("stream-delta", data));
-    this.streamManager.on("stream-end", (data) => this.emit("stream-end", data));
+    forwardStreamLifecycleEvents({
+      events: STREAM_LIFECYCLE_EVENTS_DIRECT_FORWARD,
+      listen: (event, handler) => {
+        this.streamManager.on(event, handler);
+      },
+      emit: (event, payload) => {
+        this.emit(event, payload);
+      },
+    });
 
     // Handle stream-abort: dispose of partial based on abandonPartial flag
     this.streamManager.on("stream-abort", (data: StreamAbortEvent) => {
@@ -1405,6 +1414,7 @@ export class AIService extends EventEmitter {
       // Delegate to StreamManager with model instance, system message, tools, historySequence, and initial metadata
       const streamResult = await this.streamManager.startStream(
         workspaceId,
+        assistantMessageId,
         finalMessages,
         modelResult.data,
         modelString,

src/node/services/streamManager.test.ts

@@ -84,6 +84,7 @@ describe("StreamManager - Concurrent Stream Prevention", () => {
       // Start first stream
       const result1 = await streamManager.startStream(
         workspaceId,
+        "assistant-1",
         [{ role: "user", content: "Say hello and nothing else" }],
         model,
         KNOWN_MODELS.SONNET.id,
@@ -102,6 +103,7 @@ describe("StreamManager - Concurrent Stream Prevention", () => {
       // Start second stream - should cancel first
       const result2 = await streamManager.startStream(
         workspaceId,
+        "assistant-2",
         [{ role: "user", content: "Say goodbye and nothing else" }],
         model,
         KNOWN_MODELS.SONNET.id,
@@ -164,6 +166,11 @@ describe("StreamManager - Concurrent Stream Prevention", () => {
     }
 
     const ensureStreamSafetyValue = Reflect.get(streamManager, "ensureStreamSafety") as unknown;
+    const pendingMessageIds: string[] = [];
+    streamManager.on("stream-pending", (event) => {
+      pendingMessageIds.push((event as { messageId: string }).messageId);
+    });
+
     if (typeof ensureStreamSafetyValue !== "function") {
       throw new Error("StreamManager.ensureStreamSafety is unavailable for testing");
     }
@@ -199,7 +206,10 @@ describe("StreamManager - Concurrent Stream Prevention", () => {
       "createStreamAtomically",
       (
         wsId: string,
+        assistantMessageId: string,
         streamToken: string,
+        _runtimeTempDir: string,
+        _runtime: unknown,
         messages: unknown,
         modelArg: unknown,
         modelString: string,
@@ -228,7 +238,7 @@ describe("StreamManager - Concurrent Stream Prevention", () => {
             providerMetadata: Promise.resolve(undefined),
           },
           abortController,
-          messageId: `test-${Math.random().toString(36).slice(2)}`,
+          messageId: assistantMessageId,
           token: streamToken,
           startTime: Date.now(),
           model: modelString,
@@ -274,6 +284,7 @@ describe("StreamManager - Concurrent Stream Prevention", () => {
     const promises = [
       streamManager.startStream(
         workspaceId,
+        "assistant-mutex-1",
         [{ role: "user", content: "test 1" }],
         model,
         KNOWN_MODELS.SONNET.id,
@@ -285,6 +296,7 @@ describe("StreamManager - Concurrent Stream Prevention", () => {
       ),
       streamManager.startStream(
         workspaceId,
+        "assistant-mutex-2",
         [{ role: "user", content: "test 2" }],
         model,
         KNOWN_MODELS.SONNET.id,
@@ -296,6 +308,7 @@ describe("StreamManager - Concurrent Stream Prevention", () => {
       ),
       streamManager.startStream(
         workspaceId,
+        "assistant-mutex-3",
         [{ role: "user", content: "test 3" }],
         model,
         KNOWN_MODELS.SONNET.id,
@@ -317,6 +330,12 @@ describe("StreamManager - Concurrent Stream Prevention", () => {
       expect(ensureOperations[i]).toBe("ensure-start");
       expect(ensureOperations[i + 1]).toBe("ensure-end");
     }
+
+    expect(pendingMessageIds).toEqual([
+      "assistant-mutex-1",
+      "assistant-mutex-2",
+      "assistant-mutex-3",
+    ]);
   });
 });
 

src/node/services/streamManager.ts

@@ -600,6 +600,7 @@ export class StreamManager extends EventEmitter {
    */
   private createStreamAtomically(
     workspaceId: WorkspaceId,
+    assistantMessageId: string,
     streamToken: StreamToken,
     runtimeTempDir: string,
     runtime: Runtime,
@@ -684,7 +685,7 @@ export class StreamManager extends EventEmitter {
       throw error;
     }
 
-    const messageId = `assistant-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
+    const messageId = assistantMessageId;
 
     const streamInfo: WorkspaceStreamInfo = {
       state: StreamState.STARTING,
@@ -1380,6 +1381,7 @@ export class StreamManager extends EventEmitter {
    */
   async startStream(
     workspaceId: string,
+    assistantMessageId: string,
     messages: ModelMessage[],
     model: LanguageModel,
     modelString: string,
@@ -1426,6 +1428,7 @@ export class StreamManager extends EventEmitter {
       // Step 4: Atomic stream creation and registration
       const streamInfo = this.createStreamAtomically(
         typedWorkspaceId,
+        assistantMessageId,
         streamToken,
         runtimeTempDir,
         runtime,