mcp: setup benchmark test for MCPProxy (envoyproxy#1593)

zirain · nacx · web-flow · commit e55c5ae23cb4 · 2025-12-04T11:18:50.000Z
**Description** This PR create a benchmark test for AIGW withe diffrent `mcp-session-encryption-iterations` Run this benchmark with following: ```console go test -timeout=15m -run='^$$' -bench=. -benchmem -benchtime=1x ./tests/bench/... 2025/12/03 13:50:12 starting DUMB MCP Streamable-HTTP server on :8080 at /mcp goos: darwin goarch: arm64 pkg: github.com/envoyproxy/ai-gateway/tests/bench cpu: Apple M1 Pro BenchmarkMCP/BaseLine-10 1 31305583 ns/op 49736 B/op 720 allocs/op BenchmarkMCP/Iterations_100-10 1 1394125 ns/op 22640 B/op 300 allocs/op PASS ok github.com/envoyproxy/ai-gateway/tests/bench 5.394s ``` ~~This result that the default `SessionCrypto` cost a lot, we'd better improve it or change the default setting.~~ cc @nacx --------- Signed-off-by: zirain <zirain2009@gmail.com> Signed-off-by: Ignasi Barrera <nacx@apache.org> Co-authored-by: Ignasi Barrera <nacx@apache.org> Co-authored-by: Ignasi Barrera <ignasi@tetrate.io>
diff --git a/cmd/aigw/run.go b/cmd/aigw/run.go
@@ -138,14 +138,15 @@ func run(ctx context.Context, c cmdRun, o *runOpts, stdout, stderr io.Writer) er
 	// Do the translation of the given AI Gateway resources Yaml into Envoy Gateway resources and write them to the file.
 	resourcesBuf := &bytes.Buffer{}
 	runCtx := &runCmdContext{
-		isDebug:                  c.Debug,
-		envoyGatewayResourcesOut: resourcesBuf,
-		stderrLogger:             debugLogger,
-		stderr:                   stderr,
-		tmpdir:                   filepath.Dir(o.logPath), // runDir
-		udsPath:                  o.extprocUDSPath,
-		adminPort:                c.AdminPort,
-		extProcLauncher:          o.extProcLauncher,
+		isDebug:                        c.Debug,
+		envoyGatewayResourcesOut:       resourcesBuf,
+		stderrLogger:                   debugLogger,
+		stderr:                         stderr,
+		tmpdir:                         filepath.Dir(o.logPath), // runDir
+		udsPath:                        o.extprocUDSPath,
+		adminPort:                      c.AdminPort,
+		extProcLauncher:                o.extProcLauncher,
+		mcpSessionEncryptionIterations: c.MCPSessionEncryptionIterations,
 	}
 	// If any of the configured MCP servers is using stdio, set up the streamable HTTP proxies for them
 	if err = proxyStdioMCPServers(ctx, debugLogger, c.mcpConfig); err != nil {
diff --git a/tests/bench/aigw.yaml b/tests/bench/aigw.yaml
@@ -0,0 +1,123 @@
+# Copyright Envoy AI Gateway Authors
+# SPDX-License-Identifier: Apache-2.0
+# The full text of the Apache license is available in the LICENSE file at
+# the root of the repo.
+
+---
+apiVersion: aigateway.envoyproxy.io/v1alpha1
+kind: MCPRoute
+metadata:
+  name: mcp-route
+  namespace: default
+spec:
+  parentRefs:
+    - name: aigw-run
+      kind: Gateway
+      group: gateway.networking.k8s.io
+  path: "/mcp"
+  backendRefs:
+    - name: test
+      kind: Backend
+      group: gateway.envoyproxy.io
+      path: "/mcp"
+---
+###################################################################################
+############################### Backend Definitions ###############################
+###################################################################################
+apiVersion: gateway.envoyproxy.io/v1alpha1
+kind: Backend
+metadata:
+  name: test
+  namespace: default
+spec:
+  endpoints:
+    - ip:
+        address: 127.0.0.1
+        port: 8080
+---
+###################################################################################
+############################### Gateway Definitions ###############################
+###################################################################################
+apiVersion: gateway.networking.k8s.io/v1
+kind: GatewayClass
+metadata:
+  name: aigw-run
+spec:
+  controllerName: gateway.envoyproxy.io/gatewayclass-controller
+---
+apiVersion: gateway.networking.k8s.io/v1
+kind: Gateway
+metadata:
+  name: aigw-run
+  namespace: default
+spec:
+  gatewayClassName: aigw-run
+  listeners:
+    - name: http
+      protocol: HTTP
+      port: 1975
+  infrastructure:
+    parametersRef:
+      group: gateway.envoyproxy.io
+      kind: EnvoyProxy
+      name: envoy-ai-gateway
+---
+apiVersion: gateway.envoyproxy.io/v1alpha1
+kind: EnvoyProxy
+metadata:
+  name: envoy-ai-gateway
+  namespace: default
+spec:
+  logging:
+    level:
+      default: error
+  bootstrap:
+    type: Merge
+    value: |-
+      admin:
+        address:
+          socket_address:
+            address: 127.0.0.1
+            port_value: 9901
+  telemetry:
+    accessLog:
+      settings:
+        - sinks:
+            - type: File
+              file:
+                path: /dev/stdout
+          format:
+            type: JSON
+            json:
+              # MCP specific fields
+              mcp_request_id: "%DYNAMIC_METADATA(io.envoy.ai_gateway:mcp_request_id)%"
+              mcp_session_id: "%REQ(MCP-SESSION-ID)%"
+              mcp_method: "%DYNAMIC_METADATA(io.envoy.ai_gateway:mcp_method)%"
+              mcp_backend: "%DYNAMIC_METADATA(io.envoy.ai_gateway:mcp_backend)%"
+              # Default fields
+              start_time: "%START_TIME%"
+              method: "%REQ(:METHOD)%"
+              x-envoy-origin-path: "%REQ(X-ENVOY-ORIGINAL-PATH?:PATH)%"
+              protocol: "%PROTOCOL%"
+              response_code: "%RESPONSE_CODE%"
+              response_flags: "%RESPONSE_FLAGS%"
+              response_code_details: "%RESPONSE_CODE_DETAILS%"
+              connection_termination_details: "%CONNECTION_TERMINATION_DETAILS%"
+              upstream_transport_failure_reason: "%UPSTREAM_TRANSPORT_FAILURE_REASON%"
+              bytes_received: "%BYTES_RECEIVED%"
+              bytes_sent: "%BYTES_SENT%"
+              duration: "%DURATION%"
+              x-envoy-upstream-service-time: "%RESP(X-ENVOY-UPSTREAM-SERVICE-TIME)%"
+              x-forwarded-for: "%REQ(X-FORWARDED-FOR)%"
+              user-agent: "%REQ(USER-AGENT)%"
+              x-request-id: "%REQ(X-REQUEST-ID)%"
+              ":authority": "%REQ(:AUTHORITY)%"
+              upstream_host: "%UPSTREAM_HOST%"
+              upstream_cluster: "%UPSTREAM_CLUSTER%"
+              upstream_local_address: "%UPSTREAM_LOCAL_ADDRESS%"
+              downstream_local_address: "%DOWNSTREAM_LOCAL_ADDRESS%"
+              downstream_remote_address: "%DOWNSTREAM_REMOTE_ADDRESS%"
+              requested_server_name: "%REQUESTED_SERVER_NAME%"
+              route_name: "%ROUTE_NAME%"
+---
+
diff --git a/tests/bench/bench_test.go b/tests/bench/bench_test.go
@@ -0,0 +1,167 @@
+// Copyright Envoy AI Gateway Authors
+// SPDX-License-Identifier: Apache-2.0
+// The full text of the Apache license is available in the LICENSE file at
+// the root of the repo.
+
+// 1. Build AIGW
+//  	make clean build.aigw
+// 2. Run the bench test
+//   	go test -timeout=15m -run='^$' -bench=. ./tests/bench/...
+
+package bench
+
+import (
+	"context"
+	_ "embed"
+	"fmt"
+	"net"
+	"net/url"
+	"os/exec"
+	"runtime"
+	"strings"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	"github.com/stretchr/testify/require"
+
+	"github.com/envoyproxy/ai-gateway/tests/internal/testmcp"
+)
+
+const (
+	writeTimeout  = 120 * time.Second
+	mcpServerPort = 8080
+	aigwPort      = 1975
+)
+
+var aigwBinary = fmt.Sprintf("../../out/aigw-%s-%s", runtime.GOOS, runtime.GOARCH)
+
+type MCPBenchCase struct {
+	Name        string
+	ProxyBinary string
+	ProxyArgs   []string
+	TestAddr    string
+}
+
+// setupBenchmark sets up the client connection.
+func setupBenchmark(b *testing.B) []MCPBenchCase {
+	b.Helper() // Treat this as a helper function
+
+	// setup MCP server
+	mcpServer := testmcp.NewServer(&testmcp.Options{
+		Port:              mcpServerPort,
+		ForceJSONResponse: false,
+		DumbEchoServer:    true,
+		WriteTimeout:      writeTimeout,
+		DisableLog:        true,
+	})
+	b.Cleanup(func() {
+		_ = mcpServer.Close()
+	})
+
+	return []MCPBenchCase{
+		{
+			Name:     "Baseline_NoProxy",
+			TestAddr: fmt.Sprintf("http://localhost:%d", mcpServerPort),
+		},
+		{
+			Name:        "EAIGW_Default",
+			TestAddr:    fmt.Sprintf("http://localhost:%d/mcp", aigwPort),
+			ProxyBinary: aigwBinary,
+			ProxyArgs:   []string{"run", "./aigw.yaml"},
+		},
+		{
+			Name:        "EAIGW_Config_100",
+			TestAddr:    fmt.Sprintf("http://localhost:%d/mcp", aigwPort),
+			ProxyBinary: aigwBinary,
+			ProxyArgs:   []string{"run", "./aigw.yaml", "--mcp-session-encryption-iterations=100"},
+		},
+		{
+			Name:        "EAIGW_Inline_100",
+			TestAddr:    fmt.Sprintf("http://localhost:%d/mcp", aigwPort),
+			ProxyBinary: aigwBinary,
+			ProxyArgs: []string{
+				"run",
+				"--mcp-session-encryption-iterations=100",
+				`--mcp-json={"mcpServers":{"aigw":{"type":"http","url":"http://localhost:8080/mcp"}}}`,
+			},
+		},
+	}
+}
+
+func BenchmarkMCP(b *testing.B) {
+	cases := setupBenchmark(b)
+	for _, tc := range cases {
+		var proxy *exec.Cmd
+		if tc.ProxyBinary != "" {
+			proxy = startProxy(b, &tc)
+		}
+
+		b.Run(tc.Name, func(b *testing.B) {
+			mcpClient := mcp.NewClient(&mcp.Implementation{Name: "bench-http-client", Version: "0.1.0"}, nil)
+			cs, err := mcpClient.Connect(b.Context(), &mcp.StreamableClientTransport{Endpoint: tc.TestAddr}, nil)
+			if err != nil {
+				b.Fatalf("Failed to connect server: %v", err)
+			}
+
+			tools, err := cs.ListTools(b.Context(), &mcp.ListToolsParams{})
+			if err != nil {
+				b.Fatalf("Failed to list tools: %v", err)
+			}
+			var toolName string
+			for _, t := range tools.Tools {
+				if strings.Contains(t.Name, "echo") {
+					toolName = t.Name
+					break
+				}
+			}
+			if toolName == "" {
+				b.Fatalf("no echo tool found")
+			}
+
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				ctx, cancel := context.WithTimeout(b.Context(), 5*time.Second)
+				res, err := cs.CallTool(ctx, &mcp.CallToolParams{
+					Name:      toolName,
+					Arguments: testmcp.ToolEchoArgs{Text: "hello MCP"},
+				})
+				cancel()
+				if err != nil {
+					b.Fatalf("MCP Tool call name %s failed at iteration %d: %v", toolName, i, err)
+				}
+
+				txt, ok := res.Content[0].(*mcp.TextContent)
+				if !ok {
+					b.Fatalf("unexpected content type")
+				}
+				if txt.Text != "dumb echo: hello MCP" {
+					b.Fatalf("unexpected text: %q", txt.Text)
+				}
+			}
+		})
+
+		if proxy != nil && proxy.Process != nil {
+			_ = syscall.Kill(-proxy.Process.Pid, syscall.SIGKILL)
+		}
+	}
+}
+
+func startProxy(b testing.TB, tc *MCPBenchCase) *exec.Cmd {
+	addr, err := url.Parse(tc.TestAddr)
+	require.NoError(b, err)
+
+	cmd := exec.CommandContext(b.Context(), tc.ProxyBinary, tc.ProxyArgs...) // nolint: gosec
+	// put into new process group so we can kill the entire process tree (and children)
+	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+	require.NoError(b, cmd.Start())
+
+	// Wait until we can connect to the proxy
+	require.Eventually(b, func() bool {
+		_, err = (&net.Dialer{}).DialContext(b.Context(), "tcp", addr.Host)
+		return err == nil
+	}, 30*time.Second, 500*time.Millisecond, "proxy %s did not become ready in time", tc.Name)
+
+	return cmd
+}
diff --git a/tests/internal/testmcp/server.go b/tests/internal/testmcp/server.go
@@ -26,6 +26,7 @@ type Options struct {
 	Port                              int
 	ForceJSONResponse, DumbEchoServer bool
 	WriteTimeout                      time.Duration
+	DisableLog                        bool
 }
 
 // NewServer starts a demo MCP server with two tools: echo and sum.
@@ -130,6 +131,9 @@ func NewServer(opts *Options) *http.Server {
 		WriteTimeout: opts.WriteTimeout,
 		Handler:      handler,
 		ConnState: func(conn net.Conn, state http.ConnState) {
+			if opts.DisableLog {
+				return
+			}
 			log.Printf("MCP SERVER connection [%s] %s -> %s\n", state, conn.RemoteAddr(), conn.LocalAddr())
 		},
 	}
