internal/agent/wsclient/client.go

@@ -103,15 +103,14 @@ func connectOnce(ctx context.Context, cfg Config, handle Handler) error {
 	}
 
 	dialCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
-	conn, res, err := websocket.Dial(dialCtx, wsURL, dialOpts)
+	conn, _, err := websocket.Dial(dialCtx, wsURL, dialOpts)
 	cancel()
 	if err != nil {
 		return fmt.Errorf("dial: %w", err)
 	}
-	// websocket.Dial returns the upgrade response separately from the
-	// conn. Body is empty on a successful upgrade but Go's net/http
-	// still expects it closed to release the connection.
-	defer func() { _ = res.Body.Close() }()
+	// On a successful upgrade coder/websocket transfers ownership of the
+	// response stream to conn and deliberately sets res.Body to nil. Closing
+	// the connection below releases that stream.
 	defer conn.CloseNow() //nolint:errcheck
 
 	// Send hello.

internal/agent/wsclient/client_test.go

@@ -0,0 +1,46 @@
+package wsclient
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/coder/websocket"
+)
+
+func TestConnectOnceCleanDisconnectDoesNotPanic(t *testing.T) {
+	serverErr := make(chan error, 1)
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := websocket.Accept(w, r, nil)
+		if err != nil {
+			serverErr <- err
+			return
+		}
+		defer conn.CloseNow() //nolint:errcheck
+
+		// Wait for the agent hello so Dial and the first client write have both
+		// completed before ending the connection normally.
+		if _, _, err := conn.Read(r.Context()); err != nil {
+			serverErr <- err
+			return
+		}
+		serverErr <- conn.Close(websocket.StatusNormalClosure, "test complete")
+	}))
+	defer srv.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	err := connectOnce(ctx, Config{
+		ServerURL:       srv.URL,
+		AgentToken:      "test-token",
+		HeartbeatPeriod: time.Hour,
+	}, nil)
+	if err == nil {
+		t.Fatal("connectOnce returned nil after server disconnected")
+	}
+	if err := <-serverErr; err != nil {
+		t.Fatalf("server websocket: %v", err)
+	}
+}