fix: isolate chat SSE broadcasts per generation

After stopping a streaming chat reply, the next message would appear
sent but show no Stop button or "Connecting…" indicator. The cancel
broadcast from the prior generation was leaking into the new SSE
subscription and 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. `sendMessage`'s cleanup also stops
deleting a newer generation's `activeGenerations` slot when an older
one finally unwinds.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
gsxdsm
2026-05-05 07:35:18 -07:00
parent f57a692eeb
commit f49ec30c54
6 changed files with 267 additions and 39 deletions

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

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

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