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:
@@ -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(() => {
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user