diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index f13b52f5b5..8738bdd2e1 100644 --- a/coderd/chatd/chatd.go +++ b/coderd/chatd/chatd.go @@ -2613,12 +2613,31 @@ func (p *Server) runChat( var insertedMessages []database.ChatMessage err := p.db.InTx(func(tx database.Store) error { + // Verify this worker still owns the chat before + // inserting messages. This closes the race where + // EditMessage truncates history and clears worker_id + // while persistInterruptedStep (which uses an + // uncancelable context) is still running. + // + // When the chat is in "waiting" status (set by + // InterruptChat / setChatWaiting), the worker_id has + // already been cleared but we still want to persist + // the partial assistant response. We allow the write + // because the history has NOT been truncated — the + // user simply asked to stop. In contrast, EditMessage + // sets the chat to "pending" after truncating, so the + // pending check still correctly blocks stale writes. lockedChat, lockErr := tx.GetChatByIDForUpdate(persistCtx, chat.ID) if lockErr != nil { return xerrors.Errorf("lock chat for persist: %w", lockErr) } if !lockedChat.WorkerID.Valid || lockedChat.WorkerID.UUID != p.workerID { - return chatloop.ErrInterrupted + // The worker_id was cleared. Only allow the persist + // if the chat transitioned to "waiting" (interrupt), + // not "pending" (edit) or any other status. + if lockedChat.Status != database.ChatStatusWaiting { + return chatloop.ErrInterrupted + } } if assistantContent.Valid { diff --git a/coderd/chatd/chatd_test.go b/coderd/chatd/chatd_test.go index efafc8b6fb..7f5d71874d 100644 --- a/coderd/chatd/chatd_test.go +++ b/coderd/chatd/chatd_test.go @@ -2696,3 +2696,139 @@ func TestComputerUseSubagentToolsAndModel(t *testing.T) { require.Equal(t, database.ChatModeComputerUse, children[0].Mode.ChatMode) } + +func TestInterruptChatPersistsPartialResponse(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + // Set up a mock OpenAI that streams a partial response and then + // blocks until the request context is canceled (simulating an + // interrupt mid-stream). + chunksDelivered := 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) + // Send two partial text chunks so there is meaningful + // content to persist. + for _, c := range chattest.OpenAITextChunks("hello world") { + chunks <- c + } + // Signal that chunks have been written to the HTTP response. + select { + case <-chunksDelivered: + default: + close(chunksDelivered) + } + // Block until interrupt cancels the context. + <-req.Context().Done() + }() + return chattest.OpenAIResponse{StreamingChunks: chunks} + }) + + 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, + }) + 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-persist-test", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, + }) + require.NoError(t, err) + + // Subscribe to the chat's event stream so we can observe + // message_part events — proof the chatloop has actually + // processed the streamed chunks. + _, events, subCancel, ok := server.Subscribe(ctx, chat.ID, nil, 0) + require.True(t, ok) + defer subCancel() + + // Wait for the mock to finish sending chunks. + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + select { + case <-chunksDelivered: + return true + default: + return false + } + }, testutil.IntervalFast) + + // Drain the event channel until we see a message_part event, + // which means the chatloop has consumed and published the chunk. + gotMessagePart := false + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + for { + select { + case ev := <-events: + if ev.Type == codersdk.ChatStreamEventTypeMessagePart { + gotMessagePart = true + return true + } + default: + return gotMessagePart + } + } + }, testutil.IntervalFast) + require.True(t, gotMessagePart, "should have received at least one message_part event") + + // Now interrupt the chat — the chatloop has processed content. + updated := server.InterruptChat(ctx, chat) + require.Equal(t, database.ChatStatusWaiting, updated.Status) + + // Wait for the partial assistant message to be persisted. + // After the interrupt, the chatloop runs persistInterruptedStep + // which inserts the message and publishes a "message" event. + // We poll the DB directly for the assistant message rather than + // relying on the chat status (which transitions to "waiting" + // before the persist completes). + var assistantMsg *database.ChatMessage + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + msgs, dbErr := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: 0, + }) + if dbErr != nil { + return false + } + for i := range msgs { + if msgs[i].Role == database.ChatMessageRoleAssistant { + assistantMsg = &msgs[i] + return true + } + } + return false + }, testutil.IntervalFast) + require.NotNilf(t, assistantMsg, "expected a persisted assistant message after interrupt") + + // Parse the content and verify it contains the partial text. + parts, err := chatprompt.ParseContent(*assistantMsg) + require.NoError(t, err) + + var foundText string + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeText { + foundText += part.Text + } + } + require.Contains(t, foundText, "hello world", + "partial assistant response should contain the streamed text") +} diff --git a/coderd/chatd/chatloop/chatloop.go b/coderd/chatd/chatloop/chatloop.go index 3afed86c5a..cd0ecbc0c2 100644 --- a/coderd/chatd/chatloop/chatloop.go +++ b/coderd/chatd/chatloop/chatloop.go @@ -610,10 +610,12 @@ func processStepStream( result.providerMetadata = part.ProviderMetadata case fantasy.StreamPartTypeError: - // Detect interruption: context canceled with - // ErrInterrupted as the cause. - if errors.Is(part.Error, context.Canceled) && - errors.Is(context.Cause(ctx), ErrInterrupted) { + // Detect interruption: the stream may surface the + // cancel as context.Canceled or propagate the + // ErrInterrupted cause directly, depending on + // the provider implementation. + if errors.Is(context.Cause(ctx), ErrInterrupted) && + (errors.Is(part.Error, context.Canceled) || errors.Is(part.Error, ErrInterrupted)) { // Flush in-progress content so that // persistInterruptedStep has access to partial // text, reasoning, and tool calls that were @@ -631,6 +633,23 @@ func processStepStream( } } + // The stream iterator may stop yielding parts without + // producing a StreamPartTypeError when the context is + // canceled (e.g. some providers close the response body + // silently). Detect this case and flush partial content + // so that persistInterruptedStep can save it. + if ctx.Err() != nil && + errors.Is(context.Cause(ctx), ErrInterrupted) { + flushActiveState( + &result, + activeTextContent, + activeReasoningContent, + activeToolCalls, + toolNames, + ) + return result, ErrInterrupted + } + hasLocalToolCalls := false for _, tc := range result.toolCalls { if !tc.ProviderExecuted {