@@ -817,12 +817,13 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re
817817var (
818818ctx = r .Context ()
819819workspaceAgent = httpmw .WorkspaceAgentParam (r )
820+ logger = api .Logger .Named ("agent_container_watcher" ).With (slog .F ("agent_id" ,workspaceAgent .ID ))
820821)
821822
822823// If the agent is unreachable, the request will hang. Assume that if we
823824// don't get a response after 30s that the agent is unreachable.
824- dialCtx ,cancel := context .WithTimeout (ctx ,30 * time .Second )
825- defer cancel ()
825+ dialCtx ,dialCancel := context .WithTimeout (ctx ,30 * time .Second )
826+ defer dialCancel ()
826827apiAgent ,err := db2sdk .WorkspaceAgent (
827828api .DERPMap (),
828829* api .TailnetCoordinator .Load (),
@@ -857,8 +858,7 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re
857858}
858859defer release ()
859860
860- watcherLogger := api .Logger .Named ("agent_container_watcher" ).With (slog .F ("agent_id" ,workspaceAgent .ID ))
861- containersCh ,closer ,err := agentConn .WatchContainers (ctx ,watcherLogger )
861+ containersCh ,closer ,err := agentConn .WatchContainers (ctx ,logger )
862862if err != nil {
863863httpapi .Write (ctx ,rw ,http .StatusInternalServerError , codersdk.Response {
864864Message :"Internal error watching agent's containers." ,
@@ -877,14 +877,17 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re
877877return
878878}
879879
880+ ctx ,cancel := context .WithCancel (r .Context ())
881+ defer cancel ()
882+
880883// Here we close the websocket for reading, so that the websocket library will handle pings and
881884// close frames.
882885_ = conn .CloseRead (context .Background ())
883886
884887ctx ,wsNetConn := codersdk .WebsocketNetConn (ctx ,conn ,websocket .MessageText )
885888defer wsNetConn .Close ()
886889
887- go httpapi .Heartbeat (ctx ,conn )
890+ go httpapi .HeartbeatClose (ctx , logger , cancel ,conn )
888891
889892encoder := json .NewEncoder (wsNetConn )
890893