@@ -197,13 +197,12 @@ func (api *API) workspaceAgentListen(rw http.ResponseWriter, r *http.Request) {
197
197
return
198
198
}
199
199
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.
203
202
204
203
config := yamux .DefaultConfig ()
205
204
config .LogOutput = io .Discard
206
- session , err := yamux .Server (websocket . NetConn ( r . Context (), conn , websocket . MessageBinary ) , config )
205
+ session , err := yamux .Server (wsNetConn , config )
207
206
if err != nil {
208
207
_ = conn .Close (websocket .StatusAbnormalClosure , err .Error ())
209
208
return
@@ -233,7 +232,7 @@ func (api *API) workspaceAgentListen(rw http.ResponseWriter, r *http.Request) {
233
232
}
234
233
disconnectedAt := workspaceAgent .DisconnectedAt
235
234
updateConnectionTimes := func () error {
236
- err = api .Database .UpdateWorkspaceAgentConnectionByID (r . Context () , database.UpdateWorkspaceAgentConnectionByIDParams {
235
+ err = api .Database .UpdateWorkspaceAgentConnectionByID (ctx , database.UpdateWorkspaceAgentConnectionByIDParams {
237
236
ID : workspaceAgent .ID ,
238
237
FirstConnectedAt : firstConnectedAt ,
239
238
LastConnectedAt : lastConnectedAt ,
@@ -259,7 +258,7 @@ func (api *API) workspaceAgentListen(rw http.ResponseWriter, r *http.Request) {
259
258
return
260
259
}
261
260
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 ))
263
262
264
263
ticker := time .NewTicker (api .AgentConnectionUpdateFrequency )
265
264
defer ticker .Stop ()
0 commit comments