opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: drop buffered chat parts after their durable message commits (#25164)

Browse Source
This commit is contained in:
Kyle Carberry
2026-05-12 00:30:38 -04:00
committed by GitHub
parent 07ff3b3f90
commit 5a5cd79c4c
4 changed files with 356 additions and 752 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/coderd/x/chatd/chatd_test.go
+25 -242
View File
@@ -1257,11 +1257,9 @@ func TestSubscribeRelayMultipleReconnects(t *testing.T) {
require.GreaterOrEqual(t, int(callCount.Load()), 3)
}
// TestSubscribeRelayDialCanceledOnFastCompletion demonstrates a race
// condition in multi-replica chat streaming where the relay connection
// from the subscriber replica to the worker replica is canceled before
// it can be established because the worker completes processing before
// the async relay dial finishes.
// TestSubscribeRelayDialCanceledOnFastCompletion verifies that a
// subscriber on a remote replica still sees the committed assistant
// response when the worker completes faster than the relay dial.
//
// Scenario:
// 1. Subscriber subscribes to a chat while it's in waiting state (no relay).
@@ -1269,12 +1267,15 @@ func TestSubscribeRelayMultipleReconnects(t *testing.T) {
// 3. Subscriber receives status=running via pubsub → enterprise opens relay async.
// 4. Worker completes quickly → publishes committed message + status=waiting.
// 5. Subscriber receives status=waiting → enterprise cancels the in-progress relay dial.
// 6. The relay was never established, so no message_part events were delivered.
// 7. The committed message arrives via pubsub (durable path), but streaming is lost.
// 6. Even though the relay never delivered streaming parts, the
// committed assistant message arrives via pubsub so the user
// does not need to refresh to see the response.
//
// This reproduces the user-facing issue where refreshing the page is needed
// to see a response: the streaming tokens never arrive via the relay, and
// the response only appears after the full committed message is delivered.
// Streaming parts for committed turns are intentionally NOT replayed
// via the relay: they would duplicate the durable message on the
// user's screen. The buffer retains in-progress parts only; once an
// assistant turn commits, the parts that built it are claimed by
// the durable message ID and dropped from new buffer snapshots.
func TestSubscribeRelayDialCanceledOnFastCompletion(t *testing.T) {
t.Parallel()
@@ -1336,8 +1337,10 @@ func TestSubscribeRelayDialCanceledOnFastCompletion(t *testing.T) {
return nil, nil, nil, ctx.Err()
}
// Connect to the worker. The buffer is retained for a
// grace period after processing, so the relay still gets
// the message_part snapshot.
// grace period after processing, so the relay session
// can complete (control events, status updates) even
// though every part has been claimed by its durable
// message and the snapshot is empty.
snapshot, relayEvents, cancel, ok := worker.Subscribe(ctx, chatID, requestHeader, math.MaxInt64)
if !ok {
return nil, nil, nil, xerrors.New("worker subscribe failed")
@@ -1381,27 +1384,22 @@ func TestSubscribeRelayDialCanceledOnFastCompletion(t *testing.T) {
// Release the relay dial now that the worker is done.
close(workerDone)
// Collect all events that arrived at the subscriber.
var messageParts []string
// Collect events that arrived at the subscriber. The committed
// assistant message is guaranteed to arrive via pubsub even when
// the relay dial races worker completion; streaming parts are
// best-effort and are not asserted here because the buffer drops
// already-committed parts to prevent duplicate UI rendering.
var committedAssistantMsgs int
// Drain events until we see both the committed message (via
// pubsub) and at least one streaming part (via relay
// drain-and-close).
require.Eventually(t, func() bool {
select {
case event := <-events:
switch event.Type {
case codersdk.ChatStreamEventTypeMessagePart:
if event.MessagePart != nil {
messageParts = append(messageParts, event.MessagePart.Part.Text)
}
case codersdk.ChatStreamEventTypeMessage:
if event.Message != nil && event.Message.Role == codersdk.ChatMessageRoleAssistant {
committedAssistantMsgs++
}
if event.Type == codersdk.ChatStreamEventTypeMessage &&
event.Message != nil &&
event.Message.Role == codersdk.ChatMessageRoleAssistant {
committedAssistantMsgs++
}
return committedAssistantMsgs > 0 && len(messageParts) > 0
return committedAssistantMsgs > 0
default:
return false
}
@@ -1415,221 +1413,6 @@ func TestSubscribeRelayDialCanceledOnFastCompletion(t *testing.T) {
// The relay dial was attempted when status=running arrived.
require.True(t, dialAttempted.Load(),
"relay dial should have been attempted when status changed to running")
// Streaming parts are now received even though the relay was
// slower than the worker: the OSS buffer retention grace period
// keeps parts available, and the enterprise relay completes the
// dial (drain-and-close) instead of canceling it immediately.
require.NotEmpty(t, messageParts,
"streaming parts should be received via the relay even when the "+
"worker completes before the relay is established")
}
// TestSubscribeRelayDrainWithinGraceLeavesBufferRetained characterizes
// the multi-replica trigger for the retained-buffer leak: an enterprise
// relay drain (relayDrainTimeout = 200ms) always fires inside the
// worker's 5s grace window, so the worker-side subscriber-detach hits
// cleanupStreamIfIdle's early-return and the buffer stays mapped.
// streamJanitorLoop is the timer-driven backstop.
//
// The assertion is behavioral (a fresh worker.Subscribe sees the
// retained message_parts) rather than a chatStreams-size check because
// _test.go identifiers in coderd/x/chatd do not link into the
// enterprise test binary, and adding a production accessor for this
// isn't justified. The matching reap assertion lives in the OSS unit
// tests in coderd/x/chatd/chatd_internal_test.go.
func TestSubscribeRelayDrainWithinGraceLeavesBufferRetained(t *testing.T) {
t.Parallel()
db, ps := dbtestutil.NewDB(t)
workerID := uuid.New()
subscriberID := uuid.New()
ctx := testutil.Context(t, testutil.WaitLong)
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
if !req.Stream {
return chattest.OpenAINonStreamingResponse("relay-drain-characterization")
}
return chattest.OpenAIStreamingResponse(
chattest.OpenAITextChunks("hello ", "from ", "worker")...,
)
})
workerLogger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
// Freeze the worker's clock so streamJanitorLoop cannot race the
// buffer-retained assertion on slow CI.
workerClock := quartz.NewMock(t)
trapAcquire := workerClock.Trap().NewTicker("chatd", "acquire")
defer trapAcquire.Close()
worker := osschatd.New(osschatd.Config{
Logger: workerLogger,
Database: db,
ReplicaID: workerID,
Pubsub: ps,
PendingChatAcquireInterval: time.Millisecond,
InFlightChatStaleAfter: testutil.WaitSuperLong,
Clock: workerClock,
})
worker.Start()
trapAcquire.MustWait(ctx).MustRelease(ctx)
t.Cleanup(func() {
require.NoError(t, worker.Close())
})
// Use a mock clock for the subscriber so the relay drain
// timer never fires until we explicitly advance it. This
// removes the nondeterministic 200ms race between the drain
// timer and the multi-hop snapshot forwarding pipeline.
subscriberClock := quartz.NewMock(t)
trapDrain := subscriberClock.Trap().NewTimer("drain")
defer trapDrain.Close()
// Subscriber dials through to the worker. On cancel the relay
// drain fires well inside the worker's 5s grace, exercising the
// cleanupStreamIfIdle early-return path.
subscriber := newTestServer(t, db, ps, subscriberID, func(
ctx context.Context,
chatID uuid.UUID,
targetWorkerID uuid.UUID,
requestHeader http.Header,
) (
[]codersdk.ChatStreamEvent,
<-chan codersdk.ChatStreamEvent,
func(),
error,
) {
snapshot, relayEvents, cancel, ok := worker.Subscribe(ctx, chatID, requestHeader, math.MaxInt64)
if !ok {
return nil, nil, nil, xerrors.New("worker subscribe failed")
}
return snapshot, relayEvents, cancel, nil
}, subscriberClock)
user, org, model := seedChatDependencies(t, db)
setOpenAIProviderBaseURL(ctx, t, db, openAIURL)
chat := seedWaitingChat(t, db, org.ID, user, model, "relay-drain-characterization")
// Seed the pending turn directly instead of using SendMessage.
// SendMessage publishes a pending control notification that is
// irrelevant to this relay-retention case. Under CI that
// notification can arrive after processChat arms its control
// subscription and interrupt the worker before it emits parts.
dbgen.ChatMessage(t, db, database.ChatMessage{
ChatID: chat.ID,
CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true},
Role: database.ChatMessageRoleUser,
Content: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`[{"type":"text","text":"hello"}]`),
Valid: true,
},
})
_, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{
ID: chat.ID,
Status: database.ChatStatusPending,
})
require.NoError(t, err)
// Attach before processing so the relay opens as soon as
// status=running arrives.
_, events, subCancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
require.True(t, ok)
// Wake the worker with the acquire ticker. This keeps the
// setup free of pending control notifications while still
// exercising the normal processing loop.
workerClock.Advance(time.Millisecond).MustWait(ctx)
// Drain events until processing has clearly completed: we need
// the assistant message and at least one message_part so we know
// processChat's defer has flipped buffering=false and populated
// bufferRetainedAt before the subscriber detaches.
//
// Each Eventually gets its own context so one slow assertion
// cannot starve subsequent ones of their deadline.
var committedAssistantMsgs int
var messagePartsSeen int
evCtx1 := testutil.Context(t, testutil.WaitLong)
testutil.Eventually(evCtx1, t, func(context.Context) bool {
select {
case event := <-events:
switch event.Type {
case codersdk.ChatStreamEventTypeMessagePart:
messagePartsSeen++
case codersdk.ChatStreamEventTypeMessage:
if event.Message != nil && event.Message.Role == codersdk.ChatMessageRoleAssistant {
committedAssistantMsgs++
}
}
return committedAssistantMsgs > 0 && messagePartsSeen > 0
default:
return false
}
}, testutil.IntervalFast)
// Drain all NewTimer("drain") calls in a background goroutine.
// The merge loop may create one or two drain timers depending
// on the relative ordering of the status=WAITING pubsub
// notification and the async relay dial completion. Each
// trapped call must be released so the production goroutine
// is unblocked, and the clock must be advanced past the
// 200ms drain timeout to fire the timer.
var drainsFired atomic.Int32
go func() {
for {
call, err := trapDrain.Wait(ctx)
if err != nil {
return
}
if err := call.Release(ctx); err != nil {
return
}
subscriberClock.Advance(200 * time.Millisecond)
drainsFired.Add(1)
}
}()
// Wait for DB status=waiting AND at least one drain timer to
// have fired. Checking drainsFired proves the relay was torn
// down by the drain path, not by context cancellation.
evCtx2 := testutil.Context(t, testutil.WaitLong)
testutil.Eventually(evCtx2, t, func(ctx context.Context) bool {
if drainsFired.Load() == 0 {
return false
}
fromDB, dbErr := db.GetChatByID(ctx, chat.ID)
if dbErr != nil {
return false
}
return fromDB.Status == database.ChatStatusWaiting
}, testutil.IntervalFast)
// Tear the subscriber down inside the worker's grace window.
subCancel()
// A fresh worker.Subscribe still sees the retained
// message_parts: the buffer was not reaped when the relay
// drained. Eventually absorbs the short window before the
// worker observes the teardown. The retry itself re-enters
// cleanupStreamIfIdle via its own cancel defer but still
// early-returns because grace is still open.
evCtx3 := testutil.Context(t, testutil.WaitLong)
testutil.Eventually(evCtx3, t, func(ctx context.Context) bool {
snap, _, snapCancel, ok := worker.Subscribe(ctx, chat.ID, nil, math.MaxInt64)
if !ok {
return false
}
defer snapCancel()
for _, e := range snap {
if e.Type == codersdk.ChatStreamEventTypeMessagePart {
return true
}
}
return false
}, testutil.IntervalFast,
"retained buffer must still contain message_parts after the "+
"relay drains within grace")
}
// TestSubscribeRelayEstablishedMidStream demonstrates that when the