diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 6483b18559..7393722e6d 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -5040,9 +5040,7 @@ func (p *Server) tryAutoPromoteQueuedMessage( ).withCreatedBy(chat.OwnerID)) msgs, err := insertChatMessageWithStore(ctx, tx, msgParams) if err != nil { - logger.Error(ctx, "failed to promote queued message", - slog.F("queued_message_id", nextQueued.ID), slog.Error(err)) - return nil, nil, false, nil + return nil, nil, false, xerrors.Errorf("insert promoted message: %w", err) } msg := msgs[0] @@ -5148,7 +5146,8 @@ func (p *Server) finishActiveChat( var promoteErr error result.promotedMessage, result.remainingQueuedMessages, result.shouldPublishQueueUpdate, promoteErr = p.tryAutoPromoteQueuedMessage(ctx, tx, latestChat) if promoteErr != nil { - logger.Error(ctx, "failed to auto-promote queued message", slog.Error(promoteErr)) + logger.Error(ctx, "auto-promote queued message failed, rolling back", slog.Error(promoteErr)) + return xerrors.Errorf("auto-promote queued message: %w", promoteErr) } else if result.promotedMessage != nil { status = database.ChatStatusPending } diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 0ad47577d2..6cc6956ab5 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -3524,9 +3524,9 @@ func TestProcessChat_IgnoresStaleControlNotification(t *testing.T) { clock: clock, workerID: workerID, chatHeartbeatInterval: time.Minute, + metrics: chatloop.NopMetrics(), configCache: newChatConfigCache(ctx, db, clock), heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), - metrics: chatloop.NopMetrics(), } // Publish a stale "pending" notification on the control channel @@ -3680,6 +3680,7 @@ func TestHeartbeatTick_StolenChatIsInterrupted(t *testing.T) { clock: clock, workerID: workerID, chatHeartbeatInterval: time.Minute, + metrics: chatloop.NopMetrics(), heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), } @@ -3760,6 +3761,7 @@ func TestHeartbeatTick_DBErrorDoesNotInterruptChats(t *testing.T) { clock: clock, workerID: uuid.New(), chatHeartbeatInterval: time.Minute, + metrics: chatloop.NopMetrics(), heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), } @@ -4717,3 +4719,224 @@ func TestGetWorkspaceConn_DialErrorNotMisclassifiedAsTimeout(t *testing.T) { // The original dial error should propagate. require.ErrorContains(t, err, "authentication failed") } + +// TestAutoPromote_InsertFailureRollsBackTransaction verifies that when +// tryAutoPromoteQueuedMessage pops a queued message but the subsequent +// insert fails, the error propagates to the InTx callback, causing the +// transaction to roll back and preserving the queued message. +func TestAutoPromote_InsertFailureRollsBackTransaction(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + tx := dbmock.NewMockStore(ctrl) + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + ps := dbpubsub.NewInMemory() + clock := quartz.NewReal() + + chatID := uuid.New() + workerID := uuid.New() + ownerID := uuid.New() + modelConfigID := uuid.New() + + waitingChat := database.Chat{ + ID: chatID, + OwnerID: ownerID, + LastModelConfigID: modelConfigID, + Status: database.ChatStatusWaiting, + WorkerID: uuid.NullUUID{UUID: workerID, Valid: true}, + } + queuedMsg := database.ChatQueuedMessage{ + ID: 1, + ChatID: chatID, + Content: []byte(`[{"type":"text","text":"queued"}]`), + } + insertErr := xerrors.New("insert failed") + + server := &Server{ + db: db, + logger: logger, + pubsub: ps, + configCache: newChatConfigCache(ctx, db, clock), + } + + // The caller runs tryAutoPromoteQueuedMessage inside InTx. + // Wire the mock to execute the callback against the TX mock. + var txErr error + db.EXPECT().InTx(gomock.Any(), gomock.Any()).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { + txErr = fn(tx) + return txErr + }, + ) + + // Inside the TX: lock chat, get queued messages, resolve model + // config, pop queued message, insert fails. + tx.EXPECT().GetChatByIDForUpdate(gomock.Any(), chatID).Return(waitingChat, nil) + tx.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return([]database.ChatQueuedMessage{queuedMsg}, nil) + tx.EXPECT().GetChatModelConfigByID(gomock.Any(), modelConfigID).Return(database.ChatModelConfig{ID: modelConfigID}, nil) + tx.EXPECT().PopNextQueuedMessage(gomock.Any(), chatID).Return(queuedMsg, nil) + tx.EXPECT().InsertChatMessages(gomock.Any(), gomock.Any()).Return(nil, insertErr) + + // Invoke tryAutoPromoteQueuedMessage through the same InTx + // pattern the processChat defer uses. The test directly calls + // the production path to verify error propagation. + _ = db.InTx(func(txStore database.Store) error { + latestChat, err := txStore.GetChatByIDForUpdate(ctx, chatID) + if err != nil { + return err + } + + _, _, _, promoteErr := server.tryAutoPromoteQueuedMessage(ctx, txStore, latestChat) + if promoteErr != nil { + return promoteErr + } + + // This code path should not be reached when the insert + // fails, because promoteErr should be non-nil. + return nil + }, nil) + + // The InTx callback must return a non-nil error so the + // transaction rolls back, preserving the queued message. + require.Error(t, txErr, "InTx callback should return error when insert fails") +} + +// TestAutoPromote_WakesRunLoopAfterPromotion verifies that after the +func TestAutoPromote_InsertFailureSkipsStatusUpdate(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + tx := dbmock.NewMockStore(ctrl) + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + ps := dbpubsub.NewInMemory() + clock := quartz.NewReal() + + chatID := uuid.New() + workerID := uuid.New() + ownerID := uuid.New() + modelConfigID := uuid.New() + + waitingChat := database.Chat{ + ID: chatID, + OwnerID: ownerID, + LastModelConfigID: modelConfigID, + Status: database.ChatStatusWaiting, + WorkerID: uuid.NullUUID{UUID: workerID, Valid: true}, + } + queuedMsg := database.ChatQueuedMessage{ + ID: 1, + ChatID: chatID, + Content: []byte(`[{"type":"text","text":"queued"}]`), + } + + wakeCh := make(chan struct{}, 1) + server := &Server{ + db: db, + logger: logger, + pubsub: ps, + clock: clock, + workerID: workerID, + wakeCh: wakeCh, + chatHeartbeatInterval: time.Minute, + metrics: chatloop.NopMetrics(), + configCache: newChatConfigCache(ctx, db, clock), + heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), + } + + // Block model resolution until the control subscriber fires. + modelBlocked := make(chan struct{}) + db.EXPECT().GetChatModelConfigByID(gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, _ uuid.UUID) (database.ChatModelConfig, error) { + <-modelBlocked + return database.ChatModelConfig{}, xerrors.New("no model") + }, + ).AnyTimes() + db.EXPECT().GetEnabledChatProviders(gomock.Any()).Return(nil, nil).AnyTimes() + db.EXPECT().GetEnabledChatModelConfigs(gomock.Any()).Return(nil, nil).AnyTimes() + db.EXPECT().GetChatUsageLimitConfig(gomock.Any()).Return( + database.ChatUsageLimitConfig{}, sql.ErrNoRows, + ).AnyTimes() + db.EXPECT().GetChatMessagesForPromptByChatID(gomock.Any(), chatID).Return(nil, nil).AnyTimes() + + // The deferred cleanup transaction: InsertChatMessages fails, + // so UpdateChatStatus must NOT be called. + db.EXPECT().InTx(gomock.Any(), gomock.Any()).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { + return fn(tx) + }, + ) + tx.EXPECT().GetChatByIDForUpdate(gomock.Any(), chatID).Return(waitingChat, nil) + tx.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return([]database.ChatQueuedMessage{queuedMsg}, nil) + tx.EXPECT().GetChatModelConfigByID(gomock.Any(), modelConfigID).Return(database.ChatModelConfig{ID: modelConfigID}, nil) + tx.EXPECT().PopNextQueuedMessage(gomock.Any(), chatID).Return(queuedMsg, nil) + tx.EXPECT().InsertChatMessages(gomock.Any(), gomock.Any()).Return( + nil, xerrors.New("insert failed"), + ) + tx.EXPECT().UpdateChatStatus(gomock.Any(), gomock.Any()).Times(0) + + // Subscribe BEFORE launching the goroutine. + runningCh := make(chan struct{}, 1) + unsubRunning, err := ps.SubscribeWithErr( + coderdpubsub.ChatStreamNotifyChannel(chatID), + func(_ context.Context, msg []byte, err error) { + if err != nil { + return + } + var notify coderdpubsub.ChatStreamNotifyMessage + if json.Unmarshal(msg, ¬ify) != nil { + return + } + if notify.Status == string(database.ChatStatusRunning) { + select { + case runningCh <- struct{}{}: + default: + } + } + }, + ) + require.NoError(t, err) + defer unsubRunning() + + chat := database.Chat{ID: chatID, OwnerID: ownerID, LastModelConfigID: modelConfigID} + processDone := make(chan struct{}) + go func() { + defer close(processDone) + server.processChat(ctx, chat) + }() + + select { + case <-runningCh: + case <-ctx.Done(): + t.Fatal("timed out waiting for running status") + } + + // Publish an interrupt so processChat exits runChat. + interruptMsg, err := json.Marshal(coderdpubsub.ChatStreamNotifyMessage{ + Status: string(database.ChatStatusWaiting), + }) + require.NoError(t, err) + err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chatID), interruptMsg) + require.NoError(t, err) + + // Unblock model resolution so runChat can exit. + close(modelBlocked) + + select { + case <-processDone: + case <-ctx.Done(): + t.Fatal("processChat did not complete") + } + + // The wake channel should NOT have a signal because the + // transaction failed before reaching UpdateChatStatus. + select { + case <-wakeCh: + t.Fatal("wake channel should not have a signal after insert failure") + default: + // No signal, as expected. + } +} diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx index 73d2c3224a..9f9de11388 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx @@ -731,7 +731,14 @@ describe("useChatStore", () => { }); await waitFor(() => { - expect(result.current.streamState).toBeNull(); + // Stream state is preserved after status=pending (the + // durable message event handles cleanup via + // needsStreamReset). Only new message_parts should be + // blocked by the shouldApplyMessagePart gate. + expect(result.current.streamState).not.toBeNull(); + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "first" }, + ]); }); act(() => { @@ -749,7 +756,12 @@ describe("useChatStore", () => { }); await waitFor(() => { - expect(result.current.streamState).toBeNull(); + // The late message_part should not be applied because + // shouldApplyMessagePart gates on pending/waiting. + // Stream state still shows the original "first". + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "first" }, + ]); }); }); @@ -3036,6 +3048,181 @@ describe("useChatStore", () => { expect(result.current.chatStatus).toBe("running"); }); }); + + it("preserves stream state when status transitions to waiting", async () => { + immediateAnimationFrame(); + + const chatID = "chat-preserve-stream"; + const existingMessage = makeMessage(chatID, 1, "user", "hello"); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + 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), + chatMessagesData: { + messages: [existingMessage], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + // Build up stream state with a message_part. + act(() => { + mockSocket.emitData({ + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "thinking..." }, + }, + }); + }); + + await waitFor(() => { + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "thinking..." }, + ]); + }); + + // Deliver a status=waiting event (interrupt). Stream state + // should be preserved so the user continues to see the + // partial response until the durable message arrives. + act(() => { + mockSocket.emitData({ + type: "status", + chat_id: chatID, + status: { status: "waiting" }, + }); + }); + + await waitFor(() => { + expect(result.current.streamState).not.toBeNull(); + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "thinking..." }, + ]); + }); + }); + + it("clears stream state when durable message follows waiting status", async () => { + immediateAnimationFrame(); + + const chatID = "chat-durable-clears"; + const existingMessage = makeMessage(chatID, 1, "user", "hello"); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + 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), + chatMessagesData: { + messages: [existingMessage], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + orderedIDs: useChatSelector(store, selectOrderedMessageIDs), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + // Build up stream state. + act(() => { + mockSocket.emitData({ + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "partial response" }, + }, + }); + }); + + await waitFor(() => { + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "partial response" }, + ]); + }); + + // Deliver status=waiting (interrupt). Stream state should be + // preserved so the user continues to see the partial response + // until the durable message arrives. + act(() => { + mockSocket.emitData({ + type: "status", + chat_id: chatID, + status: { status: "waiting" }, + }); + }); + + // Stream state must still be present after the status change. + await waitFor(() => { + expect(result.current.streamState).not.toBeNull(); + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "partial response" }, + ]); + }); + + // Now deliver the durable assistant message. This should + // clear stream state via the needsStreamReset path. + act(() => { + mockSocket.emitData({ + type: "message", + chat_id: chatID, + message: makeMessage(chatID, 2, "assistant", "partial response"), + }); + }); + + // Stream state should now be null and the durable message + // should be in the message store. + await waitFor(() => { + expect(result.current.streamState).toBeNull(); + expect(result.current.orderedIDs).toContain(2); + }); + }); }); describe("thinking indicator event ordering", () => { diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index 5d2cda47eb..0e6e18a917 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts @@ -390,9 +390,9 @@ export const useChatStore = ( }; // Discard buffered parts without applying them. Used when - // stream state is about to be cleared (pending, waiting, - // retry) — flushing would re-populate the state that the - // event is about to clear. + // the stream is no longer active (pending, waiting, retry) + // so stale buffered parts are not applied after the + // status transition. const discardBufferedParts = () => { partsBuf.length = 0; if (partsFlushTimer !== null) { @@ -508,7 +508,6 @@ export const useChatStore = ( store.setChatStatus(nextStatus); if (nextStatus === "pending" || nextStatus === "waiting") { discardBufferedParts(); - store.clearStreamState(); store.clearRetryState(); } if (nextStatus === "running") {