root/fusion
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Reapply "fix: isolate chat SSE broadcasts per generation"

Browse Source 
This reverts commit 46f1f042b0.
This commit is contained in:
gsxdsm
2026-05-05 08:21:33 -07:00
parent 46f1f042b0
commit 41375736e7
6 changed files with 267 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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(() => {

121
packages/dashboard/src/__tests__/chat-manager.test.ts
View File

@@ -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);
});
});

5
packages/dashboard/src/__tests__/chat-routes.test.ts
View File

@@ -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,
};
}

151
packages/dashboard/src/chat.ts
View File

@@ -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;
}

15
packages/dashboard/src/routes/register-chat-routes.ts
View File

@@ -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) {