fix(coderd): ensure inbox WebSocket is closed when client disconnects… (#21685)

johnstcn · DanielleMaywood · web-flow · commit cc88b0f479e3 · 2026-01-26T09:27:59.000-06:00
… (#21652) Relates to #19715 This is similar to #19711 This endpoint works by doing the following: - Subscribing to the database's with pubsub - Accepts a WebSocket upgrade - Starts a `httpapi.Heartbeat` - Creates a json encoder - **Infinitely loops waiting for notification until request context cancelled** The critical issue here is that `httpapi.Heartbeat` silently fails when the client has disconnected. This means we never cancel the request context, leaving the WebSocket alive until we receive a notification from the database and fail to write that down the pipe. By replacing usage of `httpapi.Heartbeat` with `httpapi.HeartbeatClose`, we cancel the context _when the heartbeat fails to write_ due to the client disconnecting. This allows us to cleanup without waiting for a notification to come through the pubsub channel. (cherry picked from commit 409360c)  Co-authored-by: Danielle Maywood <danielle@themaywoods.com>
diff --git a/coderd/inboxnotifications.go b/coderd/inboxnotifications.go
@@ -21,7 +21,6 @@ import (
 	"github.com/coder/coder/v2/coderd/pubsub"
 	markdown "github.com/coder/coder/v2/coderd/render"
 	"github.com/coder/coder/v2/codersdk"
-	"github.com/coder/coder/v2/codersdk/wsjson"
 	"github.com/coder/websocket"
 )
 
@@ -127,6 +126,7 @@ func (api *API) watchInboxNotifications(rw http.ResponseWriter, r *http.Request)
 		templates  = p.UUIDs(vals, []uuid.UUID{}, "templates")
 		readStatus = p.String(vals, "all", "read_status")
 		format     = p.String(vals, notificationFormatMarkdown, "format")
+		logger     = api.Logger.Named("inbox_notifications_watcher")
 	)
 	p.ErrorExcessParams(vals)
 	if len(p.Errors) > 0 {
@@ -214,11 +214,17 @@ func (api *API) watchInboxNotifications(rw http.ResponseWriter, r *http.Request)
 		return
 	}
 
-	go httpapi.Heartbeat(ctx, conn)
-	defer conn.Close(websocket.StatusNormalClosure, "connection closed")
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
 
-	encoder := wsjson.NewEncoder[codersdk.GetInboxNotificationResponse](conn, websocket.MessageText)
-	defer encoder.Close(websocket.StatusNormalClosure)
+	_ = conn.CloseRead(context.Background())
+
+	ctx, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageText)
+	defer wsNetConn.Close()
+
+	go httpapi.HeartbeatClose(ctx, logger, cancel, conn)
+
+	encoder := json.NewEncoder(wsNetConn)
 
 	// Log the request immediately instead of after it completes.
 	if rl := loggermw.RequestLoggerFromContext(ctx); rl != nil {
@@ -227,8 +233,12 @@ func (api *API) watchInboxNotifications(rw http.ResponseWriter, r *http.Request)
 
 	for {
 		select {
+		case <-api.ctx.Done():
+			return
+
 		case <-ctx.Done():
 			return
+
 		case notif := <-notificationCh:
 			unreadCount, err := api.Database.CountUnreadInboxNotificationsByUserID(ctx, apikey.UserID)
 			if err != nil {
