fix(adkrest): limit /run_live WebSocket message size (#1664)

* fix(adkrest): limit /run_live WebSocket message size

* Address review: 16 MiB limit, unexported constant, mechanism test

Matches adk-python's uvicorn ws_max_size default (16 MiB). Unexports

the constant since nothing currently overrides it. Adds a

mechanism-level regression test for the SetReadLimit call: this

checkout's Agent interface has no way, via its public API, to build a

fake agent that supports live sessions (agent.New()'s returned type

does not implement the liveAgent interface RunLive() requires), so a

full RunLiveHandler integration test isn't achievable here without

separate, larger live-agent test scaffolding. This test instead

exercises the exact websocket.Upgrader + SetReadLimit configuration

the handler uses, in isolation.

* Address review: replace mechanism test with handler-level test

Removes runtime_live_limit_test.go (needed only because this branch

lacked live-agent test scaffolding at the time it was written) and

adds TestRunLiveHandlerEnforcesMessageSizeLimit to runtime_live_test.go

instead, using mockLiveAgent/dialRunLiveHandler to exercise

RunLiveHandler end to end, per review. golangci-lint's two reported

issues are not addressed by this commit -- see PR discussion.

---------

Co-authored-by: prasanna8585 <prasanna8585@users.noreply.github.com>
Co-authored-by: wolo <wolo@google.com>

Author: prasanna8585

server/adkrest/controllers/runtime.go

@@ -100,6 +100,16 @@ type RuntimeAPIControllerConfig struct {
 	CheckOrigin func(*http.Request) bool
 }
 
+// maxLiveMessageBytes is the read limit RunLiveHandler applies to a single
+// client-sent WebSocket message. Matches uvicorn's ws_max_size default,
+// which adk-python's dev servers (adk web, adk api_server) leave unset, so
+// a message either server accepts, the other does too.
+//
+// Unexported: nothing currently overrides it. If that's needed later, add
+// a field to RuntimeAPIControllerConfig where zero means this default, the
+// same way ServerConfig.MaxPayloadSize works.
+const maxLiveMessageBytes = 16 << 20 // 16 MiB
+
 // NewRuntimeAPIController creates the controller for the Runtime API.
 //
 // Deprecated: use [NewRuntimeAPIControllerWithConfig], which does not have to
@@ -423,6 +433,9 @@ func (c *RuntimeAPIController) RunLiveHandler(rw http.ResponseWriter, req *http.
 		_ = ws.Close()
 	}()
 
+	// The upgrade bypasses MaxBytesMiddleware, and gorilla/websocket has no default limit.
+	ws.SetReadLimit(maxLiveMessageBytes)
+
 	sendClose := func(code int, reason string) {
 		_ = ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, truncateCloseReason(reason)))
 		_ = ws.SetReadDeadline(time.Now().Add(time.Second))

server/adkrest/controllers/runtime_live_test.go

@@ -605,3 +605,42 @@ func TestTruncateCloseReason(t *testing.T) {
 		t.Errorf("close frame payload = %d bytes, want at most 125", got)
 	}
 }
+
+func TestRunLiveHandlerEnforcesMessageSizeLimit(t *testing.T) {
+	for _, tc := range []struct {
+		name    string
+		size    int
+		forward bool
+	}{
+		{name: "at limit is forwarded", size: maxLiveMessageBytes, forward: true},
+		{name: "one byte over is rejected", size: maxLiveMessageBytes + 1},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			liveSession := newRecordingLiveSession()
+			conn, handlerDone := dialRunLiveHandler(t, func(agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) {
+				return liveSession, func(yield func(*session.Event, error) bool) { <-liveSession.closed }, nil
+			})
+			if err := conn.WriteMessage(websocket.BinaryMessage, make([]byte, tc.size)); err != nil {
+				t.Fatalf("WriteMessage() failed: %v", err)
+			}
+
+			if tc.forward {
+				got := waitForLiveRequest(t, liveSession)
+				if blob, ok := got.RealtimeInput.(*genai.Blob); !ok || len(blob.Data) != tc.size {
+					t.Fatalf("RealtimeInput = %T, want a *genai.Blob of %d bytes", got.RealtimeInput, tc.size)
+				}
+				return
+			}
+
+			if closeErr := readCloseError(t, conn); closeErr.Code != websocket.CloseMessageTooBig {
+				t.Fatalf("close code = %d, want %d", closeErr.Code, websocket.CloseMessageTooBig)
+			}
+			select {
+			case req := <-liveSession.requests:
+				t.Fatalf("oversized message reached the live session: %v", req)
+			default:
+			}
+			waitForHandlerExit(t, handlerDone)
+		})
+	}
+}