Merge pull request 'Fix agent panic on WebSocket disconnect' (#38)

The previous wording told operators to leave the slot blank and then
warned that doing so disables prune, which is contradictory. Prune is
gated on the admin slot for every backend, on both paths:
scheduled (maintenance_dispatch.go:52) skips silently, manual
(repo_ops.go:39) returns 400 admin_creds_required.

Also reworded the opening line, which still said 'only needed for
rest-server repos' while the new text asks other backends to fill it in.

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) //nolint:bodyclose // successful upgrades have a nil response body owned by conn
 	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)
+	}
+}

internal/server/http/pending_drain_test.go

@@ -512,11 +512,27 @@ func TestDrainPendingSerializesPerHost(t *testing.T) {
 	// Connect the agent so DrainPending can dispatch.
 	c := agentDial(t, srv, ts, hostID, token)
 	sendHello(t, c, "serialise-host")
-	// Drain the on-hello goroutine's pass first (no pending rows yet),
-	// then wait for the schedule.set so the connection is fully settled.
+	// Wait for the on-hello push to settle.
 	_ = drainUntil(t, c, api.MsgScheduleSet)
 
-	// Insert 5 pending rows now that the on-hello drain has already run.
+	// A real agent is always in a read loop. Keep this test client
+	// reading in the background for the rest of the test: without an
+	// active reader the server-side conn can be dropped under parallel
+	// load, which unregisters it from the hub and makes DrainPending
+	// no-op (conn == nil) — the historical source of this test's
+	// flakiness (it would observe 0 or a partial drain). The reader also
+	// consumes the command.run envelopes our drains emit.
+	readerCtx, stopReader := context.WithCancel(context.Background())
+	defer stopReader()
+	go func() {
+		for {
+			if _, _, err := c.Read(readerCtx); err != nil {
+				return
+			}
+		}
+	}()
+
+	// Insert 5 due pending rows.
 	now := time.Now().UTC()
 	for i := range 5 {
 		pid := ulid.Make().String()
@@ -533,7 +549,8 @@ func TestDrainPendingSerializesPerHost(t *testing.T) {
 		}
 	}
 
-	// Spawn 10 goroutines all calling DrainPending concurrently.
+	// Fire 10 concurrent DrainPending calls. The per-host mutex must
+	// ensure each row is dispatched at most once (no double-dispatch).
 	var wg sync.WaitGroup
 	for range 10 {
 		wg.Add(1)
@@ -544,24 +561,26 @@ func TestDrainPendingSerializesPerHost(t *testing.T) {
 	}
 	wg.Wait()
 
-	// Drain any envelopes the agent received so we don't block below.
-	// We read with short timeouts and stop when the connection goes quiet.
-	drainDeadline := time.Now().Add(500 * time.Millisecond)
-	for time.Now().Before(drainDeadline) {
-		ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
-		_, _, err := c.Read(ctx)
-		cancel()
-		if err != nil {
-			break
-		}
+	// Drain to completion. The fire-and-forget on-hello DrainPending
+	// shares the same per-host mutex and can hold it during the burst,
+	// leaving rows for a later pass — exactly how production drains
+	// (repeatedly, via the 30s tick / on reconnect). Re-drain until the
+	// queue is empty; because every drain is still serialised, each row
+	// is dispatched at most once, so the exactly-5 job count below proves
+	// there was no double-dispatch.
+	deadline := time.Now().Add(5 * time.Second)
+	for countPendingForHost(t, st, hostID) > 0 && time.Now().Before(deadline) {
+		srv.DrainPending(context.Background(), hostID)
+		time.Sleep(10 * time.Millisecond)
 	}
 
-	// All 5 pending rows must be gone.
+	// All 5 pending rows must be drained.
 	if n := countPendingForHost(t, st, hostID); n != 0 {
-		t.Errorf("pending rows after concurrent drain: got %d, want 0", n)
+		t.Errorf("pending rows after drain-to-completion: got %d, want 0", n)
 	}
 
-	// Exactly 5 backup job rows (one per pending row), not 10+ from a race.
+	// Exactly 5 backup job rows (one per pending row) — never more, which
+	// would mean the per-host mutex failed to prevent double-dispatch.
 	var n int
 	_ = st.DB().QueryRow(
 		`SELECT COUNT(*) FROM jobs WHERE host_id = ? AND kind = 'backup' AND actor_kind = 'schedule'`,

web/templates/pages/host_repo.html

@@ -82,10 +82,14 @@
         <div class="text-[12px] text-ok mb-3 mono">✓ saved</div>
       {{end}}
       <p class="text-[12.5px] text-ink-mid leading-[1.6] mb-4 max-w-[640px]">
-        Only needed for rest-server repos that distinguish an append-only
-        user (everyday backups) from a delete-capable user (prune /
-        forget). For S3 / B2 / SFTP / local, leave this blank — the
-        everyday repo credentials handle prune too.
+        Required for prune. On rest-server repos this is the
+        delete-capable user, as distinct from the append-only user used
+        for everyday backups. Note that <strong>forget</strong> always
+        runs with the everyday repo credentials, so those must have
+        delete authority. For S3 / B2 / SFTP / local, enter the same
+        delete-capable repository credentials here if you want prune
+        enabled. <strong>Prune is skipped when admin credentials are
+        unset</strong>, on any backend.
       </p>
       <div class="grid grid-cols-2 gap-4">
         <div>