diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index 39e72b0687..c8ed16c740 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -175,8 +175,61 @@ func (api *API) watchChats(rw http.ResponseWriter, r *http.Request) { apiKey := httpmw.APIKey(r) logger := api.Logger.Named("chat_watcher") + // Subscribe before accepting the websocket so the subscription + // is active when the client's Dial returns. + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + var ( + encoder *json.Encoder + encoderReady = make(chan struct{}) + // Capture before WebsocketNetConn reassigns ctx (data race). + ctxDone = ctx.Done() + ) + + cancelSubscribe, err := api.Pubsub.SubscribeWithErr(pubsub.ChatWatchEventChannel(apiKey.UserID), + pubsub.HandleChatWatchEvent( + func(cbCtx context.Context, payload codersdk.ChatWatchEvent, err error) { + if err != nil { + logger.Error(cbCtx, "chat watch event subscription error", slog.Error(err)) + return + } + select { + case <-encoderReady: + case <-ctxDone: + return + case <-cbCtx.Done(): + return + } + + // encoderReady may close with encoder still nil on error paths. + if encoder == nil { + return + } + // The encoder is only written from the pubsub delivery + // goroutine, which processes messages serially. Do not + // add a second write path without synchronization. + if err := encoder.Encode(payload); err != nil { + logger.Debug(cbCtx, "failed to send chat watch event", slog.Error(err)) + cancel() + return + } + }, + )) + if err != nil { + close(encoderReady) + logger.Error(ctx, "failed to subscribe to chat watch events", slog.Error(err)) + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to subscribe to chat events.", + Detail: err.Error(), + }) + return + } + defer cancelSubscribe() + conn, err := websocket.Accept(rw, r, nil) if err != nil { + close(encoderReady) httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Failed to open chat watch stream.", Detail: err.Error(), @@ -184,9 +237,6 @@ func (api *API) watchChats(rw http.ResponseWriter, r *http.Request) { return } - ctx, cancel := context.WithCancel(ctx) - defer cancel() - _ = conn.CloseRead(context.Background()) ctx, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageText) @@ -194,31 +244,8 @@ func (api *API) watchChats(rw http.ResponseWriter, r *http.Request) { ctx = api.wsWatcher.Watch(ctx, logger, conn) - // The encoder is only written from the SubscribeWithErr callback, - // which delivers serially per subscription. Do not add a second - // write path without introducing synchronization. - encoder := json.NewEncoder(wsNetConn) - - cancelSubscribe, err := api.Pubsub.SubscribeWithErr(pubsub.ChatWatchEventChannel(apiKey.UserID), - pubsub.HandleChatWatchEvent( - func(ctx context.Context, payload codersdk.ChatWatchEvent, err error) { - if err != nil { - logger.Error(ctx, "chat watch event subscription error", slog.Error(err)) - return - } - if err := encoder.Encode(payload); err != nil { - logger.Debug(ctx, "failed to send chat watch event", slog.Error(err)) - cancel() - return - } - }, - )) - if err != nil { - logger.Error(ctx, "failed to subscribe to chat watch events", slog.Error(err)) - _ = conn.Close(websocket.StatusInternalError, "Failed to subscribe to chat events.") - return - } - defer cancelSubscribe() + encoder = json.NewEncoder(wsNetConn) + close(encoderReady) <-ctx.Done() } diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index fdbc1160d8..a676baa7f5 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -1804,13 +1804,6 @@ func TestWatchChats(t *testing.T) { t.Run("CreatedEventIncludesAllChatFields", func(t *testing.T) { t.Parallel() - // This test verifies that the pubsub "created" event - // carries a fully-populated codersdk.Chat. Exhaustive - // field-level coverage of the converter is handled by - // TestChat_AllFieldsPopulated (db2sdk) and - // TestChat_JSONRoundTrip (codersdk). This integration - // test only checks that key fields survive the full - // API → pubsub → websocket pipeline. ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) firstUser := coderdtest.CreateFirstUser(t, client.Client) @@ -1929,31 +1922,11 @@ func TestWatchChats(t *testing.T) { payload, err := json.Marshal(event) require.NoError(t, err) - // Publish the event in a goroutine that keeps retrying. - // When the WebSocket Dial returns, the server has completed - // the HTTP upgrade but may not have called SubscribeWithErr - // yet. If we publish only once, the message can arrive - // before the subscription is active and be silently dropped, - // causing the read loop to block until the context deadline. - // Re-publishing on a short ticker guarantees that at least - // one publish lands after the subscription is ready. - publishDone := make(chan struct{}) - go func() { - ticker := time.NewTicker(testutil.IntervalFast) - defer ticker.Stop() - for { - // Publish immediately on the first iteration, - // then again on each tick. - _ = api.Pubsub.Publish(coderdpubsub.ChatWatchEventChannel(user.UserID), payload) - select { - case <-publishDone: - return - case <-ctx.Done(): - return - case <-ticker.C: - } - } - }() + // A single publish is sufficient because the subscription + // is active before websocket.Accept (and thus before Dial + // returns). This serves as a regression test for the fix. + err = api.Pubsub.Publish(coderdpubsub.ChatWatchEventChannel(user.UserID), payload) + require.NoError(t, err) var received codersdk.ChatWatchEvent for { @@ -1965,7 +1938,6 @@ func TestWatchChats(t *testing.T) { break } } - close(publishDone) // Verify the event carries the full DiffStatus. require.NotNil(t, received.Chat.DiffStatus, "diff_status_change event must include DiffStatus")