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:
Kyle Carberry
2026-03-18 11:48:28 +00:00
committed by GitHub
parent aa3cee6410
commit d42008e93d
3 changed files with 179 additions and 5 deletions
+20 -1
View File
@@ -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 {