mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix(site): don't wipe stream state when non-assistant message arrives (#23295)
scheduleStreamReset() fired for every durable message event with changed=true, including user messages (e.g. promoted queued messages). When a batch contained trailing message_parts followed by a user message event, the batch loop flushed the parts (building stream state), then scheduleStreamReset cleared it immediately. Restrict the reset to assistant messages, which are the only role that ends a streaming turn.
This commit is contained in:
@@ -2450,6 +2450,111 @@ describe("useChatStore", () => {
|
||||
expect(result.current.queuedMessages).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
|
||||
it("does not wipe in-progress stream state when user message arrives in batch", async () => {
|
||||
immediateAnimationFrame();
|
||||
|
||||
const chatID = "chat-promote-stream";
|
||||
const msg1 = makeMessage(chatID, 1, "user", "hello");
|
||||
const msg2 = makeMessage(chatID, 2, "assistant", "hi");
|
||||
|
||||
const mockSocket = createMockSocket();
|
||||
vi.mocked(watchChat).mockReturnValue(mockSocket as never);
|
||||
|
||||
const queryClient = createTestQueryClient();
|
||||
const wrapper: FC<PropsWithChildren> = ({ children }) => (
|
||||
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
|
||||
);
|
||||
const setChatErrorReason = vi.fn();
|
||||
const clearChatErrorReason = vi.fn();
|
||||
|
||||
const queuedMsg = makeQueuedMessage(chatID, 10, "follow-up");
|
||||
const initialMessages = [msg1, msg2];
|
||||
|
||||
const initialOptions = {
|
||||
chatID,
|
||||
chatMessages: initialMessages,
|
||||
chatRecord: makeChat(chatID),
|
||||
chatMessagesData: {
|
||||
messages: initialMessages,
|
||||
queued_messages: [queuedMsg],
|
||||
has_more: false,
|
||||
},
|
||||
chatQueuedMessages: [queuedMsg],
|
||||
setChatErrorReason,
|
||||
clearChatErrorReason,
|
||||
};
|
||||
|
||||
const { result } = renderHook(
|
||||
(options: Parameters<typeof useChatStore>[0]) => {
|
||||
const { store } = useChatStore(options);
|
||||
return {
|
||||
store,
|
||||
streamState: useChatSelector(store, selectStreamState),
|
||||
chatStatus: useChatSelector(store, selectChatStatus),
|
||||
orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs),
|
||||
};
|
||||
},
|
||||
{ initialProps: initialOptions, wrapper },
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(result.current.orderedMessageIDs).toEqual([1, 2]);
|
||||
});
|
||||
|
||||
// Open the WebSocket and set the chat to running.
|
||||
act(() => {
|
||||
mockSocket.emitOpen();
|
||||
});
|
||||
act(() => {
|
||||
mockSocket.emitData({
|
||||
type: "status",
|
||||
chat_id: chatID,
|
||||
status: { status: "running" },
|
||||
});
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(result.current.chatStatus).toBe("running");
|
||||
});
|
||||
|
||||
// Deliver a batch containing trailing message_parts for
|
||||
// the current response followed by the promoted user
|
||||
// message. The batch loop flushes pending parts when it
|
||||
// hits the message event (building stream state). Before
|
||||
// the fix, scheduleStreamReset would fire for the user
|
||||
// message because it only checked `changed`, and with
|
||||
// immediateAnimationFrame the RAF fires synchronously,
|
||||
// wiping the stream state that was just built.
|
||||
const promotedUser = makeMessage(chatID, 3, "user", "follow-up");
|
||||
|
||||
act(() => {
|
||||
mockSocket.emitDataBatch([
|
||||
{
|
||||
type: "message_part",
|
||||
chat_id: chatID,
|
||||
message_part: {
|
||||
part: { type: "text", text: "I am helping you" },
|
||||
},
|
||||
},
|
||||
{
|
||||
type: "message",
|
||||
chat_id: chatID,
|
||||
message: promotedUser,
|
||||
},
|
||||
]);
|
||||
});
|
||||
|
||||
// The stream state must survive: the promoted user message
|
||||
// should not wipe the in-progress assistant stream.
|
||||
await waitFor(() => {
|
||||
expect(result.current.orderedMessageIDs).toContain(3);
|
||||
expect(result.current.streamState).not.toBeNull();
|
||||
const blocks = result.current.streamState?.blocks ?? [];
|
||||
const textBlock = blocks.find((b) => b.type === "response");
|
||||
expect(textBlock).toBeDefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("updateSidebarChat via stream events", () => {
|
||||
|
||||
@@ -725,7 +725,7 @@ export const useChatStore = (
|
||||
) {
|
||||
lastMessageIdRef.current = message.id;
|
||||
}
|
||||
if (changed) {
|
||||
if (changed && message.role === "assistant") {
|
||||
scheduleStreamReset();
|
||||
}
|
||||
// Do not update updated_at here. The global
|
||||
|
||||
Reference in New Issue
Block a user