fix(agent): avoid panic on websocket disconnect

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)
+	}
+}