fix(coderd/x/chatd): checkpoint buffered message_parts to avoid stale replay (#25145)

coderd/x/chatd/chatd.go

@@ -9,6 +9,7 @@ import (
 	"errors"
 	"fmt"
 	"maps"
+	"math"
 	"net/http"
 	"slices"
 	"strconv"
@@ -83,6 +84,12 @@ const (
 	// per chat during a single LLM step. When exceeded the oldest event is
 	// evicted so memory stays bounded.
 	maxStreamBufferSize = 10000
+	// RelaySentinelAfterID is the after_id sentinel used by cross-replica
+	// relay subscribers. It instructs the peer to skip the durable DB
+	// snapshot and deliver only in-flight buffered parts. The sentinel
+	// also disables snapshotBufferLocked's redundant-part filter so
+	// relays receive every part the worker has buffered (see PR #24031).
+	RelaySentinelAfterID = math.MaxInt64
 	// maxDurableMessageCacheSize caps the number of recent durable message
 	// events cached per chat for same-replica stream catch-up.
 	maxDurableMessageCacheSize = 256
@@ -1092,9 +1099,26 @@ type SubscribeFnParams struct {
 	Logger              slog.Logger
 }
 
+// bufferedStreamPart is a buffered message_part event tagged with the
+// most recently committed assistant message ID at the moment it was
+// appended. Subscribers can use the checkpoint to skip parts that
+// belong to turns they have already received via durable
+// `message` events.
+type bufferedStreamPart struct {
+	event codersdk.ChatStreamEvent
+	// checkpoint is the chatStreamState.lastCommittedAssistantMessageID
+	// value at the time this part was buffered. A subscriber whose
+	// cursor is past this checkpoint already has the durable assistant
+	// message for the turn this part belongs to, so the part is
+	// redundant. The cursor is clamped to the current checkpoint at
+	// snapshot time, so tool/user message IDs in the cursor cannot
+	// over-drop parts from an in-progress assistant turn.
+	checkpoint int64
+}
+
 type chatStreamState struct {
 	mu                   sync.Mutex
-	buffer               []codersdk.ChatStreamEvent
+	buffer               []bufferedStreamPart
 	buffering            bool
 	durableMessages      []codersdk.ChatStreamEvent
 	durableEvictedBefore int64 // highest message ID evicted from durable cache
@@ -1114,6 +1138,12 @@ type chatStreamState struct {
 	// period expires so cross-replica relays can still
 	// snapshot the buffer.
 	bufferRetainedAt time.Time
+	// lastCommittedAssistantMessageID tracks the highest assistant
+	// durable message ID published for this chat on this replica.
+	// publishToStream tags each buffered message_part with this
+	// value so subscribeToStream can filter out parts belonging to
+	// already-committed turns.
+	lastCommittedAssistantMessageID int64
 }
 
 // heartbeatEntry tracks a single chat's cancel function and workspace
@@ -4144,10 +4174,13 @@ func (p *Server) publishToStream(chatID uuid.UUID, event codersdk.ChatStreamEven
 			// 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[0] = bufferedStreamPart{}
 			state.buffer = state.buffer[1:]
 		}
-		state.buffer = append(state.buffer, event)
+		state.buffer = append(state.buffer, bufferedStreamPart{
+			event:      event,
+			checkpoint: state.lastCommittedAssistantMessageID,
+		})
 	}
 	subscribers := make([]chan codersdk.ChatStreamEvent, 0, len(state.subscribers))
 	for _, ch := range state.subscribers {
@@ -4230,7 +4263,78 @@ func (p *Server) getCachedDurableMessages(
 	return result
 }
 
-func (p *Server) subscribeToStream(chatID uuid.UUID) (
+// snapshotBufferLocked returns the buffered message_part events that
+// the caller should receive in their initial snapshot, filtered by
+// the requested cursor.
+//
+// The cursor is clamped to lastCommittedAssistantMessageID so a
+// cursor that points at a tool or user message ID past the last
+// committed assistant turn cannot drop parts from the in-progress
+// turn. The filter then drops parts whose checkpoint is below the
+// clamped cursor; those parts belong to assistant turns the
+// subscriber already has via durable `message` events.
+//
+// The caller must hold the per-chat stream state lock. See
+// subscribeToStream for the documented afterMessageID semantics.
+func snapshotBufferLocked(
+	buffer []bufferedStreamPart,
+	afterMessageID int64,
+	lastCommittedAssistantMessageID int64,
+) []codersdk.ChatStreamEvent {
+	if len(buffer) == 0 {
+		return nil
+	}
+	// Compute the effective cursor used to drop redundant parts.
+	//   - afterMessageID <= 0 ("no cursor; deliver everything") and
+	//     the RelaySentinelAfterID both disable filtering.
+	//   - Otherwise clamp the cursor to lastCommittedAssistantMessageID
+	//     so a tool/user cursor past the last assistant turn cannot
+	//     over-drop parts from the in-progress assistant turn. We can
+	//     only be confident a buffered part is redundant when the
+	//     cursor is at or past its checkpoint AND the checkpoint maps
+	//     to a turn the subscriber already has via durable messages.
+	//   - If lastCommittedAssistantMessageID is still zero (e.g.
+	//     fresh state after cleanup), no buffered part can be proven
+	//     redundant, so deliver everything.
+	var effectiveCursor int64
+	switch {
+	case afterMessageID <= 0, afterMessageID == RelaySentinelAfterID:
+		effectiveCursor = 0
+	case lastCommittedAssistantMessageID < afterMessageID:
+		effectiveCursor = lastCommittedAssistantMessageID
+	default:
+		effectiveCursor = afterMessageID
+	}
+	snapshot := make([]codersdk.ChatStreamEvent, 0, len(buffer))
+	for _, part := range buffer {
+		if effectiveCursor > 0 && part.checkpoint < effectiveCursor {
+			continue
+		}
+		snapshot = append(snapshot, part.event)
+	}
+	return snapshot
+}
+
+// subscribeToStream registers a subscriber to the per-chat in-memory
+// stream and returns a filtered snapshot of currently-buffered
+// message_part events plus the current retry phase, the live
+// subscriber channel, and a cancel func.
+//
+// afterMessageID semantics:
+//   - 0: no filter; the full buffer snapshot is returned.
+//     New browser sessions use this and only see parts for the
+//     currently-streaming turn (the buffer is cleared at the
+//     start of each processChat run).
+//   - RelaySentinelAfterID: no filter; cross-replica relays pass
+//     this sentinel to skip the durable DB snapshot while still
+//     receiving all in-flight buffered parts.
+//   - 0 < afterMessageID < RelaySentinelAfterID: parts whose
+//     checkpoint is less than the cursor are dropped from the
+//     snapshot. The cursor is clamped to the per-chat
+//     lastCommittedAssistantMessageID before filtering so cursors
+//     that point at tool/user message IDs past the last committed
+//     assistant turn cannot over-drop in-progress parts.
+func (p *Server) subscribeToStream(chatID uuid.UUID, afterMessageID int64) (
 	[]codersdk.ChatStreamEvent,
 	*codersdk.ChatStreamRetry,
 	<-chan codersdk.ChatStreamEvent,
@@ -4238,7 +4342,7 @@ func (p *Server) subscribeToStream(chatID uuid.UUID) (
 ) {
 	state := p.getOrCreateStreamState(chatID)
 	state.mu.Lock()
-	snapshot := append([]codersdk.ChatStreamEvent(nil), state.buffer...)
+	snapshot := snapshotBufferLocked(state.buffer, afterMessageID, state.lastCommittedAssistantMessageID)
 	var currentRetry *codersdk.ChatStreamRetry
 	if state.currentRetry != nil {
 		retryCopy := *state.currentRetry
@@ -4541,7 +4645,7 @@ func (p *Server) SubscribeAuthorized(
 	// persisted messages. Capture the current retry phase under the same
 	// lock so the transient snapshot and subscriber registration reflect
 	// a single moment in time.
-	localSnapshot, localRetry, localParts, localCancel := p.subscribeToStream(chatID)
+	localSnapshot, localRetry, localParts, localCancel := p.subscribeToStream(chatID, afterMessageID)
 
 	// Merge all event sources.
 	mergedCtx, mergedCancel := context.WithCancel(ctx)
@@ -5222,12 +5326,87 @@ func (p *Server) publishMessage(chatID uuid.UUID, message database.ChatMessage)
 		Message: &sdkMessage,
 	}
 	p.cacheDurableMessage(chatID, event)
+	p.advanceAssistantCheckpoint(chatID, message)
 	p.publishEvent(chatID, event)
 	p.publishChatStreamNotify(chatID, coderdpubsub.ChatStreamNotifyMessage{
 		AfterMessageID: message.ID - 1,
 	})
 }
 
+// seedAssistantCheckpoint initializes the per-chat checkpoint from
+// the last durable assistant message ID before any parts are
+// buffered for this run. This closes the cleanup-and-recreate race
+// where a freshly stored chatStreamState would start with
+// lastCommittedAssistantMessageID = 0, tagging every part with
+// checkpoint 0 and forcing snapshotBufferLocked to deliver the
+// entire buffer to every reconnecting subscriber.
+//
+// On lookup error or when there is no prior assistant message, the
+// checkpoint stays at its current value (either zero for a brand-new
+// state, or the value carried forward from a prior run on this
+// replica). This is safe: a zero checkpoint produces an over-
+// inclusive snapshot, not data loss.
+func (p *Server) seedAssistantCheckpoint(
+	ctx context.Context,
+	chatID uuid.UUID,
+	state *chatStreamState,
+	logger slog.Logger,
+) {
+	// Use a short timeout so a stalled DB does not block the
+	// start of a chat run. The seed is best-effort: missing it
+	// only degrades the snapshot filter to "deliver everything",
+	// which is the prior behavior.
+	//
+	// The seed reads the last assistant message ID to bound the
+	// in-memory checkpoint; it never returns user data. The
+	// system context is required because processChat runs without
+	// an actor and the durable read is part of the chat worker
+	// loop. There is no authorization to skip; the chat row was
+	// already authorized before processChat was scheduled.
+	//nolint:gocritic // chatd worker reads its own durable state to seed the in-memory checkpoint; no user context exists here.
+	lookupCtx, cancel := context.WithTimeout(
+		dbauthz.AsSystemRestricted(ctx),
+		5*time.Second,
+	)
+	defer cancel()
+	last, err := p.db.GetLastChatMessageByRole(lookupCtx, database.GetLastChatMessageByRoleParams{
+		ChatID: chatID,
+		Role:   database.ChatMessageRoleAssistant,
+	})
+	if errors.Is(err, sql.ErrNoRows) {
+		return
+	}
+	if err != nil {
+		logger.Warn(ctx, "failed to seed assistant checkpoint", slog.Error(err))
+		return
+	}
+	state.mu.Lock()
+	defer state.mu.Unlock()
+	if last.ID > state.lastCommittedAssistantMessageID {
+		state.lastCommittedAssistantMessageID = last.ID
+	}
+}
+
+// advanceAssistantCheckpoint bumps the per-chat checkpoint when an
+// assistant durable message is published. Subsequent buffered
+// message_part events are tagged with the new checkpoint so
+// subscribeToStream can filter parts belonging to already-committed
+// turns when the subscriber's cursor is past the checkpoint.
+//
+// Tool and user messages do not end an assistant streaming turn, so
+// the checkpoint is only advanced for assistant-role messages.
+func (p *Server) advanceAssistantCheckpoint(chatID uuid.UUID, message database.ChatMessage) {
+	if message.Role != database.ChatMessageRoleAssistant {
+		return
+	}
+	state := p.getOrCreateStreamState(chatID)
+	state.mu.Lock()
+	defer state.mu.Unlock()
+	if message.ID > state.lastCommittedAssistantMessageID {
+		state.lastCommittedAssistantMessageID = message.ID
+	}
+}
+
 // publishEditedMessage is like publishMessage but uses FullRefresh
 // so remote subscribers re-fetch from the beginning, ensuring the
 // edit is never silently dropped. The durable cache is replaced
@@ -5678,7 +5857,19 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) {
 	streamState.bufferRetainedAt = time.Time{}
 	streamState.resetDropCounters()
 	streamState.buffering = true
+	// lastCommittedAssistantMessageID is intentionally NOT reset
+	// here: the checkpoint is lifetime-scoped across runs so that
+	// after a state was reaped and a new run starts, reconnecting
+	// subscribers can still filter parts from prior turns once we
+	// re-seed it below.
 	streamState.mu.Unlock()
+	// Seed the checkpoint from the durable store so that after a
+	// cleanupStreamIfIdle reaped the previous state, the very
+	// first parts buffered by this run are not tagged with
+	// checkpoint=0 (which would make snapshotBufferLocked deliver
+	// the full buffer to every reconnecting subscriber regardless
+	// of their cursor).
+	p.seedAssistantCheckpoint(ctx, chat.ID, streamState, logger)
 	defer func() {
 		streamState.mu.Lock()
 		// Fallback cleanup for exit paths that return before a

coderd/x/chatd/chatd_internal_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"encoding/json"
+	"math"
 	"sync"
 	"testing"
 	"time"
@@ -2511,12 +2512,14 @@ func TestSubscribeAuthorizedFallsBackToStaleRowWhenRefreshFails(t *testing.T) {
 
 	state := server.getOrCreateStreamState(chatID)
 	state.mu.Lock()
-	state.buffer = []codersdk.ChatStreamEvent{{
-		Type:   codersdk.ChatStreamEventTypeMessagePart,
-		ChatID: chatID,
-		MessagePart: &codersdk.ChatStreamMessagePart{
-			Role: "assistant",
-			Part: codersdk.ChatMessageText("thinking"),
+	state.buffer = []bufferedStreamPart{{
+		event: codersdk.ChatStreamEvent{
+			Type:   codersdk.ChatStreamEventTypeMessagePart,
+			ChatID: chatID,
+			MessagePart: &codersdk.ChatStreamMessagePart{
+				Role: "assistant",
+				Part: codersdk.ChatMessageText("thinking"),
+			},
 		},
 	}}
 	state.mu.Unlock()
@@ -2734,7 +2737,7 @@ func TestPublishToStream_DropWarnRateLimiting(t *testing.T) {
 	// buffering enabled, one saturated subscriber.
 	state := &chatStreamState{
 		buffering: true,
-		buffer:    make([]codersdk.ChatStreamEvent, maxStreamBufferSize),
+		buffer:    make([]bufferedStreamPart, maxStreamBufferSize),
 		subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{
 			uuid.New(): subCh,
 		},
@@ -2785,7 +2788,7 @@ func TestPublishToStream_DropWarnRateLimiting(t *testing.T) {
 
 	// --- Phase 3: counter reset (simulates step persist) ---
 	state.mu.Lock()
-	state.buffer = make([]codersdk.ChatStreamEvent, maxStreamBufferSize)
+	state.buffer = make([]bufferedStreamPart, maxStreamBufferSize)
 	state.resetDropCounters()
 	state.mu.Unlock()
 
@@ -3500,6 +3503,9 @@ func TestProcessChat_IgnoresStaleControlNotification(t *testing.T) {
 		database.ChatUsageLimitConfig{}, sql.ErrNoRows,
 	).AnyTimes()
 	db.EXPECT().GetChatMessagesForPromptByChatID(gomock.Any(), chatID).Return(nil, nil).AnyTimes()
+	db.EXPECT().GetLastChatMessageByRole(gomock.Any(), gomock.Any()).Return(
+		database.ChatMessage{}, sql.ErrNoRows,
+	).AnyTimes()
 
 	chat := database.Chat{ID: chatID, LastModelConfigID: uuid.New()}
 	done := make(chan struct{})
@@ -3739,10 +3745,12 @@ func TestSubscribeCancelDuringGrace_ReapedBySweep(t *testing.T) {
 		buffering:        false,
 		bufferRetainedAt: start,
 		subscribers:      map[uuid.UUID]chan codersdk.ChatStreamEvent{},
-		buffer: []codersdk.ChatStreamEvent{{
-			Type: codersdk.ChatStreamEventTypeMessagePart,
-			MessagePart: &codersdk.ChatStreamMessagePart{
-				Role: codersdk.ChatMessageRoleAssistant,
+		buffer: []bufferedStreamPart{{
+			event: codersdk.ChatStreamEvent{
+				Type: codersdk.ChatStreamEventTypeMessagePart,
+				MessagePart: &codersdk.ChatStreamMessagePart{
+					Role: codersdk.ChatMessageRoleAssistant,
+				},
 			},
 		}},
 	}
@@ -3750,7 +3758,7 @@ func TestSubscribeCancelDuringGrace_ReapedBySweep(t *testing.T) {
 
 	// Real subscribeToStream cancel path: the WS subscriber detach
 	// that leaks in prod.
-	snapshot, currentRetry, events, cancelSub := server.subscribeToStream(chatID)
+	snapshot, currentRetry, events, cancelSub := server.subscribeToStream(chatID, 0)
 	require.Len(t, snapshot, 1)
 	require.Nil(t, currentRetry)
 	require.NotNil(t, events)
@@ -3786,9 +3794,11 @@ func TestSweepIdleStreams_ReapsStaleRetainedBuffer(t *testing.T) {
 		buffering:        false,
 		bufferRetainedAt: mClock.Now(),
 		subscribers:      map[uuid.UUID]chan codersdk.ChatStreamEvent{},
-		buffer: []codersdk.ChatStreamEvent{{
-			Type:        codersdk.ChatStreamEventTypeMessagePart,
-			MessagePart: &codersdk.ChatStreamMessagePart{},
+		buffer: []bufferedStreamPart{{
+			event: codersdk.ChatStreamEvent{
+				Type:        codersdk.ChatStreamEventTypeMessagePart,
+				MessagePart: &codersdk.ChatStreamMessagePart{},
+			},
 		}},
 	}
 	server.chatStreams.Store(chatID, state)
@@ -3815,9 +3825,11 @@ func TestSweepIdleStreams_DoesNotReapActiveBuffering(t *testing.T) {
 	state := &chatStreamState{
 		buffering:   true,
 		subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{},
-		buffer: []codersdk.ChatStreamEvent{{
-			Type:        codersdk.ChatStreamEventTypeMessagePart,
-			MessagePart: &codersdk.ChatStreamMessagePart{},
+		buffer: []bufferedStreamPart{{
+			event: codersdk.ChatStreamEvent{
+				Type:        codersdk.ChatStreamEventTypeMessagePart,
+				MessagePart: &codersdk.ChatStreamMessagePart{},
+			},
 		}},
 	}
 	server.chatStreams.Store(chatID, state)
@@ -3847,9 +3859,11 @@ func TestSweepIdleStreams_DoesNotReapWithSubscribers(t *testing.T) {
 		subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{
 			uuid.New(): make(chan codersdk.ChatStreamEvent, 1),
 		},
-		buffer: []codersdk.ChatStreamEvent{{
-			Type:        codersdk.ChatStreamEventTypeMessagePart,
-			MessagePart: &codersdk.ChatStreamMessagePart{},
+		buffer: []bufferedStreamPart{{
+			event: codersdk.ChatStreamEvent{
+				Type:        codersdk.ChatStreamEventTypeMessagePart,
+				MessagePart: &codersdk.ChatStreamMessagePart{},
+			},
 		}},
 	}
 	server.chatStreams.Store(chatID, state)
@@ -3878,9 +3892,11 @@ func TestSweepIdleStreams_DefersDuringGracePeriod(t *testing.T) {
 		buffering:        false,
 		bufferRetainedAt: start,
 		subscribers:      map[uuid.UUID]chan codersdk.ChatStreamEvent{},
-		buffer: []codersdk.ChatStreamEvent{{
-			Type:        codersdk.ChatStreamEventTypeMessagePart,
-			MessagePart: &codersdk.ChatStreamMessagePart{},
+		buffer: []bufferedStreamPart{{
+			event: codersdk.ChatStreamEvent{
+				Type:        codersdk.ChatStreamEventTypeMessagePart,
+				MessagePart: &codersdk.ChatStreamMessagePart{},
+			},
 		}},
 	}
 	server.chatStreams.Store(chatID, state)
@@ -3914,11 +3930,13 @@ func TestPublishToStream_DropZeroesBackingSlot(t *testing.T) {
 
 	// 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)
+	buf := make([]bufferedStreamPart, maxStreamBufferSize, maxStreamBufferSize+1)
 	for i := range buf {
-		buf[i] = codersdk.ChatStreamEvent{
-			Type:        codersdk.ChatStreamEventTypeMessagePart,
-			MessagePart: &codersdk.ChatStreamMessagePart{},
+		buf[i] = bufferedStreamPart{
+			event: codersdk.ChatStreamEvent{
+				Type:        codersdk.ChatStreamEventTypeMessagePart,
+				MessagePart: &codersdk.ChatStreamMessagePart{},
+			},
 		}
 	}
 	// Sentinel in slot 0 distinguishes "slot was zeroed" from "slot
@@ -3926,9 +3944,11 @@ func TestPublishToStream_DropZeroesBackingSlot(t *testing.T) {
 	sentinel := &codersdk.ChatStreamMessagePart{
 		Role: codersdk.ChatMessageRoleAssistant,
 	}
-	buf[0] = codersdk.ChatStreamEvent{
-		Type:        codersdk.ChatStreamEventTypeMessagePart,
-		MessagePart: sentinel,
+	buf[0] = bufferedStreamPart{
+		event: 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.
@@ -3949,14 +3969,14 @@ func TestPublishToStream_DropZeroesBackingSlot(t *testing.T) {
 		MessagePart: newPart,
 	})
 
-	require.Equal(t, codersdk.ChatStreamEvent{}, origBacking[0],
+	require.Equal(t, bufferedStreamPart{}, 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,
+	require.Same(t, newPart, origBacking[len(origBacking)-1].event.MessagePart,
 		"append must have landed in the original backing array; the "+
 			"zero-out invariant only matters when cap > len")
 }
@@ -5260,6 +5280,9 @@ func TestAutoPromote_InsertFailureSkipsStatusUpdate(t *testing.T) {
 		database.ChatUsageLimitConfig{}, sql.ErrNoRows,
 	).AnyTimes()
 	db.EXPECT().GetChatMessagesForPromptByChatID(gomock.Any(), chatID).Return(nil, nil).AnyTimes()
+	db.EXPECT().GetLastChatMessageByRole(gomock.Any(), gomock.Any()).Return(
+		database.ChatMessage{}, sql.ErrNoRows,
+	).AnyTimes()
 
 	// The deferred cleanup transaction: InsertChatMessages fails,
 	// so UpdateChatStatus must NOT be called.
@@ -5343,3 +5366,476 @@ func TestAutoPromote_InsertFailureSkipsStatusUpdate(t *testing.T) {
 		// No signal, as expected.
 	}
 }
+
+// makeBufferedPart is a small constructor for buffered message_part
+// fixtures used by snapshotBufferLocked / subscribeToStream tests. It
+// embeds the checkpoint and a recognizable text body so failing
+// assertions can identify which part survived the filter.
+func makeBufferedPart(checkpoint int64, text string) bufferedStreamPart {
+	return bufferedStreamPart{
+		event: codersdk.ChatStreamEvent{
+			Type: codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{
+				Role: codersdk.ChatMessageRoleAssistant,
+				Part: codersdk.ChatMessageText(text),
+			},
+		},
+		checkpoint: checkpoint,
+	}
+}
+
+func partText(event codersdk.ChatStreamEvent) string {
+	if event.MessagePart == nil {
+		return ""
+	}
+	return event.MessagePart.Part.Text
+}
+
+// TestSnapshotBufferLocked_FiltersStaleParts is the core contract:
+// when a subscriber passes a real cursor, parts whose checkpoint is
+// less than the cursor are dropped from the snapshot. Parts at or
+// past the cursor are delivered.
+func TestSnapshotBufferLocked_FiltersStaleParts(t *testing.T) {
+	t.Parallel()
+
+	buffer := []bufferedStreamPart{
+		makeBufferedPart(10, "stale-1"),
+		makeBufferedPart(10, "stale-2"),
+		makeBufferedPart(20, "boundary-1"),
+		makeBufferedPart(20, "boundary-2"),
+		makeBufferedPart(30, "fresh-1"),
+	}
+
+	// Cursor matches a real assistant checkpoint, so the effective
+	// cursor is the requested cursor unchanged.
+	snapshot := snapshotBufferLocked(buffer, 20, 30)
+
+	require.Len(t, snapshot, 3,
+		"only parts checkpointed at >= afterMessageID should be kept")
+	require.Equal(t, "boundary-1", partText(snapshot[0]))
+	require.Equal(t, "boundary-2", partText(snapshot[1]))
+	require.Equal(t, "fresh-1", partText(snapshot[2]))
+}
+
+// TestSnapshotBufferLocked_ClampsCursorToLastCommittedCheckpoint
+// guards against DEREM-1: a subscriber whose cursor points at a
+// tool or user message ID past the most recent committed assistant
+// turn must not over-drop parts from the in-progress assistant
+// turn. The filter clamps the cursor down to the latest assistant
+// checkpoint so those in-progress parts survive.
+func TestSnapshotBufferLocked_ClampsCursorToLastCommittedCheckpoint(t *testing.T) {
+	t.Parallel()
+
+	// Turn A committed at assistant message 100, then tool
+	// messages 101..103 followed. Turn B is now streaming and its
+	// parts are tagged with checkpoint=100 (no new assistant turn
+	// has been committed yet on this replica).
+	buffer := []bufferedStreamPart{
+		makeBufferedPart(100, "turnB-part-1"),
+		makeBufferedPart(100, "turnB-part-2"),
+	}
+
+	// Client reloaded chat via REST and saw the latest message
+	// (a tool result at id=103), then reconnected with cursor=103.
+	// Without clamping, the filter would drop every turn B part
+	// because checkpoint (100) < afterMessageID (103).
+	snapshot := snapshotBufferLocked(buffer, 103, 100)
+
+	require.Len(t, snapshot, 2,
+		"cursor past the last assistant checkpoint must be clamped down so in-progress parts survive")
+	require.Equal(t, "turnB-part-1", partText(snapshot[0]))
+	require.Equal(t, "turnB-part-2", partText(snapshot[1]))
+}
+
+// TestSnapshotBufferLocked_ZeroCheckpointReturnsAll guards against
+// DEREM-2: a freshly created chatStreamState (after
+// cleanupStreamIfIdle reaped the previous state and seeding from DB
+// has not yet run) has lastCommittedAssistantMessageID = 0. With a
+// zero checkpoint, no buffered part can be proven redundant, so the
+// full buffer must be returned regardless of the requested cursor.
+func TestSnapshotBufferLocked_ZeroCheckpointReturnsAll(t *testing.T) {
+	t.Parallel()
+
+	buffer := []bufferedStreamPart{
+		makeBufferedPart(0, "a"),
+		makeBufferedPart(0, "b"),
+		makeBufferedPart(0, "c"),
+	}
+
+	snapshot := snapshotBufferLocked(buffer, 999, 0)
+
+	require.Len(t, snapshot, 3,
+		"lastCommittedAssistantMessageID==0 must disable the filter to avoid losing the entire in-progress turn")
+	require.Equal(t, "a", partText(snapshot[0]))
+	require.Equal(t, "b", partText(snapshot[1]))
+	require.Equal(t, "c", partText(snapshot[2]))
+}
+
+// TestSnapshotBufferLocked_ZeroCursorReturnsAll covers the
+// fresh-load convention: callers without a cursor get the full
+// buffer. Buffering is reset at the start of every processChat run,
+// so the buffer only ever contains parts from the current turn in
+// this path.
+func TestSnapshotBufferLocked_ZeroCursorReturnsAll(t *testing.T) {
+	t.Parallel()
+
+	buffer := []bufferedStreamPart{
+		makeBufferedPart(10, "a"),
+		makeBufferedPart(20, "b"),
+		makeBufferedPart(30, "c"),
+	}
+
+	snapshot := snapshotBufferLocked(buffer, 0, 30)
+
+	require.Len(t, snapshot, 3,
+		"afterMessageID == 0 means 'no cursor'; the full buffer must be returned")
+	require.Equal(t, "a", partText(snapshot[0]))
+	require.Equal(t, "b", partText(snapshot[1]))
+	require.Equal(t, "c", partText(snapshot[2]))
+}
+
+// TestSnapshotBufferLocked_RelaySentinelReturnsAll: cross-replica
+// relay dials with after_id=RelaySentinelAfterID to skip the durable
+// DB snapshot. The buffer snapshot must NOT be filtered for that
+// sentinel; otherwise the relay race PR #24031 fixed comes back.
+func TestSnapshotBufferLocked_RelaySentinelReturnsAll(t *testing.T) {
+	t.Parallel()
+
+	buffer := []bufferedStreamPart{
+		makeBufferedPart(10, "a"),
+		makeBufferedPart(20, "b"),
+		makeBufferedPart(30, "c"),
+	}
+
+	snapshot := snapshotBufferLocked(buffer, RelaySentinelAfterID, 30)
+
+	require.Len(t, snapshot, 3,
+		"the relay sentinel must NOT filter the buffer")
+	require.Equal(t, "a", partText(snapshot[0]))
+	require.Equal(t, "b", partText(snapshot[1]))
+	require.Equal(t, "c", partText(snapshot[2]))
+}
+
+// TestSnapshotBufferLocked_EmptyBufferReturnsNil documents that
+// snapshotBufferLocked returns nil (not an empty slice) for an
+// empty buffer, matching the prior append-from-nil behavior.
+func TestSnapshotBufferLocked_EmptyBufferReturnsNil(t *testing.T) {
+	t.Parallel()
+
+	require.Nil(t, snapshotBufferLocked(nil, 0, 0))
+	require.Nil(t, snapshotBufferLocked(nil, 42, 30))
+	require.Nil(t, snapshotBufferLocked([]bufferedStreamPart{}, 42, 30))
+}
+
+// TestPublishToStream_TagsPartsWithCurrentCheckpoint verifies that
+// parts buffered while the chat is streaming carry the current
+// committed-assistant-message-ID checkpoint. Subscribers can then
+// filter against this value.
+func TestPublishToStream_TagsPartsWithCurrentCheckpoint(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{},
+		lastCommittedAssistantMessageID: 100,
+	}
+	server.chatStreams.Store(chatID, state)
+
+	server.publishToStream(chatID, codersdk.ChatStreamEvent{
+		Type: codersdk.ChatStreamEventTypeMessagePart,
+		MessagePart: &codersdk.ChatStreamMessagePart{
+			Role: codersdk.ChatMessageRoleAssistant,
+			Part: codersdk.ChatMessageText("hello"),
+		},
+	})
+
+	state.mu.Lock()
+	defer state.mu.Unlock()
+	require.Len(t, state.buffer, 1)
+	require.Equal(t, int64(100), state.buffer[0].checkpoint,
+		"part must be tagged with the current checkpoint at append time")
+	require.Equal(t, "hello", partText(state.buffer[0].event))
+}
+
+// TestAdvanceAssistantCheckpoint covers the per-role behavior of
+// advanceAssistantCheckpoint:
+//   - assistant messages advance the checkpoint monotonically.
+//   - tool / user messages leave the checkpoint untouched.
+//   - older assistant IDs (out-of-order publication) do not move
+//     the checkpoint backwards.
+func TestAdvanceAssistantCheckpoint(t *testing.T) {
+	t.Parallel()
+
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  quartz.NewMock(t),
+	}
+
+	chatID := uuid.New()
+	state := server.getOrCreateStreamState(chatID)
+
+	requireCheckpoint := func(want int64) {
+		t.Helper()
+		state.mu.Lock()
+		got := state.lastCommittedAssistantMessageID
+		state.mu.Unlock()
+		require.Equal(t, want, got)
+	}
+
+	server.advanceAssistantCheckpoint(chatID, database.ChatMessage{
+		ID:   100,
+		Role: database.ChatMessageRoleAssistant,
+	})
+	requireCheckpoint(100)
+
+	server.advanceAssistantCheckpoint(chatID, database.ChatMessage{
+		ID:   200,
+		Role: database.ChatMessageRoleAssistant,
+	})
+	requireCheckpoint(200)
+
+	// Out-of-order: an older ID must not move the checkpoint
+	// backwards (defends against publish reordering).
+	server.advanceAssistantCheckpoint(chatID, database.ChatMessage{
+		ID:   150,
+		Role: database.ChatMessageRoleAssistant,
+	})
+	requireCheckpoint(200)
+
+	// Tool messages do not end an assistant streaming turn.
+	server.advanceAssistantCheckpoint(chatID, database.ChatMessage{
+		ID:   300,
+		Role: database.ChatMessageRoleTool,
+	})
+	requireCheckpoint(200)
+
+	// User messages do not end an assistant streaming turn either.
+	server.advanceAssistantCheckpoint(chatID, database.ChatMessage{
+		ID:   400,
+		Role: database.ChatMessageRoleUser,
+	})
+	requireCheckpoint(200)
+}
+
+// TestSeedAssistantCheckpoint covers the three behaviors of
+// seedAssistantCheckpoint:
+//   - success: a durable assistant message exists and its ID is
+//     installed as the checkpoint.
+//   - monotonic guard: an older ID does not move the checkpoint
+//     backwards (defends against concurrent advance from another
+//     publish path racing with the seed).
+//   - db error: a non sql.ErrNoRows failure must not change the
+//     checkpoint and must not panic.
+func TestSeedAssistantCheckpoint(t *testing.T) {
+	t.Parallel()
+
+	t.Run("InstallsLatestAssistantID", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+		ctrl := gomock.NewController(t)
+		db := dbmock.NewMockStore(ctrl)
+		server := &Server{
+			db:     db,
+			logger: slogtest.Make(t, nil),
+			clock:  quartz.NewMock(t),
+		}
+		chatID := uuid.New()
+		state := server.getOrCreateStreamState(chatID)
+
+		db.EXPECT().GetLastChatMessageByRole(gomock.Any(), database.GetLastChatMessageByRoleParams{
+			ChatID: chatID,
+			Role:   database.ChatMessageRoleAssistant,
+		}).Return(database.ChatMessage{
+			ID:   500,
+			Role: database.ChatMessageRoleAssistant,
+		}, nil)
+
+		server.seedAssistantCheckpoint(ctx, chatID, state, server.logger)
+
+		state.mu.Lock()
+		defer state.mu.Unlock()
+		require.Equal(t, int64(500), state.lastCommittedAssistantMessageID,
+			"seed must install the latest durable assistant message ID as the checkpoint")
+	})
+
+	t.Run("DoesNotMoveCheckpointBackwards", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+		ctrl := gomock.NewController(t)
+		db := dbmock.NewMockStore(ctrl)
+		server := &Server{
+			db:     db,
+			logger: slogtest.Make(t, nil),
+			clock:  quartz.NewMock(t),
+		}
+		chatID := uuid.New()
+		state := server.getOrCreateStreamState(chatID)
+		state.mu.Lock()
+		state.lastCommittedAssistantMessageID = 1000
+		state.mu.Unlock()
+
+		// DB reports an older assistant message ID. The monotonic
+		// guard must keep the existing higher checkpoint.
+		db.EXPECT().GetLastChatMessageByRole(gomock.Any(), gomock.Any()).Return(
+			database.ChatMessage{ID: 500, Role: database.ChatMessageRoleAssistant},
+			nil,
+		)
+
+		server.seedAssistantCheckpoint(ctx, chatID, state, server.logger)
+
+		state.mu.Lock()
+		defer state.mu.Unlock()
+		require.Equal(t, int64(1000), state.lastCommittedAssistantMessageID,
+			"seed must not move the checkpoint backwards")
+	})
+
+	t.Run("DBErrorLeavesCheckpointUntouched", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+		ctrl := gomock.NewController(t)
+		db := dbmock.NewMockStore(ctrl)
+		server := &Server{
+			db:     db,
+			logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
+			clock:  quartz.NewMock(t),
+		}
+		chatID := uuid.New()
+		state := server.getOrCreateStreamState(chatID)
+		state.mu.Lock()
+		state.lastCommittedAssistantMessageID = 42
+		state.mu.Unlock()
+
+		db.EXPECT().GetLastChatMessageByRole(gomock.Any(), gomock.Any()).Return(
+			database.ChatMessage{}, xerrors.New("database explode"),
+		)
+
+		server.seedAssistantCheckpoint(ctx, chatID, state, server.logger)
+
+		state.mu.Lock()
+		defer state.mu.Unlock()
+		require.Equal(t, int64(42), state.lastCommittedAssistantMessageID,
+			"a non-ErrNoRows DB error must not change the checkpoint")
+	})
+
+	t.Run("NoRowsLeavesCheckpointAtZero", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+		ctrl := gomock.NewController(t)
+		db := dbmock.NewMockStore(ctrl)
+		server := &Server{
+			db:     db,
+			logger: slogtest.Make(t, nil),
+			clock:  quartz.NewMock(t),
+		}
+		chatID := uuid.New()
+		state := server.getOrCreateStreamState(chatID)
+
+		db.EXPECT().GetLastChatMessageByRole(gomock.Any(), gomock.Any()).Return(
+			database.ChatMessage{}, sql.ErrNoRows,
+		)
+
+		server.seedAssistantCheckpoint(ctx, chatID, state, server.logger)
+
+		state.mu.Lock()
+		defer state.mu.Unlock()
+		require.Equal(t, int64(0), state.lastCommittedAssistantMessageID,
+			"a fresh chat with no prior assistant messages must leave the checkpoint at zero")
+	})
+}
+
+// TestSubscribeToStream_FiltersBufferedParts_Integration wires
+// publishToStream, advanceAssistantCheckpoint, and subscribeToStream
+// together to confirm the end-to-end contract: a subscriber with a
+// known cursor only receives parts from turns the cursor does not
+// already cover.
+func TestSubscribeToStream_FiltersBufferedParts_Integration(t *testing.T) {
+	t.Parallel()
+
+	mClock := quartz.NewMock(t)
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  mClock,
+	}
+	chatID := uuid.New()
+
+	// Start buffering, then simulate the lifecycle:
+	//   1. Stream parts of turn A (checkpoint = 0, no commit yet).
+	//   2. Commit turn A's durable message with ID 100.
+	//   3. Stream parts of turn B (checkpoint now = 100).
+	//   4. Commit turn B's durable message with ID 200.
+	//   5. Stream parts of turn C (checkpoint now = 200).
+	state := server.getOrCreateStreamState(chatID)
+	state.mu.Lock()
+	state.buffering = true
+	state.mu.Unlock()
+
+	publish := func(text string) {
+		server.publishToStream(chatID, codersdk.ChatStreamEvent{
+			Type: codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{
+				Role: codersdk.ChatMessageRoleAssistant,
+				Part: codersdk.ChatMessageText(text),
+			},
+		})
+	}
+
+	publish("A-1")
+	publish("A-2")
+	server.advanceAssistantCheckpoint(chatID, database.ChatMessage{
+		ID:   100,
+		Role: database.ChatMessageRoleAssistant,
+	})
+	publish("B-1")
+	publish("B-2")
+	server.advanceAssistantCheckpoint(chatID, database.ChatMessage{
+		ID:   200,
+		Role: database.ChatMessageRoleAssistant,
+	})
+	publish("C-1")
+
+	// Subscriber that already has turn A (cursor = 100) should
+	// receive only turn B and turn C parts.
+	snapshot, _, _, cancel := server.subscribeToStream(chatID, 100)
+	defer cancel()
+
+	texts := make([]string, 0, len(snapshot))
+	for _, ev := range snapshot {
+		texts = append(texts, partText(ev))
+	}
+	require.Equal(t, []string{"B-1", "B-2", "C-1"}, texts,
+		"subscriber past turn A must not receive turn A parts")
+
+	// Subscriber that already has both A and B (cursor = 200)
+	// should receive only turn C parts.
+	snapshot2, _, _, cancel2 := server.subscribeToStream(chatID, 200)
+	defer cancel2()
+	texts2 := make([]string, 0, len(snapshot2))
+	for _, ev := range snapshot2 {
+		texts2 = append(texts2, partText(ev))
+	}
+	require.Equal(t, []string{"C-1"}, texts2,
+		"subscriber past turn B must not receive turn A or B parts")
+
+	// Fresh subscriber (cursor = 0) receives the entire buffer.
+	snapshot3, _, _, cancel3 := server.subscribeToStream(chatID, 0)
+	defer cancel3()
+	require.Len(t, snapshot3, 5,
+		"fresh subscriber must receive every buffered part")
+
+	// Relay subscriber (sentinel) receives the entire buffer.
+	snapshot4, _, _, cancel4 := server.subscribeToStream(chatID, math.MaxInt64)
+	defer cancel4()
+	require.Len(t, snapshot4, 5,
+		"relay sentinel must receive every buffered part")
+}

coderd/x/chatd/streamcollector_internal_test.go

@@ -44,11 +44,11 @@ func TestStreamStateCollector(t *testing.T) {
 		server := &Server{}
 
 		server.chatStreams.Store(uuid.New(), &chatStreamState{
-			buffer:      make([]codersdk.ChatStreamEvent, 10),
+			buffer:      make([]bufferedStreamPart, 10),
 			subscribers: newSubscribers(t, 2),
 		})
 		server.chatStreams.Store(uuid.New(), &chatStreamState{
-			buffer:      make([]codersdk.ChatStreamEvent, 25),
+			buffer:      make([]bufferedStreamPart, 25),
 			subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{},
 		})
 		server.chatStreams.Store(uuid.New(), &chatStreamState{
@@ -74,7 +74,7 @@ func TestStreamStateCollector(t *testing.T) {
 
 		server.chatStreams.Store(uuid.New(), "garbage")
 		server.chatStreams.Store(uuid.New(), &chatStreamState{
-			buffer:      make([]codersdk.ChatStreamEvent, 5),
+			buffer:      make([]bufferedStreamPart, 5),
 			subscribers: newSubscribers(t, 1),
 		})
 
@@ -97,7 +97,7 @@ func TestStreamStateCollector(t *testing.T) {
 
 		server := &Server{}
 		state := &chatStreamState{
-			buffer:      make([]codersdk.ChatStreamEvent, 0, 100),
+			buffer:      make([]bufferedStreamPart, 0, 100),
 			subscribers: newSubscribers(t, 1),
 		}
 		server.chatStreams.Store(uuid.New(), state)
@@ -110,7 +110,7 @@ func TestStreamStateCollector(t *testing.T) {
 		wg.Go(func() {
 			for range iterations {
 				state.mu.Lock()
-				state.buffer = append(state.buffer, codersdk.ChatStreamEvent{})
+				state.buffer = append(state.buffer, bufferedStreamPart{})
 				if len(state.buffer) > 50 {
 					state.buffer = state.buffer[10:]
 				}
@@ -185,7 +185,7 @@ func TestStreamStateCollector_BufferDroppedIncrementsOnCapacity(t *testing.T) {
 	chatID := uuid.New()
 	server.chatStreams.Store(chatID, &chatStreamState{
 		buffering: true,
-		buffer:    make([]codersdk.ChatStreamEvent, maxStreamBufferSize),
+		buffer:    make([]bufferedStreamPart, maxStreamBufferSize),
 	})
 
 	partEvent := codersdk.ChatStreamEvent{

enterprise/coderd/x/chatd/chatd.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"math"
 	"net/http"
 	"net/url"
 	"strconv"
@@ -853,8 +852,9 @@ func buildRelayURL(address string, chatID uuid.UUID) (string, error) {
 	u.Path = fmt.Sprintf("/api/experimental/chats/%s/stream", chatID)
 	q := u.Query()
 	// Relays only need live message_part events, not the full
-	// history; pass after_id=MaxInt64 so the peer skips its snapshot.
-	q.Set("after_id", strconv.FormatInt(math.MaxInt64, 10))
+	// history; pass the relay sentinel so the peer skips its
+	// durable DB snapshot and delivers in-flight parts only.
+	q.Set("after_id", strconv.FormatInt(osschatd.RelaySentinelAfterID, 10))
 	u.RawQuery = q.Encode()
 	return u.String(), nil
 }