Revert "fix: isolate chat SSE broadcasts per generation"

This reverts commit f49ec30c54.
This commit is contained in:
gsxdsm
2026-05-05 08:14:10 -07:00
parent cee8adfb59
commit aa587347d7
6 changed files with 39 additions and 267 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, mockBeginGeneration } = vi.hoisted(() => {
const { mockChatStreamManager, mockSendMessage, mockCancelGeneration } = 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,7 +39,6 @@ const { mockChatStreamManager, mockSendMessage, mockCancelGeneration, mockBeginG
chatStreamManager.broadcast(sessionId, { type: "done", data: { messageId: "msg-1" } });
}),
mockCancelGeneration: vi.fn().mockReturnValue(false),
mockBeginGeneration: vi.fn(() => ({ generationId: 1, abortController: new AbortController() })),
};
});
@@ -58,7 +57,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; beginGeneration = mockBeginGeneration; },
ChatManager: class MockChatManager { sendMessage = mockSendMessage; cancelGeneration = mockCancelGeneration; },
chatStreamManager: mockChatStreamManager,
checkRateLimit: vi.fn().mockReturnValue(true),
getRateLimitResetTime: vi.fn().mockReturnValue(null),
@@ -126,7 +125,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, beginGeneration: mockBeginGeneration } as any });
} as any, chatManager: { sendMessage: mockSendMessage, cancelGeneration: mockCancelGeneration } as any });
});
it("uploads a valid attachment", async () => {
@@ -188,7 +187,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, { generationId: 1 });
expect(mockSendMessage).toHaveBeenCalledWith(session.id, "hello", undefined, undefined, attachments);
});
afterEach(() => {

View File

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