chore: use WebsocketNetConn

DanielleMaywood · DanielleMaywood · commit 04a92a48abf8 · 2025-07-10T12:25:27.000Z
diff --git a/agent/agentcontainers/api.go b/agent/agentcontainers/api.go
@@ -2,6 +2,7 @@ package agentcontainers
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"maps"
@@ -29,7 +30,6 @@ import (
 	"github.com/coder/coder/v2/coderd/httpapi"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
-	"github.com/coder/coder/v2/codersdk/wsjson"
 	"github.com/coder/coder/v2/provisioner"
 	"github.com/coder/quartz"
 	"github.com/coder/websocket"
@@ -561,13 +561,10 @@ func (api *API) watchContainers(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	ctx = api.ctx
+	ctx, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageText)
+	defer wsNetConn.Close()
 
 	go httpapi.Heartbeat(ctx, conn)
-	defer conn.Close(websocket.StatusNormalClosure, "connection closed")
-
-	encoder := wsjson.NewEncoder[codersdk.WorkspaceAgentListContainersResponse](conn, websocket.MessageText)
-	defer encoder.Close(websocket.StatusNormalClosure)
 
 	updateCh := make(chan struct{}, 1)
 
@@ -586,6 +583,9 @@ func (api *API) watchContainers(rw http.ResponseWriter, r *http.Request) {
 
 	for {
 		select {
+		case <-api.ctx.Done():
+			return
+
 		case <-ctx.Done():
 			return
 
@@ -596,7 +596,7 @@ func (api *API) watchContainers(rw http.ResponseWriter, r *http.Request) {
 				continue
 			}
 
-			if err := encoder.Encode(ct); err != nil {
+			if err := json.NewEncoder(wsNetConn).Encode(ct); err != nil {
 				api.logger.Error(ctx, "encode container list", slog.Error(err))
 				return
 			}
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -817,7 +817,7 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re
 
 	// If the agent is unreachable, the request will hang. Assume that if we
 	// don't get a response after 30s that the agent is unreachable.
-	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	dialCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
 	defer cancel()
 	apiAgent, err := db2sdk.WorkspaceAgent(
 		api.DERPMap(),
@@ -843,7 +843,7 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re
 		return
 	}
 
-	agentConn, release, err := api.agentProvider.AgentConn(ctx, workspaceAgent.ID)
+	agentConn, release, err := api.agentProvider.AgentConn(dialCtx, workspaceAgent.ID)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 			Message: "Internal error dialing workspace agent.",
@@ -873,21 +873,21 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re
 		return
 	}
 
-	ctx = api.ctx
+	ctx, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageText)
+	defer wsNetConn.Close()
 
 	go httpapi.Heartbeat(ctx, conn)
-	defer conn.Close(websocket.StatusNormalClosure, "connection closed")
-
-	encoder := wsjson.NewEncoder[codersdk.WorkspaceAgentListContainersResponse](conn, websocket.MessageText)
-	defer encoder.Close(websocket.StatusNormalClosure)
 
 	for {
 		select {
+		case <-api.ctx.Done():
+			return
+
 		case <-ctx.Done():
 			return
 
 		case containers := <-containersCh:
-			if err := encoder.Encode(containers); err != nil {
+			if err := json.NewEncoder(wsNetConn).Encode(containers); err != nil {
 				api.Logger.Error(ctx, "encode containers", slog.Error(err))
 				return
 			}
