diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index a9b0895d6e..262bbe0828 100644 --- a/coderd/chatd/chatd.go +++ b/coderd/chatd/chatd.go @@ -1086,6 +1086,23 @@ func (p *Server) Subscribe( } } + // Include the current chat status in the snapshot so the + // frontend can gate message_part processing correctly from + // the very first batch, without waiting for a separate REST + // query. + if err == nil { + statusEvent := codersdk.ChatStreamEvent{ + Type: codersdk.ChatStreamEventTypeStatus, + ChatID: chatID, + Status: &codersdk.ChatStreamStatus{ + Status: codersdk.ChatStatus(chat.Status), + }, + } + // Prepend so the frontend sees the status before any + // message_part events. + initialSnapshot = append([]codersdk.ChatStreamEvent{statusEvent}, initialSnapshot...) + } + // Track the last message ID we've seen for DB queries var lastMessageID int64 if len(messages) > 0 { @@ -1307,16 +1324,11 @@ func (p *Server) Subscribe( } }() } else { - // No pubsub, just merge local parts + // No pubsub, just merge local parts. + // localSnapshot was already included in initialSnapshot, + // so only forward new events here. go func() { defer close(mergedEvents) - for _, event := range localSnapshot { - select { - case <-mergedCtx.Done(): - return - case mergedEvents <- event: - } - } for event := range localParts { select { case <-mergedCtx.Done(): diff --git a/coderd/chatd/chatd_test.go b/coderd/chatd/chatd_test.go index 2c1dbb33c0..294af89a7c 100644 --- a/coderd/chatd/chatd_test.go +++ b/coderd/chatd/chatd_test.go @@ -632,6 +632,74 @@ func TestUpdateChatStatusPersistsLastError(t *testing.T) { require.False(t, fromDB.LastError.Valid) } +func TestSubscribeSnapshotIncludesStatusEvent(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + replica := newTestServer(t, db, ps, uuid.New()) + + ctx := testutil.Context(t, testutil.WaitLong) + user, model := seedChatDependencies(ctx, t, db) + + chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OwnerID: user.ID, + Title: "status-snapshot", + ModelConfigID: model.ID, + InitialUserContent: []fantasy.Content{fantasy.TextContent{Text: "hello"}}, + }) + require.NoError(t, err) + + snapshot, _, cancel, ok := replica.Subscribe(ctx, chat.ID, nil) + require.True(t, ok) + t.Cleanup(cancel) + + // The first event in the snapshot must be a status event. + require.NotEmpty(t, snapshot) + require.Equal(t, codersdk.ChatStreamEventTypeStatus, snapshot[0].Type) + require.NotNil(t, snapshot[0].Status) + require.Equal(t, codersdk.ChatStatusPending, snapshot[0].Status.Status) +} + +func TestSubscribeNoPubsubNoDuplicateMessageParts(t *testing.T) { + t.Parallel() + + // Use nil pubsub to force the no-pubsub path. + db, _ := dbtestutil.NewDB(t) + replica := newTestServer(t, db, nil, uuid.New()) + + ctx := testutil.Context(t, testutil.WaitLong) + user, model := seedChatDependencies(ctx, t, db) + + chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OwnerID: user.ID, + Title: "no-dup-parts", + ModelConfigID: model.ID, + InitialUserContent: []fantasy.Content{fantasy.TextContent{Text: "hello"}}, + }) + require.NoError(t, err) + + snapshot, events, cancel, ok := replica.Subscribe(ctx, chat.ID, nil) + require.True(t, ok) + t.Cleanup(cancel) + + // Snapshot should have events (at minimum: status + message). + require.NotEmpty(t, snapshot) + + // The events channel should NOT immediately produce any + // events — the snapshot already contained everything. Before + // the fix, localSnapshot was replayed into the channel, + // causing duplicates. + select { + case event, ok := <-events: + if ok { + t.Fatalf("unexpected event from channel (would be a duplicate): type=%s", event.Type) + } + // Channel closed without events is fine. + case <-time.After(200 * time.Millisecond): + // No events — correct behavior. + } +} + func newTestServer( t *testing.T, db database.Store, diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx index 1b682dff82..79f749afeb 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx +++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx @@ -1044,6 +1044,95 @@ describe("useChatStore", () => { }); }); + it("cancels scheduled stream reset when message_part arrives after message", async () => { + immediateAnimationFrame(); + + const chatID = "chat-raf"; + const existingMessage = makeMessage(chatID, 1, "user", "hello"); + const mockSocket = createMockSocket(); + vi.mocked(watchChat).mockReturnValue(mockSocket as never); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + 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 }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID); + }); + + // Build up stream state first. + act(() => { + mockSocket.emitData({ + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "working" }, + }, + }); + }); + + await waitFor(() => { + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "working" }, + ]); + }); + + // Emit a durable message followed by a message_part in the + // same batch. The message handler calls scheduleStreamReset + // (via rAF), and the subsequent message_part handler calls + // cancelScheduledStreamReset to prevent a flash. The final + // flushMessageParts re-populates stream state. + act(() => { + mockSocket.emitDataBatch([ + { + type: "message", + chat_id: chatID, + message: makeMessage(chatID, 2, "assistant", "done"), + }, + { + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: " more" }, + }, + }, + ]); + }); + + // Stream state should be non-null because the message_part + // after the message kept it populated. + await waitFor(() => { + expect(result.current.streamState).not.toBeNull(); + }); + }); + it("startTransition deferred parts are discarded after chat switch", async () => { immediateAnimationFrame(); @@ -1214,4 +1303,74 @@ describe("useChatStore", () => { }); expect(result.current.queuedMessages).toEqual([]); }); + + it("does not apply message parts after status changes to waiting", async () => { + immediateAnimationFrame(); + + const chatID = "chat-status-guard"; + const mockSocket = createMockSocket(); + vi.mocked(watchChat).mockReturnValue(mockSocket as never); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [], + chatRecord: makeChat(chatID), + chatData: { + chat: makeChat(chatID), + messages: [], + queued_messages: [], + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID); + }); + + // Emit a batch with message_parts followed by a status change + // to "waiting". The status handler clears stream state + // synchronously, and the startTransition guard should prevent + // the deferred applyMessageParts from re-populating it. + act(() => { + mockSocket.emitDataBatch([ + { + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "should be discarded" }, + }, + }, + { + type: "status", + chat_id: chatID, + status: { status: "waiting" }, + }, + ]); + }); + + // Stream state should be null — the status change cleared it, + // and the deferred applyMessageParts should not have + // re-populated it. + await waitFor(() => { + expect(result.current.streamState).toBeNull(); + }); + }); }); diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts index b5ed72e45b..afec590fc4 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts +++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts @@ -584,6 +584,12 @@ export const useChatStore = ( if (activeChatIDRef.current !== currentChatID) { return; } + // Re-check status at execution time. A status + // event processed between scheduling and running + // this callback may have cleared stream state. + if (!shouldApplyMessagePart()) { + return; + } store.applyMessageParts(parts); }); }; @@ -599,6 +605,7 @@ export const useChatStore = ( } const part = asRecord(streamEvent.message_part?.part); if (part) { + cancelScheduledStreamReset(); pendingMessageParts.push(part); } continue;