mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix: persist partial assistant response when chat is interrupted mid-stream (#23193)
## Problem
When a user cancels a streaming chat response mid-stream, the partial
content disappears entirely — both from the UI and the database. The
streamed text vanishes as if the response never happened.
## Root Causes
Three issues combine to prevent partial message persistence on
interrupt:
### 1. StreamPartTypeError only matched `context.Canceled`
(`chatloop.go`)
The interrupt detection in `processStepStream` checked:
```go
errors.Is(part.Error, context.Canceled) && errors.Is(context.Cause(ctx), ErrInterrupted)
```
But some providers propagate `ErrInterrupted` directly as the stream
error rather than wrapping it in `context.Canceled`. This caused the
condition to fail, so `flushActiveState` was never called and partial
text accumulated in `activeTextContent` was lost.
### 2. No post-loop interrupt check (`chatloop.go`)
If the stream iterator stops yielding parts without producing a
`StreamPartTypeError` (e.g., a provider that silently closes the
response body on cancel), there was no check after the `for part :=
range stream` loop to detect the interrupt and flush active state.
### 3. Worker ownership check blocked interrupted persists (`chatd.go`)
`InterruptChat` → `setChatWaiting` clears `worker_id` in the DB
**before** the chatloop detects the interrupt. When
`persistInterruptedStep` (using `context.WithoutCancel`) tried to write
the partial message, the ownership check:
```go
if !lockedChat.WorkerID.Valid || lockedChat.WorkerID.UUID != p.workerID {
return chatloop.ErrInterrupted // always blocks!
}
```
unconditionally rejected the write. The error was silently logged as a
warning.
## Fix
- **Broaden the `StreamPartTypeError` interrupt detection** to match
both `context.Canceled` and `ErrInterrupted` as the stream error.
- **Add a post-loop interrupt check** in `processStepStream` that
flushes active state when the context was canceled with
`ErrInterrupted`.
- **Allow `persistStep` to write when the chat is in `waiting` status**
(interrupt) even if `worker_id` was cleared. The `pending` status (from
`EditMessage`, where history is truncated) still correctly blocks stale
writes.
## Testing
Added `TestInterruptChatPersistsPartialResponse` — an end-to-end
integration test that:
1. Streams partial text chunks from a mock LLM
2. Waits for the chatloop to publish `message_part` events (confirming
chunks were processed)
3. Interrupts the chat mid-stream
4. Verifies the partial assistant message is persisted in the database
with the expected text content
This commit is contained in:
+20
-1
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user