mirror of
https://github.com/coder/coder.git
synced 2026-06-03 13:08:25 +00:00
Kyle Carberry
d42008e93d
## 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