fix(site): clear stream state on WebSocket reconnect to prevent text duplication (#22838)

When the chat WebSocket reconnects, the server replays all buffered
`message_part` events in the initial snapshot. The client's `onOpen`
callback only cleared the stream error but **not** the stream state, so
replayed parts appended to the stale accumulator, doubling (or further
multiplying) the visible text with each reconnect. A page refresh would
clear the issue temporarily since it creates a fresh `ChatStore`.

This was caused by:

- **Server** (`coderd/chatd/chatd.go`): `Subscribe()` unconditionally
  includes all buffered `message_part` events in the snapshot sent to
  new connections. The `afterMessageID` parameter only filters durable
  DB messages, not ephemeral stream parts.
- **Client** (`ChatContext.ts`): The `onOpen` callback in
  `createReconnectingWebSocket` called `store.clearStreamError()` but
  not `store.clearStreamState()`. When the reconnected stream replays
  buffered `message_part` events, `applyMessagePartToStreamState`
  blindly appends text via `appendTextBlock`.

The fix was to add `store.clearStreamState()` in the `onOpen` callback
so replayed parts build from a clean slate instead of appending to stale
content.

A red/green verification test was added to ensure the fix works as
expected.
This commit is contained in:
Mathias Fredriksson
2026-03-09 16:36:17 +02:00
committed by GitHub
parent a96ec4c397
commit a072d542a5
2 changed files with 143 additions and 1 deletions
@@ -1975,6 +1975,144 @@ describe("useChatStore", () => {
);
});
it("does not duplicate streamed text after reconnect", async () => {
// The reconnect timer in createReconnectingWebSocket
// fires inside a setTimeout. With real timers the
// callback runs outside any act() boundary, so
// startTransition updates from the reconnected socket
// never commit to React state. Fake timers with
// shouldAdvanceTime let us control when the reconnect
// timer fires (via advanceTimersByTime inside act)
// while still letting waitFor's internal polling work.
vi.useFakeTimers({ shouldAdvanceTime: true });
immediateAnimationFrame();
const chatID = "chat-reconnect-dedup";
const existingMessage = makeMessage(chatID, 1, "user", "hello");
const watchMock = vi.mocked(watchChat);
// Return a fresh MockSocket for each connection attempt
// so we can control the first and second sockets
// independently.
watchMock.mockImplementation(() => createMockSocket() as never);
const queryClient = createTestQueryClient();
const wrapper = ({ children }: PropsWithChildren) => (
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
);
const setChatErrorReason = vi.fn();
const clearChatErrorReason = vi.fn();
const { result } = renderHook(
() => {
const { store } = useChatStore({
chatID,
chatMessages: [existingMessage],
chatRecord: makeChat(chatID),
chatData: {
chat: makeChat(chatID),
messages: [existingMessage],
queued_messages: [],
},
chatQueuedMessages: [],
setChatErrorReason,
clearChatErrorReason,
});
return {
streamState: useChatSelector(store, selectStreamState),
};
},
{ wrapper },
);
// Wait for the first socket to be created.
await waitFor(() => {
expect(watchMock).toHaveBeenCalledWith(chatID, 1);
});
const socket1 = watchMock.mock.results[0].value as MockSocket;
// Simulate the first socket opening successfully.
act(() => socket1.emitOpen());
// Stream "Hello world" on the first connection.
act(() => {
socket1.emitData({
type: "message_part",
chat_id: chatID,
message_part: {
role: "assistant",
part: { type: "text", text: "Hello" },
},
});
socket1.emitData({
type: "message_part",
chat_id: chatID,
message_part: {
role: "assistant",
part: { type: "text", text: " world" },
},
});
});
await waitFor(() => {
expect(result.current.streamState?.blocks).toEqual([
{ type: "response", text: "Hello world" },
]);
});
// --- Disconnect the first socket ---
act(() => socket1.emitClose());
// Advance past the reconnect backoff (1 s base delay)
// inside act() so the setTimeout callback fires within
// React's scheduling context.
await act(async () => {
vi.advanceTimersByTime(1_500);
});
// A second socket should now exist.
expect(watchMock).toHaveBeenCalledTimes(2);
const socket2 = watchMock.mock.results[1].value as MockSocket;
// Simulate the reconnected socket opening. This is
// where onOpen fires clearStreamState().
act(() => socket2.emitOpen());
// Replay the same parts the server would send on the
// new connection.
act(() => {
socket2.emitData({
type: "message_part",
chat_id: chatID,
message_part: {
role: "assistant",
part: { type: "text", text: "Hello" },
},
});
socket2.emitData({
type: "message_part",
chat_id: chatID,
message_part: {
role: "assistant",
part: { type: "text", text: " world" },
},
});
});
// Without clearStreamState() in onOpen the replayed
// parts would append to the stale accumulator, producing
// "Hello worldHello world". The fix ensures a clean
// slate so we get the correct single copy.
await waitFor(() => {
expect(result.current.streamState?.blocks).toEqual([
{ type: "response", text: "Hello world" },
]);
});
vi.useRealTimers();
});
it("clears chatErrorReason when status transitions to non-error", async () => {
immediateAnimationFrame();
@@ -777,8 +777,12 @@ export const useChatStore = (
},
onOpen() {
// Connection succeeded — clear any previous disconnect
// error.
// error and stale stream state. Clearing stream state
// is critical for reconnections: the server replays
// all buffered message_part events, so we must start
// from a clean slate to avoid duplicating text.
store.clearStreamError();
store.clearStreamState();
},
onDisconnect(attempt) {
// Show the error only on the first disconnect (not