From 39aff83837d030512ef0a745069066d6f890fc5e Mon Sep 17 00:00:00 2001
From: Steve Cliff <steve@devcloud.guru>
Date: Sat, 22 Aug 2026 09:45:20 +0100
Subject: [PATCH 1/2] fix(agent): avoid panic on websocket disconnect

---
 internal/agent/wsclient/client.go      |  9 +++--
 internal/agent/wsclient/client_test.go | 46 ++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 5 deletions(-)
 create mode 100644 internal/agent/wsclient/client_test.go

diff --git a/internal/agent/wsclient/client.go b/internal/agent/wsclient/client.go
index 77f3ee1..171912f 100644
--- a/internal/agent/wsclient/client.go
+++ b/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.
diff --git a/internal/agent/wsclient/client_test.go b/internal/agent/wsclient/client_test.go
new file mode 100644
index 0000000..5be6220
--- /dev/null
+++ b/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)
+	}
+}
-- 
2.52.0


From dfe082629f1d41e55f2cb2d19efc3773a85bdaff Mon Sep 17 00:00:00 2001
From: Steve Cliff <steve@devcloud.guru>
Date: Sat, 22 Aug 2026 09:46:50 +0100
Subject: [PATCH 2/2] chore(lint): document websocket response ownership

---
 internal/agent/wsclient/client.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/internal/agent/wsclient/client.go b/internal/agent/wsclient/client.go
index 171912f..625b7f7 100644
--- a/internal/agent/wsclient/client.go
+++ b/internal/agent/wsclient/client.go
@@ -103,7 +103,7 @@ func connectOnce(ctx context.Context, cfg Config, handle Handler) error {
 	}
 
 	dialCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
-	conn, _, err := websocket.Dial(dialCtx, wsURL, dialOpts)
+	conn, _, err := websocket.Dial(dialCtx, wsURL, dialOpts) //nolint:bodyclose // successful upgrades have a nil response body owned by conn
 	cancel()
 	if err != nil {
 		return fmt.Errorf("dial: %w", err)
-- 
2.52.0