From 0ed57ee3436492752db8ff51b924f78b28ba3de7 Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Mon, 11 May 2026 17:27:03 -0400 Subject: [PATCH] fix(coderd/x/chatd): checkpoint buffered message_parts to avoid stale replay (#25145) --- coderd/x/chatd/chatd.go | 203 ++++++- coderd/x/chatd/chatd_internal_test.go | 564 ++++++++++++++++-- .../x/chatd/streamcollector_internal_test.go | 12 +- enterprise/coderd/x/chatd/chatd.go | 6 +- 4 files changed, 736 insertions(+), 49 deletions(-) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 5c15e33ab1..5429e47ba3 100644 --- a/coderd/x/chatd/chatd.go +++ b/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 diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 0e2d9c8242..f4a804c475 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/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") +} diff --git a/coderd/x/chatd/streamcollector_internal_test.go b/coderd/x/chatd/streamcollector_internal_test.go index 417aab1f56..81dae5f133 100644 --- a/coderd/x/chatd/streamcollector_internal_test.go +++ b/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{ diff --git a/enterprise/coderd/x/chatd/chatd.go b/enterprise/coderd/x/chatd/chatd.go index e407e0e23d..8301e1d191 100644 --- a/enterprise/coderd/x/chatd/chatd.go +++ b/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 }