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

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

Browse Source
This commit is contained in:
Kyle Carberry
2026-05-12 00:30:38 -04:00
committed by GitHub
parent 07ff3b3f90
commit 5a5cd79c4c
4 changed files with 356 additions and 752 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/x/chatd/chatd.go
+114 -175
View File
@@ -86,9 +86,10 @@ const (
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).
// snapshot and only deliver buffered message_part events. The
// buffer itself filters committed parts out (see snapshotBufferLocked),
// so the sentinel resolves to "send me any in-progress streaming
// parts you have; I will receive durable messages through pubsub."
RelaySentinelAfterID = math.MaxInt64
// maxDurableMessageCacheSize caps the number of recent durable message
// events cached per chat for same-replica stream catch-up.
@@ -114,10 +115,15 @@ const (
// goroutines and lifecycle management.
streamDropWarnInterval = 10 * time.Second
// bufferRetainGracePeriod is how long the message_part
// buffer is kept after processing completes. This gives
// cross-replica relay subscribers time to connect and
// snapshot the buffer before it is garbage-collected.
// bufferRetainGracePeriod is how long the per-chat stream
// state is kept after processing completes. The retained
// state lets late-connecting cross-replica relay subscribers
// register against the live stream before the next worker
// run starts, preventing a race between cleanupStreamIfIdle
// and subscriber registration. The buffer itself is no
// longer useful at this point: every part has been claimed
// by its durable assistant message and is filtered out of
// the subscriber snapshot.
bufferRetainGracePeriod = 5 * time.Second
// chatStreamControlFetchTimeout bounds subscriber-owned
// control-path DB reads when the caller has no deadline.
@@ -1099,21 +1105,22 @@ 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.
// bufferedStreamPart is a buffered message_part event with its
// committed-message linkage. Parts that have not yet been claimed by
// a durable assistant message carry committedMessageID == 0 and are
// considered "in progress"; when an assistant message is published
// every still-in-progress part is claimed by that durable message
// ID, marking the part as redundant for any subscriber that will
// receive the durable message via REST or pubsub.
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
// committedMessageID is the durable assistant message ID that
// claimed this part, or 0 while the part belongs to the
// in-progress turn. snapshotBufferLocked drops parts with
// committedMessageID != 0 because the subscriber will receive
// the durable message through a different channel (REST snapshot,
// initial DB query in SubscribeAuthorized, or pubsub).
committedMessageID int64
}
type chatStreamState struct {
@@ -1132,18 +1139,16 @@ type chatStreamState struct {
// to retry.
currentRetry *codersdk.ChatStreamRetry
// bufferRetainedAt records when processing completed and
// the buffer was retained for late-connecting relay
// subscribers. Zero while buffering is active. When
// the per-chat stream state entered the post-completion
// grace window. Zero while buffering is active. When
// non-zero, cleanupStreamIfIdle skips GC until the grace
// period expires so cross-replica relays can still
// snapshot the buffer.
// period expires so cross-replica relay subscribers can
// register without racing state deletion. The buffer
// itself does not deliver content here: every part is
// claimed by a durable assistant message before
// bufferRetainedAt is set, so snapshotBufferLocked
// returns no parts during the grace window.
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
@@ -4178,8 +4183,10 @@ func (p *Server) publishToStream(chatID uuid.UUID, event codersdk.ChatStreamEven
state.buffer = state.buffer[1:]
}
state.buffer = append(state.buffer, bufferedStreamPart{
event: event,
checkpoint: state.lastCommittedAssistantMessageID,
event: event,
// committedMessageID stays 0 here: the part belongs to
// the in-progress turn until publishMessage claims it
// with the committed assistant message ID.
})
}
subscribers := make([]chan codersdk.ChatStreamEvent, 0, len(state.subscribers))
@@ -4264,50 +4271,30 @@ func (p *Server) getCachedDurableMessages(
}
// snapshotBufferLocked returns the buffered message_part events that
// the caller should receive in their initial snapshot, filtered by
// the requested cursor.
// the caller should receive in their initial snapshot.
//
// 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.
// Parts whose committedMessageID != 0 are dropped: those parts were
// claimed by a durable assistant message that the subscriber will
// receive through a different channel (REST snapshot, the initial DB
// query in SubscribeAuthorized, or pubsub catch-up). Delivering them
// here would render the same content twice on the client, once in the
// streaming UI and once as a durable message.
//
// 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 {
// Every caller receives the same view: in-progress parts are always
// delivered and committed parts are always dropped, regardless of
// cursor or relay sentinel. This keeps the buffer free of duplicate
// work for every subscriber, including cross-replica relay
// subscribers whose user-facing peers receive the durable message
// via pubsub.
//
// The caller must hold the per-chat stream state lock.
func snapshotBufferLocked(buffer []bufferedStreamPart) []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 {
if part.committedMessageID != 0 {
continue
}
snapshot = append(snapshot, part.event)
@@ -4316,25 +4303,17 @@ func snapshotBufferLocked(
}
// 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.
// stream and returns a snapshot of currently in-progress 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) (
// Parts that were claimed by a committed durable assistant message
// (committedMessageID != 0) are excluded from the snapshot. The
// subscriber will receive those durable messages through the REST
// snapshot, the initial DB query in SubscribeAuthorized, or pubsub,
// so re-delivering their constituent parts here would render the
// same content twice.
func (p *Server) subscribeToStream(chatID uuid.UUID) (
[]codersdk.ChatStreamEvent,
*codersdk.ChatStreamRetry,
<-chan codersdk.ChatStreamEvent,
@@ -4342,7 +4321,7 @@ func (p *Server) subscribeToStream(chatID uuid.UUID, afterMessageID int64) (
) {
state := p.getOrCreateStreamState(chatID)
state.mu.Lock()
snapshot := snapshotBufferLocked(state.buffer, afterMessageID, state.lastCommittedAssistantMessageID)
snapshot := snapshotBufferLocked(state.buffer)
var currentRetry *codersdk.ChatStreamRetry
if state.currentRetry != nil {
retryCopy := *state.currentRetry
@@ -4400,8 +4379,8 @@ func (p *Server) cleanupStreamIfIdle(chatID uuid.UUID, state *chatStreamState) b
return false
}
// Keep stream state alive during the grace period so
// late-connecting relay subscribers can snapshot the
// buffer after the worker finishes processing.
// late-connecting cross-replica relay subscribers can
// register against this chat before GC.
if !state.bufferRetainedAt.IsZero() &&
p.clock.Now().Before(state.bufferRetainedAt.Add(bufferRetainGracePeriod)) {
return false
@@ -4645,7 +4624,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, afterMessageID)
localSnapshot, localRetry, localParts, localCancel := p.subscribeToStream(chatID)
// Merge all event sources.
mergedCtx, mergedCancel := context.WithCancel(ctx)
@@ -5326,84 +5305,48 @@ func (p *Server) publishMessage(chatID uuid.UUID, message database.ChatMessage)
Message: &sdkMessage,
}
p.cacheDurableMessage(chatID, event)
p.advanceAssistantCheckpoint(chatID, message)
// Claim every still-in-progress buffered message_part for this
// durable assistant message BEFORE publishing it, so any new
// subscriber that races publishEvent below takes a buffer
// snapshot in which the parts for this turn are already
// suppressed. Existing subscribers already received the
// constituent parts on the live channel; the frontend
// dedupes those against the durable message via
// clearStreamState in the same batch.
p.claimCommittedParts(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.
// claimCommittedParts walks the chat's buffered message_part events
// and assigns every in-progress part (committedMessageID == 0) to
// the supplied assistant message ID. Subsequent subscriber snapshots
// drop those parts so a reconnecting client does not re-render the
// content of an assistant turn that has already been delivered as a
// durable message via REST or pubsub.
//
// 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) {
// only assistant-role messages claim parts.
func (p *Server) claimCommittedParts(chatID uuid.UUID, message database.ChatMessage) {
if message.Role != database.ChatMessageRoleAssistant {
return
}
state := p.getOrCreateStreamState(chatID)
val, ok := p.chatStreams.Load(chatID)
if !ok {
return
}
state, ok := val.(*chatStreamState)
if !ok {
return
}
state.mu.Lock()
defer state.mu.Unlock()
if message.ID > state.lastCommittedAssistantMessageID {
state.lastCommittedAssistantMessageID = message.ID
for i := range state.buffer {
if state.buffer[i].committedMessageID == 0 {
state.buffer[i].committedMessageID = message.ID
}
}
}
@@ -5857,19 +5800,7 @@ 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
@@ -5877,11 +5808,18 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) {
streamState.currentRetry = nil
streamState.resetDropCounters()
streamState.buffering = false
// Retain the buffer for a grace period so
// cross-replica relay subscribers can still snapshot
// it after processing completes. The buffer is
// Retain the per-chat stream state for a grace period
// so cross-replica relay subscribers can register
// against this chat after processing completes,
// without racing cleanupStreamIfIdle. The buffer is
// cleared when the next processChat starts or when
// cleanupStreamIfIdle runs after the grace period.
// cleanupStreamIfIdle runs after the grace period; on
// the normal-completion path every part has been
// claimed by its durable assistant message, so the
// snapshot is empty. On error or panic exit some parts
// may still be in-progress; those are likewise
// discarded when the buffer is cleared, and the
// frontend recovers via the next REST snapshot.
streamState.bufferRetainedAt = p.clock.Now()
streamState.mu.Unlock()
}()
@@ -7264,9 +7202,10 @@ func (p *Server) runChat(
}
}
// Do NOT clear the stream buffer here. Cross-replica
// relay subscribers may still need to snapshot buffered
// message_parts after processing completes. The buffer
// Do NOT clear the stream buffer here. The per-chat
// stream state must remain alive for the post-completion
// grace window so cross-replica relay subscribers can
// register without racing cleanupStreamIfIdle. The buffer
// is bounded by maxStreamBufferSize and is cleared when
// the next processChat starts or when the stream state
// is garbage-collected after the retention grace period.
coderd/x/chatd/chatd_internal_test.go
+165 -317
View File
@@ -4,7 +4,6 @@ import (
"context"
"database/sql"
"encoding/json"
"math"
"sync"
"testing"
"time"
@@ -3503,9 +3502,6 @@ 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{})
@@ -3758,7 +3754,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, 0)
snapshot, currentRetry, events, cancelSub := server.subscribeToStream(chatID)
require.Len(t, snapshot, 1)
require.Nil(t, currentRetry)
require.NotNil(t, events)
@@ -5280,9 +5276,6 @@ 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.
@@ -5367,11 +5360,12 @@ func TestAutoPromote_InsertFailureSkipsStatusUpdate(t *testing.T) {
}
}
// makeBufferedPart is a small constructor for buffered message_part
// makeInProgressPart 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 {
// builds an in-progress part (committedMessageID == 0) with a
// recognizable text body so failing assertions can identify which
// part survived the filter.
func makeInProgressPart(text string) bufferedStreamPart {
return bufferedStreamPart{
event: codersdk.ChatStreamEvent{
Type: codersdk.ChatStreamEventTypeMessagePart,
@@ -5380,10 +5374,17 @@ func makeBufferedPart(checkpoint int64, text string) bufferedStreamPart {
Part: codersdk.ChatMessageText(text),
},
},
checkpoint: checkpoint,
}
}
// makeCommittedPart builds a part already claimed by the given
// durable assistant message ID.
func makeCommittedPart(committedID int64, text string) bufferedStreamPart {
p := makeInProgressPart(text)
p.committedMessageID = committedID
return p
}
func partText(event codersdk.ChatStreamEvent) string {
if event.MessagePart == nil {
return ""
@@ -5391,126 +5392,46 @@ func partText(event codersdk.ChatStreamEvent) string {
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) {
// TestSnapshotBufferLocked_DropsCommittedParts asserts the core
// dedup contract: parts that were claimed by a durable assistant
// message (committedMessageID != 0) are dropped from the snapshot
// because the subscriber will receive that durable message through
// the REST snapshot, the initial DB query, or pubsub.
func TestSnapshotBufferLocked_DropsCommittedParts(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"),
makeCommittedPart(100, "turnA-1"),
makeCommittedPart(100, "turnA-2"),
makeCommittedPart(200, "turnB-1"),
makeInProgressPart("in-progress-1"),
makeInProgressPart("in-progress-2"),
}
// 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)
snapshot := snapshotBufferLocked(buffer)
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]))
"only in-progress (committedMessageID == 0) parts should be kept")
require.Equal(t, "in-progress-1", partText(snapshot[0]))
require.Equal(t, "in-progress-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) {
// TestSnapshotBufferLocked_AllInProgressReturnsAll covers the
// fresh-load convention: when no assistant message has committed
// yet, every buffered part is in-progress and must be delivered.
func TestSnapshotBufferLocked_AllInProgressReturnsAll(t *testing.T) {
t.Parallel()
buffer := []bufferedStreamPart{
makeBufferedPart(0, "a"),
makeBufferedPart(0, "b"),
makeBufferedPart(0, "c"),
makeInProgressPart("a"),
makeInProgressPart("b"),
makeInProgressPart("c"),
}
snapshot := snapshotBufferLocked(buffer, 999, 0)
snapshot := snapshotBufferLocked(buffer)
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")
"all in-progress parts must be delivered to the subscriber")
require.Equal(t, "a", partText(snapshot[0]))
require.Equal(t, "b", partText(snapshot[1]))
require.Equal(t, "c", partText(snapshot[2]))
@@ -5522,16 +5443,33 @@ func TestSnapshotBufferLocked_RelaySentinelReturnsAll(t *testing.T) {
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))
require.Nil(t, snapshotBufferLocked(nil))
require.Nil(t, snapshotBufferLocked([]bufferedStreamPart{}))
}
// 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) {
// TestSnapshotBufferLocked_AllCommittedReturnsEmpty covers the
// natural resting point after an assistant turn commits and before
// the next turn starts streaming: every buffered part has been
// claimed and must be filtered out. The snapshot must be empty so
// reconnecting subscribers do not re-render content that is already
// available as a durable message.
func TestSnapshotBufferLocked_AllCommittedReturnsEmpty(t *testing.T) {
t.Parallel()
buffer := []bufferedStreamPart{
makeCommittedPart(100, "a"),
makeCommittedPart(100, "b"),
makeCommittedPart(200, "c"),
}
require.Empty(t, snapshotBufferLocked(buffer))
}
// TestPublishToStream_AppendsAsInProgress verifies that parts
// buffered while the chat is streaming are tagged as in-progress
// (committedMessageID == 0) until publishMessage claims them via a
// committed assistant message.
func TestPublishToStream_AppendsAsInProgress(t *testing.T) {
t.Parallel()
mClock := quartz.NewMock(t)
@@ -5542,9 +5480,8 @@ func TestPublishToStream_TagsPartsWithCurrentCheckpoint(t *testing.T) {
chatID := uuid.New()
state := &chatStreamState{
buffering: true,
subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{},
lastCommittedAssistantMessageID: 100,
buffering: true,
subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{},
}
server.chatStreams.Store(chatID, state)
@@ -5559,206 +5496,138 @@ func TestPublishToStream_TagsPartsWithCurrentCheckpoint(t *testing.T) {
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, int64(0), state.buffer[0].committedMessageID,
"newly buffered parts must be in-progress until publishMessage claims them")
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) {
// TestClaimCommittedParts covers the per-role behavior of
// claimCommittedParts:
// - assistant messages claim every in-progress part with the
// committed message ID.
// - tool / user messages do not claim parts.
// - parts already claimed by an earlier assistant message are not
// re-claimed.
// - a chat with no live state is a no-op (does not panic).
func TestClaimCommittedParts(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.Run("AssistantClaimsAllInProgressParts", 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.buffer = []bufferedStreamPart{
makeCommittedPart(100, "old-1"),
makeInProgressPart("new-1"),
makeInProgressPart("new-2"),
}
state.mu.Unlock()
db.EXPECT().GetLastChatMessageByRole(gomock.Any(), database.GetLastChatMessageByRoleParams{
ChatID: chatID,
Role: database.ChatMessageRoleAssistant,
}).Return(database.ChatMessage{
ID: 500,
server.claimCommittedParts(chatID, database.ChatMessage{
ID: 200,
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")
require.Equal(t, int64(100), state.buffer[0].committedMessageID,
"already-claimed parts must keep their original message ID")
require.Equal(t, int64(200), state.buffer[1].committedMessageID,
"in-progress parts must be claimed by the new message ID")
require.Equal(t, int64(200), state.buffer[2].committedMessageID,
"in-progress parts must be claimed by the new message ID")
})
t.Run("DoesNotMoveCheckpointBackwards", func(t *testing.T) {
t.Run("ToolMessageIsNoOp", 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),
state.buffer = []bufferedStreamPart{
makeInProgressPart("a"),
makeInProgressPart("b"),
}
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)
server.claimCommittedParts(chatID, database.ChatMessage{
ID: 300,
Role: database.ChatMessageRoleTool,
})
state.mu.Lock()
defer state.mu.Unlock()
require.Equal(t, int64(42), state.lastCommittedAssistantMessageID,
"a non-ErrNoRows DB error must not change the checkpoint")
require.Equal(t, int64(0), state.buffer[0].committedMessageID,
"tool messages must not claim buffered parts")
require.Equal(t, int64(0), state.buffer[1].committedMessageID,
"tool messages must not claim buffered parts")
})
t.Run("NoRowsLeavesCheckpointAtZero", func(t *testing.T) {
t.Run("UserMessageIsNoOp", 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.buffer = []bufferedStreamPart{
makeInProgressPart("a"),
}
state.mu.Unlock()
db.EXPECT().GetLastChatMessageByRole(gomock.Any(), gomock.Any()).Return(
database.ChatMessage{}, sql.ErrNoRows,
)
server.seedAssistantCheckpoint(ctx, chatID, state, server.logger)
server.claimCommittedParts(chatID, database.ChatMessage{
ID: 400,
Role: database.ChatMessageRoleUser,
})
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")
require.Equal(t, int64(0), state.buffer[0].committedMessageID,
"user messages must not claim buffered parts")
})
t.Run("NoLiveStateIsNoOp", func(t *testing.T) {
t.Parallel()
server := &Server{
logger: slogtest.Make(t, nil),
clock: quartz.NewMock(t),
}
chatID := uuid.New()
// No state stored: claimCommittedParts must not panic and
// must not allocate a new state for an unknown chat.
require.NotPanics(t, func() {
server.claimCommittedParts(chatID, database.ChatMessage{
ID: 500,
Role: database.ChatMessageRoleAssistant,
})
})
_, ok := server.chatStreams.Load(chatID)
require.False(t, ok,
"claimCommittedParts must not create stream state for a chat that has none")
})
}
// 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.
// publishToStream, claimCommittedParts (via publishMessage), and
// subscribeToStream together to confirm the end-to-end contract: a
// reconnecting subscriber only receives parts that belong to the
// current in-progress turn, not parts that were already committed
// to durable assistant messages.
func TestSubscribeToStream_FiltersBufferedParts_Integration(t *testing.T) {
t.Parallel()
@@ -5769,18 +5638,18 @@ func TestSubscribeToStream_FiltersBufferedParts_Integration(t *testing.T) {
}
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).
// Simulate the lifecycle:
// 1. Stream parts of turn A (still in-progress, no commit yet).
// 2. Commit turn A; its parts are claimed by message 100.
// 3. Stream parts of turn B (in-progress).
// 4. Commit turn B; its parts are claimed by message 200.
// 5. Stream parts of turn C (in-progress, never committed).
state := server.getOrCreateStreamState(chatID)
state.mu.Lock()
state.buffering = true
state.mu.Unlock()
publish := func(text string) {
publishPart := func(text string) {
server.publishToStream(chatID, codersdk.ChatStreamEvent{
Type: codersdk.ChatStreamEventTypeMessagePart,
MessagePart: &codersdk.ChatStreamMessagePart{
@@ -5790,52 +5659,31 @@ func TestSubscribeToStream_FiltersBufferedParts_Integration(t *testing.T) {
})
}
publish("A-1")
publish("A-2")
server.advanceAssistantCheckpoint(chatID, database.ChatMessage{
publishPart("A-1")
publishPart("A-2")
server.claimCommittedParts(chatID, database.ChatMessage{
ID: 100,
Role: database.ChatMessageRoleAssistant,
})
publish("B-1")
publish("B-2")
server.advanceAssistantCheckpoint(chatID, database.ChatMessage{
publishPart("B-1")
publishPart("B-2")
server.claimCommittedParts(chatID, database.ChatMessage{
ID: 200,
Role: database.ChatMessageRoleAssistant,
})
publish("C-1")
publishPart("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)
// Reconnecting subscriber: only the currently in-progress turn
// (turn C) survives the filter, no matter what cursor the
// client passes through SubscribeAuthorized (the filter no
// longer depends on the cursor).
snapshot, _, _, cancel := server.subscribeToStream(chatID)
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")
require.Equal(t, []string{"C-1"}, texts,
"only in-progress (un-claimed) buffered parts must survive the filter")
}
coderd/x/chatd/chatd_test.go
+52 -18
View File
@@ -9560,6 +9560,48 @@ func TestAdvisorHappyPath_RootChat(t *testing.T) {
})
require.NoError(t, err)
// Subscribe before the worker commits any durable messages so we
// observe the advisor tool-result deltas live. Buffered parts are
// claimed by their committed durable message ID at publishMessage
// time and dropped from snapshots of late-connecting subscribers, so
// a post-completion Subscribe() would no longer see streaming
// deltas. Collecting events from the live channel covers the
// streaming UX contract this test exists to verify.
_, liveEvents, cancelLive, ok := server.Subscribe(ctx, chat.ID, nil, 0)
require.True(t, ok)
var (
livePartsMu sync.Mutex
liveAdvisorDeltas []string
liveCollectorDone = make(chan struct{})
)
go func() {
defer close(liveCollectorDone)
for {
select {
case <-ctx.Done():
return
case event, eventsOK := <-liveEvents:
if !eventsOK {
return
}
if event.Type != codersdk.ChatStreamEventTypeMessagePart ||
event.MessagePart == nil {
continue
}
part := event.MessagePart.Part
if event.MessagePart.Role != codersdk.ChatMessageRoleTool ||
part.Type != codersdk.ChatMessagePartTypeToolResult ||
part.ToolName != chatadvisor.ToolName ||
part.ResultDelta == "" {
continue
}
livePartsMu.Lock()
liveAdvisorDeltas = append(liveAdvisorDeltas, part.ResultDelta)
livePartsMu.Unlock()
}
}
}()
require.Eventually(t, func() bool {
got, getErr := db.GetChatByID(ctx, chat.ID)
if getErr != nil {
@@ -9614,24 +9656,16 @@ func TestAdvisorHappyPath_RootChat(t *testing.T) {
require.True(t, parentSawAdvisorResult,
"parent must see the advisor reply in its continuation call")
snapshot, _, cancelStream, ok := server.Subscribe(ctx, chat.ID, nil, 0)
require.True(t, ok)
cancelStream()
var streamedAdvisorDeltas []string
for _, event := range snapshot {
if event.Type != codersdk.ChatStreamEventTypeMessagePart || event.MessagePart == nil {
continue
}
part := event.MessagePart.Part
if event.MessagePart.Role == codersdk.ChatMessageRoleTool &&
part.Type == codersdk.ChatMessagePartTypeToolResult &&
part.ToolName == chatadvisor.ToolName &&
part.ResultDelta != "" {
streamedAdvisorDeltas = append(streamedAdvisorDeltas, part.ResultDelta)
}
}
require.Equal(t, advisorDeltas, streamedAdvisorDeltas,
// Stop the live collector and assert it captured the streaming
// advisor deltas during processing. Late subscribers no longer
// see committed parts because publishMessage claims them out of
// new snapshots, so the assertion must use the live collector.
cancelLive()
<-liveCollectorDone
livePartsMu.Lock()
collectedAdvisorDeltas := append([]string(nil), liveAdvisorDeltas...)
livePartsMu.Unlock()
require.Equal(t, advisorDeltas, collectedAdvisorDeltas,
"advisor nested text deltas must stream into the parent tool card")
persisted, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{