fix(agent-runner): bound unauthenticated ACP bridge handshakes

docs/ai-agent/agent-runner-service.md

@@ -171,6 +171,8 @@ There is no REST equivalent to `services/ai-agent`'s `/conversations/:id/pause`
 
 The three internal endpoints require `X-Internal-Token` matching `INTERNAL_API_KEY` — `services/api`'s `AI_AGENT_INTERNAL_KEY` must equal this value.
 
+`/agent-bridge/ws` authenticates after the upgrade, with the `hello` frame (agent id plus bridge token). Because that handshake is unauthenticated until the frame arrives, it is bounded: at most `maxPendingHellos` (128) handshakes in flight per agent-runner process — excess upgrades get `503` before the socket is accepted — the `hello` frame may not exceed 1 KiB, and a peer that sends no `hello` within 10 s is closed. The 10 MiB bridge frame limit applies only after the token checks out. Anyone exempting this path from an edge/browser authentication layer inherits that bound and must not remove it.
+
 ---
 
 ## Environment Variables

services/agent-runner/internal/acpbridge/server.go

@@ -27,6 +27,19 @@ import (
 // _HELLO_TIMEOUT_SECONDS.
 const helloTimeout = 10 * time.Second
 
+// helloReadLimit bounds the pre-auth frame. A hello carries only an agent id
+// and a token, so a peer that has not authenticated yet must not be able to
+// allocate the full bridge budget. The limit is raised once the token checks
+// out, below in handleWS.
+const helloReadLimit = 1 << 10
+
+// maxPendingHellos caps concurrent unauthenticated handshakes. Excess
+// upgrades are refused up front instead of occupying the 10s hello window.
+const maxPendingHellos = 128
+
+// helloSlots is the process-wide budget of unauthenticated handshakes.
+var helloSlots = make(chan struct{}, maxPendingHellos)
+
 // bridgeMessageReadLimit overrides coder/websocket's 32KiB default (see
 // handleWS) — generous enough for a large tool-call payload (verbose
 // command output, a sizeable file diff) without leaving this per-connection
@@ -145,19 +158,28 @@ type helloFrame struct {
 // self-describing handshake that degrades gracefully across daemon/server
 // version skew, which headers rejected pre-accept wouldn't).
 func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) {
+	select {
+	case helloSlots <- struct{}{}:
+	default:
+		// Refuse before upgrading: a saturated service must not accept more
+		// sockets it cannot authenticate or bound.
+		http.Error(w, "too many pending bridge handshakes", http.StatusServiceUnavailable)
+		return
+	}
+	held := true
+	release := func() {
+		if held {
+			<-helloSlots
+			held = false
+		}
+	}
+	defer release()
+
 	conn, err := websocket.Accept(w, r, nil)
 	if err != nil {
 		return
 	}
-	// coder/websocket defaults to a 32KiB read limit — comfortably too
-	// small for a bridge daemon relaying a raw ACP event verbatim (a large
-	// command's output, or a sizeable file diff, both routinely exceed it),
-	// which surfaced as the daemon getting disconnected mid-conversation
-	// with "received close frame: status = StatusMessageTooBig". Bridge
-	// events are the only traffic this endpoint ever reads at any real
-	// size (hello/ping frames are tiny), so this is scoped to the whole
-	// connection rather than per-message.
-	conn.SetReadLimit(bridgeMessageReadLimit)
+	conn.SetReadLimit(helloReadLimit)
 
 	helloCtx, cancel := context.WithTimeout(r.Context(), helloTimeout)
 	var hello helloFrame
@@ -186,6 +208,13 @@ func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	// Authenticated: free the handshake budget and allow bridge-sized frames.
+	release()
+	conn.SetReadLimit(bridgeMessageReadLimit)
+
+	// TODO: cover this raise with a post-auth frame above helloReadLimit
+	// (the e2e fake bridge currently sends only a small event payload).
+
 	sessionID, err := s.Registry.Register(r.Context(), agentID, cfg.ProjectID, wsConn{conn})
 	if err != nil {
 		s.Log.Error("acpbridge: failed to register bridge connection", "agent_id", agentID, "error", err)

services/agent-runner/internal/acpbridge/server_hello_test.go

@@ -0,0 +1,102 @@
+package acpbridge
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"log/slog"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/coder/websocket"
+)
+
+func newHelloTestServer(t *testing.T) *httptest.Server {
+	t.Helper()
+	s := &Server{Log: slog.New(slog.NewTextHandler(io.Discard, nil))}
+	ts := httptest.NewServer(s.Routes())
+	t.Cleanup(ts.Close)
+	return ts
+}
+
+func dialBridge(t *testing.T, ts *httptest.Server) *websocket.Conn {
+	t.Helper()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	conn, dialResp, err := websocket.Dial(ctx, "ws"+strings.TrimPrefix(ts.URL, "http")+"/agent-bridge/ws", nil)
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	if dialResp != nil && dialResp.Body != nil {
+		_ = dialResp.Body.Close()
+	}
+	t.Cleanup(func() { _ = conn.CloseNow() })
+	return conn
+}
+
+// A peer that has not authenticated must not be able to write bridge-sized
+// frames: the pre-auth limit is the only thing standing between an anonymous
+// socket and the relay's 10 MiB buffer.
+func TestHandleWSRejectsOversizedPreAuthFrame(t *testing.T) {
+	ts := newHelloTestServer(t)
+	conn := dialBridge(t, ts)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	_ = conn.Write(ctx, websocket.MessageText, bytes.Repeat([]byte("a"), 2*helloReadLimit))
+	_, _, err := conn.Read(ctx)
+	// 1009 proves the frame was refused by the pre-auth limit; a malformed
+	// hello under the bridge-sized limit closes with 1007 instead, because
+	// wsjson.Read rejects the JSON before the handler's 4400 close runs.
+	if status := websocket.CloseStatus(err); status != websocket.StatusMessageTooBig {
+		t.Fatalf("close status = %d, want %d (err=%v)", status, websocket.StatusMessageTooBig, err)
+	}
+}
+
+func TestHandleWSRefusesWhenHandshakeBudgetIsExhausted(t *testing.T) {
+	ts := newHelloTestServer(t)
+	for i := 0; i < maxPendingHellos; i++ {
+		helloSlots <- struct{}{}
+	}
+	t.Cleanup(func() {
+		for i := 0; i < maxPendingHellos; i++ {
+			<-helloSlots
+		}
+	})
+
+	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, ts.URL+"/agent-bridge/ws", nil)
+	if err != nil {
+		t.Fatalf("request: %v", err)
+	}
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		t.Fatalf("request: %v", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusServiceUnavailable {
+		t.Fatalf("status = %d, want %d", resp.StatusCode, http.StatusServiceUnavailable)
+	}
+}
+
+func TestHandleWSReleasesHandshakeSlotAfterRejectedHello(t *testing.T) {
+	ts := newHelloTestServer(t)
+	conn := dialBridge(t, ts)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	_ = conn.Write(ctx, websocket.MessageText, []byte(`{"type":"hello","agent_id":"not-a-uuid","token":"x"}`))
+	if _, _, err := conn.Read(ctx); err == nil {
+		t.Fatal("expected the connection to be closed for an invalid agent_id")
+	}
+
+	deadline := time.Now().Add(2 * time.Second)
+	for len(helloSlots) != 0 {
+		if time.Now().After(deadline) {
+			t.Fatalf("handshake slot not released, %d still held", len(helloSlots))
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+}