diff --git a/coderd/inboxnotifications.go b/coderd/inboxnotifications.go index 5437165bb71a6..ebb2a08dfe7eb 100644 --- a/coderd/inboxnotifications.go +++ b/coderd/inboxnotifications.go @@ -94,18 +94,6 @@ func (api *API) watchInboxNotifications(rw http.ResponseWriter, r *http.Request) return } - conn, err := websocket.Accept(rw, r, nil) - if err != nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Failed to upgrade connection to websocket.", - Detail: err.Error(), - }) - return - } - - go httpapi.Heartbeat(ctx, conn) - defer conn.Close(websocket.StatusNormalClosure, "connection closed") - notificationCh := make(chan codersdk.InboxNotification, 10) closeInboxNotificationsSubscriber, err := api.Pubsub.SubscribeWithErr(pubsub.InboxNotificationForOwnerEventChannel(apikey.UserID), @@ -161,9 +149,20 @@ func (api *API) watchInboxNotifications(rw http.ResponseWriter, r *http.Request) api.Logger.Error(ctx, "subscribe to inbox notification event", slog.Error(err)) return } - defer closeInboxNotificationsSubscriber() + conn, err := websocket.Accept(rw, r, nil) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to upgrade connection to websocket.", + Detail: err.Error(), + }) + return + } + + go httpapi.Heartbeat(ctx, conn) + defer conn.Close(websocket.StatusNormalClosure, "connection closed") + encoder := wsjson.NewEncoder[codersdk.GetInboxNotificationResponse](conn, websocket.MessageText) defer encoder.Close(websocket.StatusNormalClosure) diff --git a/coderd/inboxnotifications_test.go b/coderd/inboxnotifications_test.go index 81e119381d281..4253733300e14 100644 --- a/coderd/inboxnotifications_test.go +++ b/coderd/inboxnotifications_test.go @@ -122,7 +122,8 @@ func TestInboxNotification_Watch(t *testing.T) { }, "notification title", "notification content", nil) require.NoError(t, err) - dispatchFunc(ctx, uuid.New()) + _, err = dispatchFunc(ctx, uuid.New()) + require.NoError(t, err) _, message, err := wsConn.Read(ctx) require.NoError(t, err) @@ -174,7 +175,8 @@ func TestInboxNotification_Watch(t *testing.T) { }, "memory related title", "memory related content", nil) require.NoError(t, err) - dispatchFunc(ctx, uuid.New()) + _, err = dispatchFunc(ctx, uuid.New()) + require.NoError(t, err) _, message, err := wsConn.Read(ctx) require.NoError(t, err) @@ -193,7 +195,8 @@ func TestInboxNotification_Watch(t *testing.T) { }, "disk related title", "disk related title", nil) require.NoError(t, err) - dispatchFunc(ctx, uuid.New()) + _, err = dispatchFunc(ctx, uuid.New()) + require.NoError(t, err) dispatchFunc, err = inboxHandler.Dispatcher(types.MessagePayload{ UserID: memberClient.ID.String(), @@ -201,7 +204,8 @@ func TestInboxNotification_Watch(t *testing.T) { }, "second memory related title", "second memory related title", nil) require.NoError(t, err) - dispatchFunc(ctx, uuid.New()) + _, err = dispatchFunc(ctx, uuid.New()) + require.NoError(t, err) _, message, err = wsConn.Read(ctx) require.NoError(t, err) @@ -256,7 +260,8 @@ func TestInboxNotification_Watch(t *testing.T) { }, "memory related title", "memory related content", nil) require.NoError(t, err) - dispatchFunc(ctx, uuid.New()) + _, err = dispatchFunc(ctx, uuid.New()) + require.NoError(t, err) _, message, err := wsConn.Read(ctx) require.NoError(t, err) @@ -276,7 +281,8 @@ func TestInboxNotification_Watch(t *testing.T) { }, "second memory related title", "second memory related title", nil) require.NoError(t, err) - dispatchFunc(ctx, uuid.New()) + _, err = dispatchFunc(ctx, uuid.New()) + require.NoError(t, err) dispatchFunc, err = inboxHandler.Dispatcher(types.MessagePayload{ UserID: memberClient.ID.String(), @@ -285,7 +291,8 @@ func TestInboxNotification_Watch(t *testing.T) { }, "another memory related title", "another memory related title", nil) require.NoError(t, err) - dispatchFunc(ctx, uuid.New()) + _, err = dispatchFunc(ctx, uuid.New()) + require.NoError(t, err) _, message, err = wsConn.Read(ctx) require.NoError(t, err)