@@ -197,13 +197,12 @@ func (api *API) workspaceAgentListen(rw http.ResponseWriter, r *http.Request) {
197197return
198198}
199199
200- defer func () {
201- _ = conn .Close (websocket .StatusNormalClosure ,"" )
202- }()
200+ ctx ,wsNetConn := websocketNetConn (r .Context (),conn ,websocket .MessageBinary )
201+ defer wsNetConn .Close ()// Also closes conn.
203202
204203config := yamux .DefaultConfig ()
205204config .LogOutput = io .Discard
206- session ,err := yamux .Server (websocket . NetConn ( r . Context (), conn , websocket . MessageBinary ) ,config )
205+ session ,err := yamux .Server (wsNetConn ,config )
207206if err != nil {
208207_ = conn .Close (websocket .StatusAbnormalClosure ,err .Error ())
209208return
@@ -233,7 +232,7 @@ func (api *API) workspaceAgentListen(rw http.ResponseWriter, r *http.Request) {
233232}
234233disconnectedAt := workspaceAgent .DisconnectedAt
235234updateConnectionTimes := func ()error {
236- err = api .Database .UpdateWorkspaceAgentConnectionByID (r . Context () , database.UpdateWorkspaceAgentConnectionByIDParams {
235+ err = api .Database .UpdateWorkspaceAgentConnectionByID (ctx , database.UpdateWorkspaceAgentConnectionByIDParams {
237236ID :workspaceAgent .ID ,
238237FirstConnectedAt :firstConnectedAt ,
239238LastConnectedAt :lastConnectedAt ,
@@ -259,7 +258,7 @@ func (api *API) workspaceAgentListen(rw http.ResponseWriter, r *http.Request) {
259258return
260259}
261260
262- api .Logger .Info (r . Context () ,"accepting agent" ,slog .F ("resource" ,resource ),slog .F ("agent" ,workspaceAgent ))
261+ api .Logger .Info (ctx ,"accepting agent" ,slog .F ("resource" ,resource ),slog .F ("agent" ,workspaceAgent ))
263262
264263ticker := time .NewTicker (api .AgentConnectionUpdateFrequency )
265264defer ticker .Stop ()