diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index 6c992143dd..a6e35caf28 100644 --- a/coderd/chatd/chatd.go +++ b/coderd/chatd/chatd.go @@ -1673,6 +1673,7 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { // Determine the final status and last error to set when we're done. status := database.ChatStatusWaiting + wasInterrupted := false lastError := "" remainingQueuedMessages := []database.ChatQueuedMessage{} shouldPublishQueueUpdate := false @@ -1802,10 +1803,10 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { // Send a web push notification when the agent finishes // processing. We only notify for terminal states (waiting - // = success, error = failure) and skip sub-agent chats to - // avoid spamming the user with notifications for internal - // delegation. - if p.webpushDispatcher != nil && p.webpushDispatcher.PublicKey() != "" && !chat.ParentChatID.Valid { + // = success, error = failure) and skip sub-agent chats + // and user-interrupted chats to avoid unnecessary + // notifications. + if p.webpushDispatcher != nil && p.webpushDispatcher.PublicKey() != "" && !chat.ParentChatID.Valid && !wasInterrupted { if status == database.ChatStatusWaiting || status == database.ChatStatusError { pushMsg := codersdk.WebpushMessage{ Title: chat.Title, @@ -1833,6 +1834,7 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { if errors.Is(err, chatloop.ErrInterrupted) || errors.Is(context.Cause(chatCtx), chatloop.ErrInterrupted) { logger.Info(ctx, "chat interrupted") status = database.ChatStatusWaiting + wasInterrupted = true return } if isShutdownCancellation(ctx, chatCtx, err) { diff --git a/coderd/chatd/chatd_test.go b/coderd/chatd/chatd_test.go index 5b4c76af0e..075c83ecc0 100644 --- a/coderd/chatd/chatd_test.go +++ b/coderd/chatd/chatd_test.go @@ -1187,6 +1187,116 @@ func setOpenAIProviderBaseURL( require.NoError(t, err) } +func TestInterruptChatDoesNotSendWebPushNotification(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + // Set up a mock OpenAI that blocks until the request context is + // canceled (i.e. until the chat is interrupted). + streamStarted := make(chan struct{}) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + chunks := make(chan chattest.OpenAIChunk, 1) + go func() { + defer close(chunks) + chunks <- chattest.OpenAITextChunks("partial")[0] + select { + case <-streamStarted: + default: + close(streamStarted) + } + // Block until the chat context is canceled by the interrupt. + <-req.Context().Done() + }() + return chattest.OpenAIResponse{StreamingChunks: chunks} + }) + + // Mock webpush dispatcher that records calls. + mockPush := &mockWebpushDispatcher{} + + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + server := chatd.New(chatd.Config{ + Logger: logger, + Database: db, + ReplicaID: uuid.New(), + Pubsub: ps, + PendingChatAcquireInterval: 10 * time.Millisecond, + InFlightChatStaleAfter: testutil.WaitSuperLong, + WebpushDispatcher: mockPush, + }) + t.Cleanup(func() { + require.NoError(t, server.Close()) + }) + + user, model := seedChatDependencies(ctx, t, db) + setOpenAIProviderBaseURL(ctx, t, db, openAIURL) + + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OwnerID: user.ID, + Title: "interrupt-no-push", + ModelConfigID: model.ID, + InitialUserContent: []fantasy.Content{fantasy.TextContent{Text: "hello"}}, + }) + require.NoError(t, err) + + // Wait for the chat to be picked up and start streaming. + require.Eventually(t, func() bool { + fromDB, dbErr := db.GetChatByID(ctx, chat.ID) + if dbErr != nil { + return false + } + return fromDB.Status == database.ChatStatusRunning && fromDB.WorkerID.Valid + }, testutil.WaitMedium, testutil.IntervalFast) + + require.Eventually(t, func() bool { + select { + case <-streamStarted: + return true + default: + return false + } + }, testutil.WaitMedium, testutil.IntervalFast) + + // Interrupt the chat. + updated := server.InterruptChat(ctx, chat) + require.Equal(t, database.ChatStatusWaiting, updated.Status) + + // Wait for the chat to finish processing and return to waiting. + require.Eventually(t, func() bool { + fromDB, dbErr := db.GetChatByID(ctx, chat.ID) + if dbErr != nil { + return false + } + return fromDB.Status == database.ChatStatusWaiting && !fromDB.WorkerID.Valid + }, testutil.WaitMedium, testutil.IntervalFast) + + // Verify no web push notification was dispatched. + require.Equal(t, int32(0), mockPush.dispatchCount.Load(), + "expected no web push dispatch for an interrupted chat") +} + +// mockWebpushDispatcher implements webpush.Dispatcher and records Dispatch calls. +type mockWebpushDispatcher struct { + dispatchCount atomic.Int32 +} + +func (m *mockWebpushDispatcher) Dispatch(_ context.Context, _ uuid.UUID, _ codersdk.WebpushMessage) error { + m.dispatchCount.Add(1) + return nil +} + +func (*mockWebpushDispatcher) Test(_ context.Context, _ codersdk.WebpushSubscription) error { + return nil +} + +func (*mockWebpushDispatcher) PublicKey() string { + return "test-vapid-public-key" +} + func TestCloseDuringShutdownContextCanceledShouldRetryOnNewReplica(t *testing.T) { t.Parallel()