fix: reap idle chatd stream states on a timer (#24476)

* Adds `streamJanitorLoop` to clean up stale streams every 30s
* zeroes dropped slots to aid in gc-eligibliity
* Adds regression tests in coderd/x/chatd and enterprise/coderd/x/chatd

> 🤖

coderd/x/chatd/chatd.go

@@ -14,6 +14,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"charm.land/fantasy"
@@ -100,6 +101,11 @@ const (
 	// snapshot the buffer before it is garbage-collected.
 	bufferRetainGracePeriod = 5 * time.Second
 
+	// streamJanitorInterval is how often sweepIdleStreams runs.
+	// Worst-case retention is bufferRetainGracePeriod +
+	// streamJanitorInterval.
+	streamJanitorInterval = 30 * time.Second
+
 	// DefaultMaxChatsPerAcquire is the maximum number of chats to
 	// acquire in a single processOnce call. Batching avoids
 	// waiting a full polling interval between acquisitions
@@ -2902,6 +2908,8 @@ func (p *Server) start(ctx context.Context) {
 	// Single heartbeat loop for all chats on this replica.
 	go p.heartbeatLoop(ctx)
 
+	go p.streamJanitorLoop(ctx)
+
 	acquireTicker := p.clock.NewTicker(
 		p.pendingChatAcquireInterval,
 		"chatd",
@@ -3025,6 +3033,10 @@ func (p *Server) publishToStream(chatID uuid.UUID, event codersdk.ChatStreamEven
 				state.bufferDropCount = 0
 				state.bufferLastWarnAt = now
 			}
+			// Zero the dropped slot so its *ChatStreamMessagePart is
+			// GC-eligible; the later append reuses this slot in place
+			// whenever cap > len.
+			state.buffer[0] = codersdk.ChatStreamEvent{}
 			state.buffer = state.buffer[1:]
 		}
 		state.buffer = append(state.buffer, event)
@@ -3078,6 +3090,9 @@ func (p *Server) cacheDurableMessage(chatID uuid.UUID, event codersdk.ChatStream
 		if evicted := state.durableMessages[0]; evicted.Message != nil {
 			state.durableEvictedBefore = evicted.Message.ID
 		}
+		// Zero the dropped slot so the evicted *ChatMessage is
+		// GC-eligible; see publishToStream for the same pattern.
+		state.durableMessages[0] = codersdk.ChatStreamEvent{}
 		state.durableMessages = state.durableMessages[1:]
 	}
 	state.durableMessages = append(state.durableMessages, event)
@@ -3151,25 +3166,96 @@ func (p *Server) getOrCreateStreamState(chatID uuid.UUID) *chatStreamState {
 	return state
 }
 
-// cleanupStreamIfIdle removes the chat entry from the sync.Map
-// when there are no subscribers and the stream is not buffering.
-// When bufferRetainedAt is set, cleanup is deferred until the
-// grace period expires so cross-replica relay subscribers can
-// still snapshot the buffer.
-// The caller must hold state.mu.
-func (p *Server) cleanupStreamIfIdle(chatID uuid.UUID, state *chatStreamState) {
+// cleanupStreamIfIdle removes the chat entry from the sync.Map when
+// there are no subscribers, the stream is not buffering, and any
+// grace period for late-connecting relay subscribers has elapsed. If
+// the grace window is still open it returns without rescheduling.
+// streamJanitorLoop is the backstop that re-checks on a timer.
+//
+// The caller must hold state.mu. The state pointer may have been
+// captured outside this lock (sync.Map.Load or Range); we use
+// CompareAndDelete so a stale pointer cannot evict a fresh entry
+// installed by a racing getOrCreateStreamState. Returns true
+// if the state was deleted, false otherwise.
+func (p *Server) cleanupStreamIfIdle(chatID uuid.UUID, state *chatStreamState) bool {
 	if state.buffering || len(state.subscribers) > 0 {
-		return
+		return false
 	}
 	// Keep stream state alive during the grace period so
 	// late-connecting relay subscribers can snapshot the
 	// buffer after the worker finishes processing.
 	if !state.bufferRetainedAt.IsZero() &&
 		p.clock.Now().Before(state.bufferRetainedAt.Add(bufferRetainGracePeriod)) {
-		return
+		return false
+	}
+	if !p.chatStreams.CompareAndDelete(chatID, state) {
+		return false
 	}
-	p.chatStreams.Delete(chatID)
 	p.workspaceMCPToolsCache.Delete(chatID)
+	return true
+}
+
+// streamJanitorLoop periodically reaps idle chat stream states whose
+// grace period has expired. It is the backstop for the grace-window
+// early-return in cleanupStreamIfIdle; without it, a subscriber that
+// detaches inside grace (the common enterprise relay-drain case,
+// relayDrainTimeout = 200ms vs. 5s grace) pins the state forever.
+func (p *Server) streamJanitorLoop(ctx context.Context) {
+	ticker := p.clock.NewTicker(streamJanitorInterval, "chatd", "stream-janitor")
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			p.safeSweepIdleStreams(ctx)
+		}
+	}
+}
+
+// safeSweepIdleStreams runs sweepIdleStreams under a panic recovery
+// so an unexpected panic in the sweep cannot kill the janitor
+// goroutine and silently reintroduce the very leak it exists to
+// prevent. The next tick retries.
+func (p *Server) safeSweepIdleStreams(ctx context.Context) {
+	defer func() {
+		if r := recover(); r != nil {
+			p.logger.Error(ctx, "stream janitor sweep panicked, will retry next tick",
+				slog.F("panic", r))
+		}
+	}()
+	p.sweepIdleStreams()
+}
+
+// sweepIdleStreams iterates chatStreams once and delegates each entry
+// to cleanupStreamIfIdle. Range may skip entries that become reapable
+// concurrently. Any such entry is reaped on the next tick.
+func (p *Server) sweepIdleStreams() {
+	var reaped atomic.Int64
+	defer func() {
+		if count := reaped.Load(); count > 0 {
+			p.logger.Info(context.Background(), "reaped idle chat streams", slog.F("count", count))
+		}
+	}()
+	p.chatStreams.Range(func(key, value any) bool {
+		chatID, ok := key.(uuid.UUID)
+		if !ok {
+			return true
+		}
+		state, ok := value.(*chatStreamState)
+		if !ok {
+			return true
+		}
+		// guard against any panic from cleanupStreamIfIdle locking state.mu for all time
+		func() {
+			state.mu.Lock()
+			defer state.mu.Unlock()
+			if p.cleanupStreamIfIdle(chatID, state) {
+				reaped.Add(1)
+			}
+		}()
+		return true
+	})
 }
 
 // registerHeartbeat enrolls a chat in the centralized batch

coderd/x/chatd/chatd_internal_test.go

@@ -3046,3 +3046,319 @@ func TestHeartbeatTick_DBErrorDoesNotInterruptChats(t *testing.T) {
 	require.NoError(t, chatCtx.Err(),
 		"chat context should not be canceled on transient DB error")
 }
+
+// TestSubscribeCancelDuringGrace_ReapedBySweep verifies that a
+// subscriber detach inside bufferRetainGracePeriod (the OSS trigger
+// for the retained-buffer leak) leaves the state mapped, and the
+// next sweep past the grace window reaps it.
+func TestSubscribeCancelDuringGrace_ReapedBySweep(t *testing.T) {
+	t.Parallel()
+
+	logger := slogtest.Make(t, nil)
+	mClock := quartz.NewMock(t)
+
+	server := &Server{
+		logger: logger,
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+	start := mClock.Now()
+
+	// Just-finished chat: processing done, buffer retained for
+	// late-connecting relay subscribers.
+	state := &chatStreamState{
+		buffering:        false,
+		bufferRetainedAt: start,
+		subscribers:      map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+		buffer: []codersdk.ChatStreamEvent{{
+			Type: codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{
+				Role: codersdk.ChatMessageRoleAssistant,
+			},
+		}},
+	}
+	server.chatStreams.Store(chatID, state)
+
+	// Real subscribeToStream cancel path: the WS subscriber detach
+	// that leaks in prod.
+	_, _, cancelSub := server.subscribeToStream(chatID)
+
+	mClock.Advance(bufferRetainGracePeriod / 2)
+	cancelSub()
+
+	_, ok := server.chatStreams.Load(chatID)
+	require.True(t, ok,
+		"entry should remain during grace window after subscriber detach")
+
+	mClock.Advance(bufferRetainGracePeriod)
+	server.sweepIdleStreams()
+
+	_, ok = server.chatStreams.Load(chatID)
+	require.False(t, ok,
+		"entry should be reaped after grace period expires and sweep runs")
+}
+
+// TestSweepIdleStreams_ReapsStaleRetainedBuffer: grace expired, no
+// subscribers, not buffering -> reaped.
+func TestSweepIdleStreams_ReapsStaleRetainedBuffer(t *testing.T) {
+	t.Parallel()
+
+	mClock := quartz.NewMock(t)
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+	state := &chatStreamState{
+		buffering:        false,
+		bufferRetainedAt: mClock.Now(),
+		subscribers:      map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+		buffer: []codersdk.ChatStreamEvent{{
+			Type:        codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{},
+		}},
+	}
+	server.chatStreams.Store(chatID, state)
+
+	mClock.Advance(bufferRetainGracePeriod + time.Second)
+	server.sweepIdleStreams()
+
+	_, ok := server.chatStreams.Load(chatID)
+	require.False(t, ok, "stale retained state should be reaped")
+}
+
+// TestSweepIdleStreams_DoesNotReapActiveBuffering: buffering=true
+// blocks reap even long after any grace would have expired.
+func TestSweepIdleStreams_DoesNotReapActiveBuffering(t *testing.T) {
+	t.Parallel()
+
+	mClock := quartz.NewMock(t)
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+	state := &chatStreamState{
+		buffering:   true,
+		subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+		buffer: []codersdk.ChatStreamEvent{{
+			Type:        codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{},
+		}},
+	}
+	server.chatStreams.Store(chatID, state)
+
+	mClock.Advance(time.Hour)
+	server.sweepIdleStreams()
+
+	_, ok := server.chatStreams.Load(chatID)
+	require.True(t, ok, "actively-buffering state must not be reaped")
+}
+
+// TestSweepIdleStreams_DoesNotReapWithSubscribers: attached
+// subscribers block reap even when grace has expired.
+func TestSweepIdleStreams_DoesNotReapWithSubscribers(t *testing.T) {
+	t.Parallel()
+
+	mClock := quartz.NewMock(t)
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+	state := &chatStreamState{
+		buffering:        false,
+		bufferRetainedAt: mClock.Now(),
+		subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{
+			uuid.New(): make(chan codersdk.ChatStreamEvent, 1),
+		},
+		buffer: []codersdk.ChatStreamEvent{{
+			Type:        codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{},
+		}},
+	}
+	server.chatStreams.Store(chatID, state)
+
+	mClock.Advance(bufferRetainGracePeriod + time.Second)
+	server.sweepIdleStreams()
+
+	_, ok := server.chatStreams.Load(chatID)
+	require.True(t, ok, "state with subscribers must not be reaped")
+}
+
+// TestSweepIdleStreams_DefersDuringGracePeriod: sweep inside grace
+// is a no-op; the next sweep past grace reaps.
+func TestSweepIdleStreams_DefersDuringGracePeriod(t *testing.T) {
+	t.Parallel()
+
+	mClock := quartz.NewMock(t)
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+	start := mClock.Now()
+	state := &chatStreamState{
+		buffering:        false,
+		bufferRetainedAt: start,
+		subscribers:      map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+		buffer: []codersdk.ChatStreamEvent{{
+			Type:        codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{},
+		}},
+	}
+	server.chatStreams.Store(chatID, state)
+
+	mClock.Advance(bufferRetainGracePeriod / 2)
+	server.sweepIdleStreams()
+
+	_, ok := server.chatStreams.Load(chatID)
+	require.True(t, ok, "sweep inside grace window must not reap")
+
+	mClock.Advance(bufferRetainGracePeriod)
+	server.sweepIdleStreams()
+
+	_, ok = server.chatStreams.Load(chatID)
+	require.False(t, ok, "sweep after grace window must reap")
+}
+
+// TestPublishToStream_DropZeroesBackingSlot verifies that evicting
+// the oldest buffered event at capacity zeroes the dropped slot so
+// its *ChatStreamMessagePart becomes GC-eligible immediately.
+func TestPublishToStream_DropZeroesBackingSlot(t *testing.T) {
+	t.Parallel()
+
+	mClock := quartz.NewMock(t)
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+
+	// Over-allocate by one so the post-drop append fits in place and
+	// exercises the backing-array reuse this test is checking.
+	buf := make([]codersdk.ChatStreamEvent, maxStreamBufferSize, maxStreamBufferSize+1)
+	for i := range buf {
+		buf[i] = codersdk.ChatStreamEvent{
+			Type:        codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{},
+		}
+	}
+	// Sentinel in slot 0 distinguishes "slot was zeroed" from "slot
+	// was overwritten by a later append".
+	sentinel := &codersdk.ChatStreamMessagePart{
+		Role: codersdk.ChatMessageRoleAssistant,
+	}
+	buf[0] = codersdk.ChatStreamEvent{
+		Type:        codersdk.ChatStreamEventTypeMessagePart,
+		MessagePart: sentinel,
+	}
+	// Alias over the full backing array so we can still observe slot
+	// 0 after publishToStream reslices state.buffer forward.
+	origBacking := buf[:cap(buf)]
+
+	state := &chatStreamState{
+		buffering:   true,
+		buffer:      buf,
+		subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+	}
+	server.chatStreams.Store(chatID, state)
+
+	newPart := &codersdk.ChatStreamMessagePart{
+		Role: codersdk.ChatMessageRoleAssistant,
+	}
+	server.publishToStream(chatID, codersdk.ChatStreamEvent{
+		Type:        codersdk.ChatStreamEventTypeMessagePart,
+		MessagePart: newPart,
+	})
+
+	require.Equal(t, codersdk.ChatStreamEvent{}, origBacking[0],
+		"dropped slot must be zero-valued so its *ChatStreamMessagePart "+
+			"is eligible for GC; got %+v", origBacking[0])
+
+	// Sanity-check the in-place append path the fix targets: if Go's
+	// growth policy ever makes this append reallocate, this fails
+	// loudly so the test author revisits the setup.
+	require.Same(t, newPart, origBacking[len(origBacking)-1].MessagePart,
+		"append must have landed in the original backing array; the "+
+			"zero-out invariant only matters when cap > len")
+}
+
+// TestCleanupStreamIfIdle_StalePointerDoesNotDeleteFreshEntry covers
+// the race where a caller holds a pointer to a no-longer-mapped
+// state (e.g. a janitor Range callback racing a fresh
+// getOrCreateStreamState) and would otherwise evict the fresh entry.
+// With CompareAndDelete in cleanupStreamIfIdle the stale delete is
+// a no-op.
+func TestCleanupStreamIfIdle_StalePointerDoesNotDeleteFreshEntry(t *testing.T) {
+	t.Parallel()
+
+	mClock := quartz.NewMock(t)
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+
+	// Stale pointer: reapable (not buffering, no subscribers, grace
+	// expired) but no longer the map's live entry.
+	stale := &chatStreamState{
+		buffering:        false,
+		bufferRetainedAt: mClock.Now(),
+		subscribers:      map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+	}
+
+	// Fresh entry: the state getOrCreateStreamState would install
+	// after a racing processChat run. Actively buffering, so not
+	// reapable. Only this state is in the map.
+	fresh := &chatStreamState{
+		buffering:   true,
+		subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+	}
+	server.chatStreams.Store(chatID, fresh)
+
+	mClock.Advance(bufferRetainGracePeriod + time.Second)
+
+	// Stale caller mirrors the janitor Range callback after the map
+	// entry has already been replaced.
+	stale.mu.Lock()
+	server.cleanupStreamIfIdle(chatID, stale)
+	stale.mu.Unlock()
+
+	got, ok := server.chatStreams.Load(chatID)
+	require.True(t, ok,
+		"fresh entry must remain mapped when cleanup is called with a stale pointer")
+	require.Same(t, fresh, got,
+		"cleanup must not replace the fresh entry with the stale one")
+}
+
+// TestSafeSweepIdleStreams_RecoversFromPanic verifies that an
+// unexpected panic inside sweepIdleStreams is recovered rather than
+// killing the janitor goroutine. Without this guard, a panic would
+// silently reintroduce the very leak the janitor exists to prevent.
+func TestSafeSweepIdleStreams_RecoversFromPanic(t *testing.T) {
+	t.Parallel()
+
+	server := &Server{
+		logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
+		clock:  quartz.NewMock(t),
+	}
+
+	chatID := uuid.New()
+	// A nil *chatStreamState passes the type assertion in sweepIdleStreams
+	// but panics on state.mu.Lock with a nil-pointer deref. Any future
+	// panic source in the sweep would trigger the same recovery path.
+	var nilState *chatStreamState
+	server.chatStreams.Store(chatID, nilState)
+
+	require.NotPanics(t, func() {
+		server.safeSweepIdleStreams(context.Background())
+	}, "safeSweepIdleStreams must recover panics so the janitor loop keeps running")
+}

enterprise/coderd/x/chatd/chatd_test.go

@@ -1425,6 +1425,148 @@ func TestSubscribeRelayDialCanceledOnFastCompletion(t *testing.T) {
 			"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()
+
+	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)
+	worker := osschatd.New(osschatd.Config{
+		Logger:                     workerLogger,
+		Database:                   db,
+		ReplicaID:                  workerID,
+		Pubsub:                     ps,
+		PendingChatAcquireInterval: time.Hour,
+		InFlightChatStaleAfter:     testutil.WaitSuperLong,
+		Clock:                      workerClock,
+	})
+	t.Cleanup(func() {
+		require.NoError(t, worker.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
+	}, nil)
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	user, org, model := seedChatDependencies(ctx, t, db)
+	setOpenAIProviderBaseURL(ctx, t, db, openAIURL)
+
+	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-drain-characterization")
+
+	// 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)
+
+	_, err := worker.SendMessage(ctx, osschatd.SendMessageOptions{
+		ChatID:    chat.ID,
+		CreatedBy: user.ID,
+		Content:   []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
+	})
+	require.NoError(t, err)
+
+	// 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.
+	var committedAssistantMsgs int
+	var messagePartsSeen int
+	testutil.Eventually(ctx, 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)
+
+	testutil.Eventually(ctx, t, func(ctx context.Context) bool {
+		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.
+	testutil.Eventually(ctx, 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
 // relay is established while the worker is still streaming, the
 // subscriber receives buffered parts via the relay snapshot and live