Reapply "fix: isolate chat SSE broadcasts per generation"

This reverts commit 46f1f042b0.

.changeset/fix-chat-cancel-leak-into-new-stream.md

@@ -0,0 +1,5 @@
+---
+"@runfusion/fusion": patch
+---
+
+Fix chat: after stopping a streaming reply, the next message would appear sent but show no Stop button or "Connecting…" indicator. The cancellation broadcast from the previous generation was leaking into the new SSE subscription, immediately marking it as errored. Each `chatManager.sendMessage` now allocates a per-generation id; `ChatStreamManager` only delivers tagged broadcasts to subscribers from the matching generation, and `sendMessage`'s cleanup no longer deletes a newer generation's `activeGenerations` slot when an older one finally unwinds.

packages/dashboard/src/__tests__/chat-attachment-routes.test.ts

@@ -17,7 +17,7 @@ const mockGetMessages = vi.fn();
 const mockGetMessage = vi.fn();
 const mockGetLastMessageForSessions = vi.fn().mockReturnValue(new Map());
 const mockDeleteMessage = vi.fn();
-const { mockChatStreamManager, mockSendMessage, mockCancelGeneration } = vi.hoisted(() => {
+const { mockChatStreamManager, mockSendMessage, mockCancelGeneration, mockBeginGeneration } = vi.hoisted(() => {
   const subscribers = new Map<string, Set<(event: any, eventId?: number) => void>>();
   const chatStreamManager = {
     subscribe: vi.fn((sessionId: string, callback: (event: any, eventId?: number) => void) => {
@@ -39,6 +39,7 @@ const { mockChatStreamManager, mockSendMessage, mockCancelGeneration } = vi.hois
       chatStreamManager.broadcast(sessionId, { type: "done", data: { messageId: "msg-1" } });
     }),
     mockCancelGeneration: vi.fn().mockReturnValue(false),
+    mockBeginGeneration: vi.fn(() => ({ generationId: 1, abortController: new AbortController() })),
   };
 });
 
@@ -57,7 +58,7 @@ const mockGetOrCreateProjectStore = vi.fn();
 vi.mock("../project-store-resolver.js", () => ({ getOrCreateProjectStore: mockGetOrCreateProjectStore, invalidateAllGlobalSettingsCaches: vi.fn() }));
 
 vi.mock("../chat.js", () => ({
-  ChatManager: class MockChatManager { sendMessage = mockSendMessage; cancelGeneration = mockCancelGeneration; },
+  ChatManager: class MockChatManager { sendMessage = mockSendMessage; cancelGeneration = mockCancelGeneration; beginGeneration = mockBeginGeneration; },
   chatStreamManager: mockChatStreamManager,
   checkRateLimit: vi.fn().mockReturnValue(true),
   getRateLimitResetTime: vi.fn().mockReturnValue(null),
@@ -125,7 +126,7 @@ describe("chat attachment routes", () => {
     const { createServer } = await import("../server.js");
     app = createServer(store as any, { chatStore: {
       init: mockInit, createSession: mockCreateSession, getSession: mockGetSession, listSessions: mockListSessions, updateSession: mockUpdateSession, deleteSession: mockDeleteSession, addMessage: mockAddMessage, getMessages: mockGetMessages, getMessage: mockGetMessage, getLastMessageForSessions: mockGetLastMessageForSessions, deleteMessage: mockDeleteMessage,
-    } as any, chatManager: { sendMessage: mockSendMessage, cancelGeneration: mockCancelGeneration } as any });
+    } as any, chatManager: { sendMessage: mockSendMessage, cancelGeneration: mockCancelGeneration, beginGeneration: mockBeginGeneration } as any });
   });
 
   it("uploads a valid attachment", async () => {
@@ -187,7 +188,7 @@ describe("chat attachment routes", () => {
     const body = JSON.stringify({ content: "hello", attachments });
     const response = await request(app, "POST", `/api/chat/sessions/${session.id}/messages`, body, { "content-type": "application/json" });
     expect(response.status).toBe(200);
-    expect(mockSendMessage).toHaveBeenCalledWith(session.id, "hello", undefined, undefined, attachments);
+    expect(mockSendMessage).toHaveBeenCalledWith(session.id, "hello", undefined, undefined, attachments, { generationId: 1 });
   });
 
   afterEach(() => {

packages/dashboard/src/__tests__/chat-manager.test.ts

@@ -1661,3 +1661,124 @@ describe("ChatManager.getGeneratingSessionIds", () => {
     expect(chatManager.getGeneratingSessionIds()).toEqual([]);
   });
 });
+
+describe("ChatManager generation isolation", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    __resetChatState();
+    mockChatStore.getSession.mockReturnValue({
+      id: "chat-001",
+      agentId: "agent-001",
+      status: "active",
+    });
+    mockChatStore.addMessage.mockReturnValue({
+      id: "msg-001",
+      sessionId: "chat-001",
+      role: "assistant",
+      content: "",
+    });
+  });
+
+  // Regression: a "Generation cancelled" broadcast from a previous generation
+  // must not leak into a new SSE subscriber that connected after the cancel
+  // request landed. Before per-generation tagging this race silently flipped
+  // the new chat into "no streaming" state with no Stop button or indicator.
+  it("cancelGeneration broadcast does not reach a new generation's subscriber", async () => {
+    const chatManager = createChatManager();
+
+    // Start gen #1 manually so we can hold a reference to its generationId
+    // without driving a real agent loop.
+    const firstGen = chatManager.beginGeneration("chat-001");
+    expect(firstGen.generationId).toBe(1);
+
+    // Simulate the new request subscribing for gen #2 BEFORE the cancel of
+    // gen #1 has been processed by the backend.
+    const secondGen = chatManager.beginGeneration("chat-001");
+    expect(secondGen.generationId).toBe(2);
+    // beginGeneration aborts the prior controller so the old loop unwinds.
+    expect(firstGen.abortController.signal.aborted).toBe(true);
+
+    const eventsForNewSubscriber: Array<{ type: string; data: unknown }> = [];
+    const unsubscribe = chatStreamManager.subscribe(
+      "chat-001",
+      (event) => { eventsForNewSubscriber.push(event); },
+      { generationId: secondGen.generationId },
+    );
+
+    // Cancel gen #1 (which is what the in-flight HTTP cancel request would do).
+    // Today this is a no-op for activeGenerations because beginGeneration #2
+    // overwrote the entry, but in production the cancel can race with the
+    // beginGeneration. Simulate the broadcast directly as well to exercise the
+    // tagged-broadcast filtering.
+    chatStreamManager.broadcast(
+      "chat-001",
+      { type: "error", data: "Generation cancelled" },
+      { generationId: firstGen.generationId },
+    );
+
+    expect(eventsForNewSubscriber).toEqual([]);
+
+    // A broadcast tagged for gen #2 still reaches the subscriber.
+    chatStreamManager.broadcast(
+      "chat-001",
+      { type: "text", data: "hello" },
+      { generationId: secondGen.generationId },
+    );
+    expect(eventsForNewSubscriber).toEqual([{ type: "text", data: "hello" }]);
+
+    unsubscribe();
+  });
+
+  // Regression: an old generation completing its `finally` block must not
+  // delete a newer generation's activeGenerations entry, otherwise
+  // `isGenerating(sessionId)` returns false while the new request is still
+  // streaming and recovery polling cannot find it.
+  it("old generation finally does not delete a newer generation's slot", async () => {
+    const chatManager = createChatManager();
+
+    let resolvePrompt: (() => void) | undefined;
+    let promptCallCount = 0;
+    __setCreateFnAgent(async () => {
+      promptCallCount += 1;
+      const callIndex = promptCallCount;
+      return {
+        session: {
+          prompt: vi.fn().mockImplementation(() => {
+            // First call hangs until we resolve it; second resolves immediately.
+            if (callIndex === 1) {
+              return new Promise<void>((resolve) => { resolvePrompt = resolve; });
+            }
+            return Promise.resolve();
+          }),
+          dispose: vi.fn(),
+          state: { messages: [{ role: "assistant", content: "ok" }] },
+        },
+      };
+    });
+
+    // Kick off generation #1 — it will hang inside prompt().
+    const sendOne = chatManager.sendMessage("chat-001", "first");
+    // Yield enough microtasks for sendOne to set its activeGenerations entry.
+    await Promise.resolve();
+    await Promise.resolve();
+    await Promise.resolve();
+    expect(chatManager.isGenerating("chat-001")).toBe(true);
+
+    // Cancel #1 (so its prompt() will eventually unwind via the abort path).
+    chatManager.cancelGeneration("chat-001");
+
+    // Start generation #2 and let it run. Its sendMessage promise resolves
+    // synchronously (mock prompt returns immediately on the second call), but
+    // the activeGenerations slot for #2 is set during its execution.
+    const sendTwo = chatManager.sendMessage("chat-001", "second");
+    await sendTwo;
+
+    // Now release #1's prompt so its finally block runs.
+    resolvePrompt?.();
+    await sendOne;
+
+    // The slot was correctly cleaned up by sendTwo (the most recent owner)
+    // and not re-deleted/corrupted by sendOne's late finally.
+    expect(chatManager.isGenerating("chat-001")).toBe(false);
+  });
+});

packages/dashboard/src/__tests__/chat-routes.test.ts

@@ -48,7 +48,7 @@ function createSSEResponse(): {
 const mockInit = vi.fn().mockResolvedValue(undefined);
 
 // Create mock functions before vi.mock
-const { mockCreateFnAgent, mockChatStreamManager, mockSendMessage, mockCancelGeneration } = vi.hoisted(() => {
+const { mockCreateFnAgent, mockChatStreamManager, mockSendMessage, mockCancelGeneration, mockBeginGeneration } = vi.hoisted(() => {
   // Store subscribers per session for broadcast simulation
   const subscribers = new Map<string, Set<(event: any, eventId?: number) => void>>();
 
@@ -107,6 +107,7 @@ const { mockCreateFnAgent, mockChatStreamManager, mockSendMessage, mockCancelGen
     mockCreateFnAgent: vi.fn(),
     mockSendMessage: vi.fn(),
     mockCancelGeneration: vi.fn(),
+    mockBeginGeneration: vi.fn(() => ({ generationId: 1, abortController: new AbortController() })),
     mockChatStreamManager: chatStreamManager,
   };
 });
@@ -162,6 +163,7 @@ vi.mock("../chat.js", () => {
     ChatManager: class MockChatManager {
       sendMessage = mockSendMessage;
       cancelGeneration = mockCancelGeneration;
+      beginGeneration = mockBeginGeneration;
     },
     chatStreamManager: mockChatStreamManager,
     checkRateLimit: vi.fn().mockReturnValue(true),
@@ -263,6 +265,7 @@ function createMockChatManager() {
   return {
     sendMessage: mockSendMessage,
     cancelGeneration: mockCancelGeneration,
+    beginGeneration: mockBeginGeneration,
   };
 }
 

packages/dashboard/src/chat.ts

@@ -168,6 +168,12 @@ export type ChatStreamEvent =
 /** Callback function for streaming events */
 export type ChatStreamCallback = (event: ChatStreamEvent, eventId?: number) => void;
 
+/** Per-subscription record. `generationId` (if set) filters which broadcasts are delivered. */
+interface ChatStreamSubscription {
+  callback: ChatStreamCallback;
+  generationId?: number;
+}
+
 interface RateLimitEntry {
   count: number;
   firstRequestAt: Date;
@@ -285,7 +291,7 @@ export async function resolveFileReferences(content: string, rootDir: string): P
  * Follows the PlanningStreamManager pattern.
  */
 export class ChatStreamManager extends EventEmitter {
-  private readonly sessions = new Map<string, Set<ChatStreamCallback>>();
+  private readonly sessions = new Map<string, Set<ChatStreamSubscription>>();
   private readonly buffers = new Map<string, SessionEventBuffer>();
 
   constructor(private readonly bufferSize = 100) {
@@ -295,18 +301,29 @@ export class ChatStreamManager extends EventEmitter {
   /**
    * Register a client callback for a chat session.
    * Returns a function to unsubscribe.
+   *
+   * If `options.generationId` is provided, this subscriber only receives broadcasts
+   * tagged with the same generationId (or untagged broadcasts). This isolates each
+   * client SSE connection to events from its own `chatManager.sendMessage` call so
+   * that a previous generation's late "Generation cancelled" event cannot leak into
+   * a new request that has just subscribed for the same session.
    */
-  subscribe(sessionId: string, callback: ChatStreamCallback): () => void {
+  subscribe(
+    sessionId: string,
+    callback: ChatStreamCallback,
+    options?: { generationId?: number },
+  ): () => void {
     if (!this.sessions.has(sessionId)) {
       this.sessions.set(sessionId, new Set());
     }
 
-    const callbacks = this.sessions.get(sessionId)!;
-    callbacks.add(callback);
+    const subscriptions = this.sessions.get(sessionId)!;
+    const subscription: ChatStreamSubscription = { callback, generationId: options?.generationId };
+    subscriptions.add(subscription);
 
     return () => {
-      callbacks.delete(callback);
-      if (callbacks.size === 0) {
+      subscriptions.delete(subscription);
+      if (subscriptions.size === 0) {
         this.sessions.delete(sessionId);
       }
     };
@@ -324,18 +341,36 @@ export class ChatStreamManager extends EventEmitter {
   /**
    * Broadcast an event to all clients subscribed to a session.
    * Every event is buffered and assigned a monotonically increasing id.
+   *
+   * When `options.generationId` is set, the event is delivered only to subscribers
+   * that registered without a generation filter or whose generation matches.
+   * Subscribers tied to a different generation will not receive it. Untagged
+   * broadcasts (no generationId) reach every subscriber for backward compatibility.
    */
-  broadcast(sessionId: string, event: ChatStreamEvent): number {
+  broadcast(
+    sessionId: string,
+    event: ChatStreamEvent,
+    options?: { generationId?: number },
+  ): number {
     const serialized = JSON.stringify((event as { data?: unknown }).data ?? {});
     const eventData = typeof serialized === "string" ? serialized : "{}";
     const eventId = this.getBuffer(sessionId).push(event.type, eventData);
 
-    const callbacks = this.sessions.get(sessionId);
-    if (!callbacks) return eventId;
+    const subscriptions = this.sessions.get(sessionId);
+    if (!subscriptions) return eventId;
 
-    for (const callback of callbacks) {
+    const broadcastGenerationId = options?.generationId;
+
+    for (const subscription of subscriptions) {
+      if (
+        broadcastGenerationId !== undefined &&
+        subscription.generationId !== undefined &&
+        subscription.generationId !== broadcastGenerationId
+      ) {
+        continue;
+      }
       try {
-        callback(event, eventId);
+        subscription.callback(event, eventId);
       } catch (err) {
         diagnostics.error(`Error broadcasting to client for session ${sessionId}:`, err);
       }
@@ -357,8 +392,8 @@ export class ChatStreamManager extends EventEmitter {
    * Check if a session has active subscribers.
    */
   hasSubscribers(sessionId: string): boolean {
-    const callbacks = this.sessions.get(sessionId);
-    return callbacks !== undefined && callbacks.size > 0;
+    const subscriptions = this.sessions.get(sessionId);
+    return subscriptions !== undefined && subscriptions.size > 0;
   }
 
   /**
@@ -447,9 +482,11 @@ export function getRateLimitResetTime(ip: string): Date | null {
  */
 export class ChatManager {
   private agentStoreReady?: Promise<void>;
+  private generationCounter = 0;
   private activeGenerations = new Map<string, {
     abortController: AbortController;
     agentResult?: AgentResult;
+    generationId: number;
   }>();
 
   constructor(
@@ -490,6 +527,7 @@ export class ChatManager {
 
   private handleFallbackModelUsed(
     sessionId: string,
+    generationId: number,
     payload: {
       primaryModel: string;
       fallbackModel: string;
@@ -510,7 +548,38 @@ export class ChatManager {
     chatStreamManager.broadcast(sessionId, {
       type: "fallback",
       data: payload,
-    });
+    }, { generationId });
+  }
+
+  /**
+   * Allocate a fresh generation slot for a session before subscribing/streaming.
+   *
+   * Returns a monotonically increasing `generationId` plus an `AbortController` that
+   * later steps (the SSE route, `sendMessage`, `cancelGeneration`) use to drive and
+   * tear down this specific generation. Any in-flight generation for the same
+   * session is pre-emptively aborted; its lingering broadcasts will carry the old
+   * generationId, which `ChatStreamManager` filters out for new subscribers.
+   *
+   * Routes that subscribe to SSE before invoking `sendMessage` should call this
+   * first so subscription and broadcast generationIds are tied together.
+   */
+  beginGeneration(sessionId: string): { generationId: number; abortController: AbortController } {
+    const existing = this.activeGenerations.get(sessionId);
+    if (existing) {
+      existing.abortController.abort();
+      if (existing.agentResult) {
+        try {
+          existing.agentResult.session.dispose?.();
+        } catch (err) {
+          diagnostics.error(`Error disposing previous agent session during pre-emption:`, err);
+        }
+      }
+    }
+    this.generationCounter += 1;
+    const generationId = this.generationCounter;
+    const abortController = new AbortController();
+    this.activeGenerations.set(sessionId, { abortController, generationId });
+    return { generationId, abortController };
   }
 
   /**
@@ -686,9 +755,25 @@ export class ChatManager {
     modelProvider?: string,
     modelId?: string,
     attachments?: ChatAttachment[],
+    options?: { generationId?: number },
   ): Promise<void> {
-    const abortController = new AbortController();
-    this.activeGenerations.set(sessionId, { abortController });
+    // The SSE route allocates a generation via `beginGeneration` so it can subscribe
+    // with a matching filter before this method runs. Direct callers (tests, internal
+    // code) pass nothing and we allocate a generation here.
+    const preallocated = options?.generationId !== undefined
+      ? this.activeGenerations.get(sessionId)
+      : undefined;
+    let generationId: number;
+    let abortController: AbortController;
+    if (preallocated && preallocated.generationId === options?.generationId) {
+      generationId = preallocated.generationId;
+      abortController = preallocated.abortController;
+    } else {
+      const allocated = this.beginGeneration(sessionId);
+      generationId = allocated.generationId;
+      abortController = allocated.abortController;
+    }
+    const broadcastOptions = { generationId };
 
     const session = this.chatStore.getSession(sessionId);
     let agentResult: AgentResult | undefined;
@@ -712,7 +797,7 @@ export class ChatManager {
         chatStreamManager.broadcast(sessionId, {
           type: "error",
           data: `Chat session ${sessionId} not found`,
-        });
+        }, broadcastOptions);
         return;
       }
 
@@ -732,7 +817,7 @@ export class ChatManager {
         chatStreamManager.broadcast(sessionId, {
           type: "error",
           data: `Failed to save message: ${err instanceof Error ? err.message : "Unknown error"}`,
-        });
+        }, broadcastOptions);
         return;
       }
 
@@ -879,21 +964,21 @@ export class ChatManager {
           triggerPoint: "session-creation" | "prompt-time";
         }) => {
           fallbackInfo = payload;
-          this.handleFallbackModelUsed(sessionId, payload);
+          this.handleFallbackModelUsed(sessionId, generationId, payload);
         },
         onThinking: (delta: string) => {
           accumulatedThinking += delta;
           chatStreamManager.broadcast(sessionId, {
             type: "thinking",
             data: delta,
-          });
+          }, broadcastOptions);
         },
         onText: (delta: string) => {
           accumulatedText += delta;
           chatStreamManager.broadcast(sessionId, {
             type: "text",
             data: delta,
-          });
+          }, broadcastOptions);
         },
         onToolStart: (name: string, args?: Record<string, unknown>) => {
           const pendingForTool = pendingToolStarts.get(name) ?? [];
@@ -903,7 +988,7 @@ export class ChatManager {
           chatStreamManager.broadcast(sessionId, {
             type: "tool_start",
             data: { toolName: name, args },
-          });
+          }, broadcastOptions);
         },
         onToolEnd: (name: string, isError: boolean, result?: unknown) => {
           const pendingForTool = pendingToolStarts.get(name);
@@ -922,7 +1007,7 @@ export class ChatManager {
           chatStreamManager.broadcast(sessionId, {
             type: "tool_end",
             data: { toolName: name, isError, result },
-          });
+          }, broadcastOptions);
         },
       };
 
@@ -937,7 +1022,7 @@ export class ChatManager {
       } else {
         agentResult = await createFnAgent(sessionOptions);
       }
-      this.activeGenerations.set(sessionId, { abortController, agentResult });
+      this.activeGenerations.set(sessionId, { abortController, agentResult, generationId });
 
       if (abortController.signal.aborted) {
         agentResult.session.dispose?.();
@@ -960,7 +1045,7 @@ export class ChatManager {
         chatStreamManager.broadcast(sessionId, {
           type: "error",
           data: sessionErrorMessage,
-        });
+        }, broadcastOptions);
         return;
       }
 
@@ -1021,13 +1106,13 @@ export class ChatManager {
           },
           attachments,
         },
-      });
+      }, broadcastOptions);
     } catch (err) {
       if (abortController.signal.aborted) {
         chatStreamManager.broadcast(sessionId, {
           type: "error",
           data: "Generation cancelled",
-        });
+        }, broadcastOptions);
         return;
       }
 
@@ -1054,9 +1139,15 @@ export class ChatManager {
       chatStreamManager.broadcast(sessionId, {
         type: "error",
         data: errorMessage,
-      });
+      }, broadcastOptions);
     } finally {
-      this.activeGenerations.delete(sessionId);
+      // Only clear the active-generation slot if it still belongs to us. If a newer
+      // sendMessage pre-empted us via beginGeneration, the slot now holds that newer
+      // generation's controller and must not be deleted by our cleanup.
+      const current = this.activeGenerations.get(sessionId);
+      if (current?.generationId === generationId) {
+        this.activeGenerations.delete(sessionId);
+      }
 
       // Always dispose agent session
       if (agentResult) {
@@ -1088,7 +1179,7 @@ export class ChatManager {
     chatStreamManager.broadcast(sessionId, {
       type: "error",
       data: "Generation cancelled",
-    });
+    }, { generationId: entry.generationId });
 
     return true;
   }

packages/dashboard/src/routes/register-chat-routes.ts

@@ -536,7 +536,13 @@ export function registerChatRoutes(ctx: ApiRoutesContext, deps: ChatRouteDeps):
         }
       }
 
-      // Subscribe to session events
+      // Allocate a generation up front so subscription and sendMessage broadcasts
+      // share the same id. This filters out stragglers from a prior, just-cancelled
+      // generation that would otherwise hit this fresh subscriber and falsely look
+      // like an error/done for this request.
+      const { generationId } = chatManager.beginGeneration(sessionId);
+
+      // Subscribe to session events for this generation only.
       const unsubscribe = chatStreamManager.subscribe(sessionId, (event, eventId) => {
         const data = (event as { data?: unknown }).data;
         if (!writeSSEEvent(res, event.type, JSON.stringify(data ?? {}), eventId)) {
@@ -549,7 +555,7 @@ export function registerChatRoutes(ctx: ApiRoutesContext, deps: ChatRouteDeps):
           unsubscribe();
           res.end();
         }
-      });
+      }, { generationId });
 
       // Handle client disconnect
       req.on("close", () => {
@@ -577,7 +583,7 @@ export function registerChatRoutes(ctx: ApiRoutesContext, deps: ChatRouteDeps):
         chatStreamManager.broadcast(sessionId, {
           type: "error",
           data: "modelProvider and modelId must both be provided or neither",
-        });
+        }, { generationId });
         unsubscribe();
         res.end();
         return;
@@ -590,6 +596,7 @@ export function registerChatRoutes(ctx: ApiRoutesContext, deps: ChatRouteDeps):
         normalizedProvider,
         normalizedModelId,
         Array.isArray(attachments) ? attachments : undefined,
+        { generationId },
       ).catch((err: Error) => {
         chatLogger.error("Error in sendMessage", {
           error: err.message,
@@ -597,7 +604,7 @@ export function registerChatRoutes(ctx: ApiRoutesContext, deps: ChatRouteDeps):
         chatStreamManager.broadcast(sessionId, {
           type: "error",
           data: err.message || "Failed to process message",
-        });
+        }, { generationId });
       });
     } catch (err: unknown) {
       if (err instanceof ApiError) {