From 3a153ebb1555033da999d1da71cb1eba24dbb7d6 Mon Sep 17 00:00:00 2001 From: Ethan <39577870+ethanndickson@users.noreply.github.com> Date: Mon, 4 May 2026 11:48:39 +1000 Subject: [PATCH] fix(coderd/x/chatd): replay retry phase on subscribe (#24569) Retry events were previously fire-and-forget, so subscribers that connected after a retry started only saw durable history plus `status=running` and could not tell the stream was backing off. Keep the current retry phase in `chatStreamState`, capture it atomically with subscriber registration, replay it in the initial snapshot for same-chat late joiners, and clear it when streaming resumes or ends so reconnects get consistent retry state without duplicate delivery at the subscription boundary. Relates to CODAGT-139 --- coderd/x/chatd/chatd.go | 81 +++++++- coderd/x/chatd/chatd_internal_test.go | 261 ++++++++++++++++++++++++-- 2 files changed, 324 insertions(+), 18 deletions(-) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 7aa4426417..d9183bb4cb 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -975,6 +975,10 @@ type chatStreamState struct { bufferLastWarnAt time.Time subscriberDropCount int64 subscriberLastWarnAt time.Time + // currentRetry records the current retry phase for late-joining + // same-replica subscribers. Nil when the stream is not waiting + // 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 @@ -4099,9 +4103,41 @@ func (p *Server) processOnce(ctx context.Context) { p.inflightMu.Unlock() } +func shouldClearRetryPhaseForStatus(status codersdk.ChatStatus) bool { + switch status { + case codersdk.ChatStatusWaiting, + codersdk.ChatStatusPending, + codersdk.ChatStatusPaused, + codersdk.ChatStatusCompleted, + codersdk.ChatStatusError, + codersdk.ChatStatusRequiresAction: + return true + default: + return false + } +} + func (p *Server) publishToStream(chatID uuid.UUID, event codersdk.ChatStreamEvent) { state := p.getOrCreateStreamState(chatID) state.mu.Lock() + switch event.Type { + case codersdk.ChatStreamEventTypeRetry: + if event.Retry != nil { + retryCopy := *event.Retry + state.currentRetry = &retryCopy + } + case codersdk.ChatStreamEventTypeMessagePart: + // Any streamed part means the provider is making forward + // progress again, so the stream has left the retry backoff + // window regardless of role. + state.currentRetry = nil + case codersdk.ChatStreamEventTypeError: + state.currentRetry = nil + case codersdk.ChatStreamEventTypeStatus: + if event.Status != nil && shouldClearRetryPhaseForStatus(event.Status.Status) { + state.currentRetry = nil + } + } if event.Type == codersdk.ChatStreamEventTypeMessagePart { if !state.buffering { p.cleanupStreamIfIdle(chatID, state) @@ -4212,12 +4248,18 @@ func (p *Server) getCachedDurableMessages( func (p *Server) subscribeToStream(chatID uuid.UUID) ( []codersdk.ChatStreamEvent, + *codersdk.ChatStreamRetry, <-chan codersdk.ChatStreamEvent, func(), ) { state := p.getOrCreateStreamState(chatID) state.mu.Lock() snapshot := append([]codersdk.ChatStreamEvent(nil), state.buffer...) + var currentRetry *codersdk.ChatStreamRetry + if state.currentRetry != nil { + retryCopy := *state.currentRetry + currentRetry = &retryCopy + } id := uuid.New() ch := make(chan codersdk.ChatStreamEvent, 128) state.subscribers[id] = ch @@ -4235,7 +4277,7 @@ func (p *Server) subscribeToStream(chatID uuid.UUID) ( state.mu.Unlock() } - return snapshot, ch, cancel + return snapshot, currentRetry, ch, cancel } // getOrCreateStreamState returns the per-chat stream state, @@ -4456,8 +4498,10 @@ func (p *Server) Subscribe( } // Subscribe to the local stream for message_parts and same-replica - // persisted messages. - localSnapshot, localParts, localCancel := p.subscribeToStream(chatID) + // 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) // Merge all event sources. mergedCtx, mergedCancel := context.WithCancel(ctx) @@ -4521,13 +4565,24 @@ func (p *Server) Subscribe( // is already active so no notifications can be lost during this // window. initialSnapshot := make([]codersdk.ChatStreamEvent, 0) - // Add local message_parts to snapshot + // Add local same-replica message_parts to the snapshot. Retry comes + // from state.currentRetry, not the event buffer, so late joiners see + // only the latest phase rather than a stale buffered retry event. for _, event := range localSnapshot { if event.Type == codersdk.ChatStreamEventTypeMessagePart { initialSnapshot = append(initialSnapshot, event) } } + var retryEvent *codersdk.ChatStreamEvent + if localRetry != nil { + retryEvent = &codersdk.ChatStreamEvent{ + Type: codersdk.ChatStreamEventTypeRetry, + ChatID: chatID, + Retry: localRetry, + } + } + // Load initial messages from DB. When afterMessageID > 0 the // caller already has messages up to that ID (e.g. from the REST // endpoint), so we only fetch newer ones to avoid sending @@ -4602,9 +4657,18 @@ func (p *Server) Subscribe( Status: codersdk.ChatStatus(chat.Status), }, } - // Prepend so the frontend sees the status before any - // message_part events. - initialSnapshot = append([]codersdk.ChatStreamEvent{statusEvent}, initialSnapshot...) + // Prepend so the frontend sees the current stream phases + // before any message_part events. + prefix := []codersdk.ChatStreamEvent{statusEvent} + if retryEvent != nil { + prefix = append(prefix, *retryEvent) + retryEvent = nil + } + initialSnapshot = append(prefix, initialSnapshot...) + } + + if retryEvent != nil { + initialSnapshot = append(initialSnapshot, *retryEvent) } // Track the highest durable message ID delivered to this subscriber, @@ -5476,6 +5540,9 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { streamState.mu.Unlock() defer func() { streamState.mu.Lock() + // Fallback cleanup for exit paths that return before a + // terminal stream event is published. + streamState.currentRetry = nil streamState.resetDropCounters() streamState.buffering = false // Retain the buffer for a grace period so diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 3a9d84b6b0..e76c5ae24f 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -2010,16 +2010,7 @@ func TestSubscribeDeliversRetryEventViaPubsubOnce(t *testing.T) { require.True(t, ok) defer cancel() - retryingAt := time.Unix(1_700_000_000, 0).UTC() - expected := &codersdk.ChatStreamRetry{ - Attempt: 1, - DelayMs: (1500 * time.Millisecond).Milliseconds(), - Error: "OpenAI is rate limiting requests (HTTP 429).", - Kind: chaterror.KindRateLimit, - Provider: "openai", - StatusCode: 429, - RetryingAt: retryingAt, - } + expected := newTestRetryPayload() server.publishRetry(chatID, expected) @@ -2028,6 +2019,190 @@ func TestSubscribeDeliversRetryEventViaPubsubOnce(t *testing.T) { requireNoStreamEvent(t, events, 200*time.Millisecond) } +func TestSubscribeReplaysCurrentRetryPhaseInSnapshot(t *testing.T) { + t.Parallel() + + ctx, cancelCtx := context.WithCancel(context.Background()) + defer cancelCtx() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + chatID := uuid.New() + chat := database.Chat{ID: chatID, Status: database.ChatStatusRunning} + + gomock.InOrder( + db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ + ChatID: chatID, + AfterID: 0, + }).Return(nil, nil), + db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), + db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), + ) + + server := newBufferedSubscribeTestServer(t, db, chatID) + + expected := newTestRetryPayload() + server.publishRetry(chatID, expected) + + snapshot, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) + require.True(t, ok) + defer cancel() + + require.Len(t, snapshot, 2) + require.Equal(t, codersdk.ChatStreamEventTypeStatus, snapshot[0].Type) + require.Equal(t, codersdk.ChatStreamEventTypeRetry, snapshot[1].Type) + event := requireSnapshotRetryEvent(t, snapshot) + require.Equal(t, expected, event.Retry) + requireNoStreamEvent(t, events, 200*time.Millisecond) +} + +func TestSubscribeCapturesRetryPhaseAtSubscriptionBoundary(t *testing.T) { + t.Parallel() + + ctx, cancelCtx := context.WithCancel(context.Background()) + defer cancelCtx() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + chatID := uuid.New() + chat := database.Chat{ID: chatID, Status: database.ChatStatusRunning} + expected := newTestRetryPayload() + + server := newSubscribeTestServer(t, db) + + gomock.InOrder( + db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ + ChatID: chatID, + AfterID: 0, + }).DoAndReturn(func(context.Context, database.GetChatMessagesByChatIDParams) ([]database.ChatMessage, error) { + server.publishRetry(chatID, expected) + return nil, nil + }), + db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), + db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), + ) + + snapshot, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) + require.True(t, ok) + defer cancel() + + requireNoSnapshotRetryEvent(t, snapshot) + event := requireStreamRetryEvent(t, events) + require.Equal(t, expected, event.Retry) + requireNoStreamEvent(t, events, 200*time.Millisecond) +} + +func TestSubscribeDoesNotReplayRetryAfterStreamResumes(t *testing.T) { + t.Parallel() + + ctx, cancelCtx := context.WithCancel(context.Background()) + defer cancelCtx() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + chatID := uuid.New() + chat := database.Chat{ID: chatID, Status: database.ChatStatusRunning} + + gomock.InOrder( + db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ + ChatID: chatID, + AfterID: 0, + }).Return(nil, nil), + db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), + db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), + ) + + server := newBufferedSubscribeTestServer(t, db, chatID) + + server.publishRetry(chatID, newTestRetryPayload()) + server.publishMessagePart(chatID, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("retry recovered")) + + snapshot, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) + require.True(t, ok) + defer cancel() + + requireNoSnapshotRetryEvent(t, snapshot) + requireSnapshotMessagePartEvent(t, snapshot) + requireNoStreamEvent(t, events, 200*time.Millisecond) +} + +func TestSubscribeDoesNotReplayRetryAfterTerminalError(t *testing.T) { + t.Parallel() + + ctx, cancelCtx := context.WithCancel(context.Background()) + defer cancelCtx() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + chatID := uuid.New() + chat := database.Chat{ID: chatID, Status: database.ChatStatusRunning} + + gomock.InOrder( + db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ + ChatID: chatID, + AfterID: 0, + }).Return(nil, nil), + db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), + db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), + ) + + server := newBufferedSubscribeTestServer(t, db, chatID) + + server.publishRetry(chatID, newTestRetryPayload()) + server.publishError(chatID, chaterror.ClassifiedError{ + Message: "OpenAI is rate limiting requests (HTTP 429).", + Kind: chaterror.KindRateLimit, + Provider: "openai", + Retryable: true, + StatusCode: 429, + }) + + snapshot, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) + require.True(t, ok) + defer cancel() + + requireNoSnapshotRetryEvent(t, snapshot) + requireNoStreamEvent(t, events, 200*time.Millisecond) +} + +func TestSubscribeDoesNotReplayRetryAfterTerminalStatus(t *testing.T) { + t.Parallel() + + ctx, cancelCtx := context.WithCancel(context.Background()) + defer cancelCtx() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + chatID := uuid.New() + chat := database.Chat{ID: chatID, Status: database.ChatStatusCompleted} + + gomock.InOrder( + db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ + ChatID: chatID, + AfterID: 0, + }).Return(nil, nil), + db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), + db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), + ) + + server := newBufferedSubscribeTestServer(t, db, chatID) + + server.publishRetry(chatID, newTestRetryPayload()) + server.publishStatus(chatID, database.ChatStatusCompleted, uuid.NullUUID{}) + + snapshot, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) + require.True(t, ok) + defer cancel() + + requireNoSnapshotRetryEvent(t, snapshot) + requireNoStreamEvent(t, events, 200*time.Millisecond) +} + func TestSubscribePrefersStructuredErrorPayloadViaPubsub(t *testing.T) { t.Parallel() @@ -2103,6 +2278,18 @@ func TestSubscribeFallsBackToLegacyErrorStringViaPubsub(t *testing.T) { requireNoStreamEvent(t, events, 200*time.Millisecond) } +func newTestRetryPayload() *codersdk.ChatStreamRetry { + return &codersdk.ChatStreamRetry{ + Attempt: 1, + DelayMs: (1500 * time.Millisecond).Milliseconds(), + Error: "OpenAI is rate limiting requests (HTTP 429).", + Kind: chaterror.KindRateLimit, + Provider: "openai", + StatusCode: 429, + RetryingAt: time.Unix(1_700_000_000, 0).UTC(), + } +} + func newSubscribeTestServer(t *testing.T, db database.Store) *Server { t.Helper() @@ -2113,6 +2300,17 @@ func newSubscribeTestServer(t *testing.T, db database.Store) *Server { } } +func newBufferedSubscribeTestServer(t *testing.T, db database.Store, chatID uuid.UUID) *Server { + t.Helper() + + server := newSubscribeTestServer(t, db) + state := server.getOrCreateStreamState(chatID) + state.mu.Lock() + state.buffering = true + state.mu.Unlock() + return server +} + func requireStreamMessageEvent(t *testing.T, events <-chan codersdk.ChatStreamEvent) codersdk.ChatStreamEvent { t.Helper() @@ -2143,6 +2341,44 @@ func requireStreamRetryEvent(t *testing.T, events <-chan codersdk.ChatStreamEven } } +func requireSnapshotRetryEvent(t *testing.T, snapshot []codersdk.ChatStreamEvent) codersdk.ChatStreamEvent { + t.Helper() + + var retryEvents []codersdk.ChatStreamEvent + for _, event := range snapshot { + if event.Type == codersdk.ChatStreamEventTypeRetry { + retryEvents = append(retryEvents, event) + } + } + + require.Len(t, retryEvents, 1, "expected exactly one retry event in snapshot") + require.NotNil(t, retryEvents[0].Retry) + return retryEvents[0] +} + +func requireNoSnapshotRetryEvent(t *testing.T, snapshot []codersdk.ChatStreamEvent) { + t.Helper() + + for _, event := range snapshot { + require.NotEqual(t, codersdk.ChatStreamEventTypeRetry, event.Type, + "unexpected retry event in snapshot: %+v", event) + } +} + +func requireSnapshotMessagePartEvent(t *testing.T, snapshot []codersdk.ChatStreamEvent) codersdk.ChatStreamEvent { + t.Helper() + + for _, event := range snapshot { + if event.Type == codersdk.ChatStreamEventTypeMessagePart { + require.NotNil(t, event.MessagePart) + return event + } + } + + t.Fatal("expected message_part event in snapshot") + return codersdk.ChatStreamEvent{} +} + func requireStreamErrorEvent(t *testing.T, events <-chan codersdk.ChatStreamEvent) codersdk.ChatStreamEvent { t.Helper() @@ -3822,7 +4058,10 @@ func TestSubscribeCancelDuringGrace_ReapedBySweep(t *testing.T) { // Real subscribeToStream cancel path: the WS subscriber detach // that leaks in prod. - _, _, cancelSub := server.subscribeToStream(chatID) + snapshot, currentRetry, events, cancelSub := server.subscribeToStream(chatID) + require.Len(t, snapshot, 1) + require.Nil(t, currentRetry) + require.NotNil(t, events) mClock.Advance(bufferRetainGracePeriod / 2) cancelSub()