mcp: add configured header attribtues to metrics and spans (#1342)

**Description**

The LLM metrics and trace providers allow for the configuration of
custom headers to be included as metric attributes and span attributes.
The MCP implementation, however, did not honor those values.

This PR updates the MCP metrics and tracer configurations to also take
this configuration into account.
* The tracer method to create a span now receives the request headers.
* To avoid having to propagate request headers all the way down inside
the MCP Proxy so that they can be taken into account when recording
metrics, I've opted to instantiate an MCPProxy on each request, and have
it set up with the request attributes from the beginning. This allows
reusing the metrics implementation in a clean way.

**Related Issues/PRs (if applicable)**

Fixes #1335
Fixes #1348

**Special notes for reviewers (if applicable)**

N/A

---------

Signed-off-by: Ignasi Barrera <[email protected]>

Author: nacx

cmd/extproc/mainlib/main.go

@@ -233,7 +233,7 @@ func Main(ctx context.Context, args []string, stderr io.Writer) (err error) {
 	chatCompletionMetrics := metrics.NewChatCompletion(meter, metricsRequestHeaderAttributes)
 	completionMetrics := metrics.NewCompletion(meter, metricsRequestHeaderAttributes)
 	embeddingsMetrics := metrics.NewEmbeddings(meter, metricsRequestHeaderAttributes)
-	mcpMetrics := metrics.NewMCP(meter)
+	mcpMetrics := metrics.NewMCP(meter, metricsRequestHeaderAttributes)
 
 	tracing, err := tracing.NewTracingFromEnv(ctx, os.Stdout, spanRequestHeaderAttributes)
 	if err != nil {
@@ -264,13 +264,13 @@ func Main(ctx context.Context, args []string, stderr io.Writer) (err error) {
 		seed, fallbackSeed, _ := strings.Cut(flags.mcpSessionEncryptionSeed, ",")
 		mcpSessionCrypto := mcpproxy.DefaultSessionCrypto(seed, fallbackSeed)
 		var mcpProxyMux *http.ServeMux
-		var mcpProxy *mcpproxy.MCPProxy
-		mcpProxy, mcpProxyMux, err = mcpproxy.NewMCPProxy(l.With("component", "mcp-proxy"), mcpMetrics,
+		var mcpProxyConfig *mcpproxy.ProxyConfig
+		mcpProxyConfig, mcpProxyMux, err = mcpproxy.NewMCPProxy(l.With("component", "mcp-proxy"), mcpMetrics,
 			tracing.MCPTracer(), mcpSessionCrypto)
 		if err != nil {
 			return fmt.Errorf("failed to create MCP proxy: %w", err)
 		}
-		if err = extproc.StartConfigWatcher(ctx, flags.configPath, mcpProxy, l, time.Second*5); err != nil {
+		if err = extproc.StartConfigWatcher(ctx, flags.configPath, mcpProxyConfig, l, time.Second*5); err != nil {
 			return fmt.Errorf("failed to start config watcher: %w", err)
 		}
 

internal/lang/maps.go

@@ -0,0 +1,36 @@
+// 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.
+
+package lang
+
+import (
+	"fmt"
+	"maps"
+	"slices"
+	"strings"
+)
+
+// CaseInsensitiveValue retrieves a value from the meta map in a case-insensitive manner.
+// If the same key is present in different cases, the first one in alphabetical order
+// that matches is returned.
+// If the key is not found, it returns an empty string.
+func CaseInsensitiveValue(m map[string]any, key string) string {
+	if m == nil {
+		return ""
+	}
+
+	if v, ok := m[key]; ok {
+		return fmt.Sprintf("%v", v)
+	}
+
+	keys := slices.Sorted(maps.Keys(m))
+	for _, k := range keys {
+		if strings.EqualFold(k, key) {
+			return fmt.Sprintf("%v", m[k])
+		}
+	}
+
+	return ""
+}

internal/lang/maps_test.go

@@ -0,0 +1,57 @@
+// 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.
+
+package lang
+
+import "testing"
+
+func TestCaseInsensitiveValue(t *testing.T) {
+	tests := []struct {
+		name string
+		m    map[string]any
+		key  string
+		want string
+	}{
+		{
+			name: "nil map",
+			m:    nil,
+			key:  "anything",
+			want: "",
+		},
+		{
+			name: "exact match returns value",
+			m:    map[string]any{"Foo": "bar", "foo": "should-not-be-used"},
+			key:  "Foo",
+			want: "bar",
+		},
+		{
+			name: "case-insensitive match when exact not present",
+			m:    map[string]any{"FOO": "baz"},
+			key:  "foo",
+			want: "baz",
+		},
+		{
+			name: "multiple case variants - alphabetical first chosen",
+			m:    map[string]any{"ALPHA": 2, "Alpha": 1},
+			key:  "alpha",
+			want: "2", // ALPHA is alphabetically first
+		},
+		{
+			name: "nil value formatted",
+			m:    map[string]any{"key": nil},
+			key:  "key",
+			want: "<nil>",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got := CaseInsensitiveValue(tc.m, tc.key)
+			if got != tc.want {
+				t.Fatalf("CaseInsensitiveValue(%v, %q) = %q; want %q", tc.m, tc.key, got, tc.want)
+			}
+		})
+	}
+}

internal/mcpproxy/handlers.go

@@ -107,6 +107,7 @@ func (m *MCPProxy) servePOST(w http.ResponseWriter, r *http.Request) {
 		errType       metrics.MCPErrorType
 		requestMethod string
 		span          tracing.MCPSpan
+		params        mcp.Params
 	)
 	defer func() {
 		if m.l.Enabled(ctx, slog.LevelDebug) {
@@ -119,17 +120,17 @@ func (m *MCPProxy) servePOST(w http.ResponseWriter, r *http.Request) {
 			if span != nil {
 				span.EndSpanOnError(string(errType), err)
 			}
-			m.metrics.RecordMethodErrorCount(ctx)
-			m.metrics.RecordRequestErrorDuration(ctx, &startAt, errType)
+			m.metrics.RecordMethodErrorCount(ctx, params)
+			m.metrics.RecordRequestErrorDuration(ctx, &startAt, errType, params)
 			return
 		}
 
 		if span != nil {
 			span.EndSpan()
 		}
-		m.metrics.RecordRequestDuration(ctx, &startAt)
+		m.metrics.RecordRequestDuration(ctx, &startAt, params)
 		// TODO: should we special case when this request is "Response" where method is empty?
-		m.metrics.RecordMethodCount(ctx, requestMethod)
+		m.metrics.RecordMethodCount(ctx, requestMethod, params)
 	}()
 	if sessionID := r.Header.Get(sessionIDHeader); sessionID != "" {
 		s, err = m.sessionFromID(secureClientToGatewaySessionID(sessionID), secureClientToGatewayEventID(r.Header.Get(lastEventIDHeader)))
@@ -189,37 +190,37 @@ func (m *MCPProxy) servePOST(w http.ResponseWriter, r *http.Request) {
 
 		switch msg.Method {
 		case "notifications/roots/list_changed":
-			p := &mcp.RootsListChangedParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.RootsListChangedParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				onErrorResponse(w, http.StatusBadRequest, "invalid params")
 				return
 			}
 			err = m.handleNotificationsRootsListChanged(ctx, s, w, msg, span)
 		case "completion/complete":
-			p := &mcp.CompleteParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.CompleteParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				onErrorResponse(w, http.StatusBadRequest, "invalid params")
 				return
 			}
-			err = m.handleCompletionComplete(ctx, s, w, msg, p, span)
+			err = m.handleCompletionComplete(ctx, s, w, msg, params.(*mcp.CompleteParams), span)
 		case "notifications/progress":
-			m.metrics.RecordProgress(ctx)
-			p := &mcp.ProgressNotificationParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.ProgressNotificationParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
+			m.metrics.RecordProgress(ctx, params)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				onErrorResponse(w, http.StatusBadRequest, "invalid params")
 				return
 			}
-			err = m.handleClientToServerNotificationsProgress(ctx, s, w, msg, p, span)
+			err = m.handleClientToServerNotificationsProgress(ctx, s, w, msg, params.(*mcp.ProgressNotificationParams), span)
 		case "initialize":
 			// The very first request from the client to establish a session.
-			p := &mcp.InitializeParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.InitializeParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				m.l.Error("Failed to unmarshal initialize params", slog.String("error", err.Error()))
@@ -235,107 +236,107 @@ func (m *MCPProxy) servePOST(w http.ResponseWriter, r *http.Request) {
 				onErrorResponse(w, http.StatusInternalServerError, "missing route header")
 				return
 			}
-			err = m.handleInitializeRequest(ctx, w, msg, p, route, extractSubject(r), span)
+			err = m.handleInitializeRequest(ctx, w, msg, params.(*mcp.InitializeParams), route, extractSubject(r), span)
 		case "notifications/initialized":
 			// According to the MCP spec, when the server receives a JSON-RPC response or notification from the client
 			// and accepts it, the server MUST return HTTP 202 Accepted with an empty body.
 			// https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#sending-messages-to-the-server
 			w.WriteHeader(http.StatusAccepted)
 		case "logging/setLevel":
-			p := &mcp.SetLoggingLevelParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.SetLoggingLevelParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				m.l.Error("Failed to unmarshal set logging level params", slog.String("error", err.Error()))
 				onErrorResponse(w, http.StatusBadRequest, "invalid set logging level params")
 				return
 			}
-			err = m.handleSetLoggingLevel(ctx, s, w, msg, p, span)
+			err = m.handleSetLoggingLevel(ctx, s, w, msg, params.(*mcp.SetLoggingLevelParams), span)
 		case "ping":
 			// Ping is intentionally not traced as it's a lightweight health check.
 			err = m.handlePing(ctx, w, msg)
 		case "prompts/list":
-			p := &mcp.ListPromptsParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.ListPromptsParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				onErrorResponse(w, http.StatusBadRequest, "invalid params")
 				return
 			}
-			err = m.handlePromptListRequest(ctx, s, w, msg, p, span)
+			err = m.handlePromptListRequest(ctx, s, w, msg, params.(*mcp.ListPromptsParams), span)
 		case "prompts/get":
-			p := &mcp.GetPromptParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.GetPromptParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				onErrorResponse(w, http.StatusBadRequest, "invalid params")
 				return
 			}
-			err = m.handlePromptGetRequest(ctx, s, w, msg, p)
+			err = m.handlePromptGetRequest(ctx, s, w, msg, params.(*mcp.GetPromptParams))
 		case "tools/call":
-			p := &mcp.CallToolParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.CallToolParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				m.l.Error("Failed to unmarshal params", slog.String("method", msg.Method), slog.String("error", err.Error()))
 				onErrorResponse(w, http.StatusBadRequest, "invalid params")
 				return
 			}
-			err = m.handleToolCallRequest(ctx, s, w, msg, p, span)
+			err = m.handleToolCallRequest(ctx, s, w, msg, params.(*mcp.CallToolParams), span)
 		case "tools/list":
-			p := &mcp.ListToolsParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.ListToolsParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				onErrorResponse(w, http.StatusBadRequest, "invalid params")
 				return
 			}
-			err = m.handleToolsListRequest(ctx, s, w, msg, p, span)
+			err = m.handleToolsListRequest(ctx, s, w, msg, params.(*mcp.ListToolsParams), span)
 		case "resources/list":
-			p := &mcp.ListResourcesParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.ListResourcesParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				onErrorResponse(w, http.StatusBadRequest, "invalid params")
 				return
 			}
-			err = m.handleResourceListRequest(ctx, s, w, msg, p, span)
+			err = m.handleResourceListRequest(ctx, s, w, msg, params.(*mcp.ListResourcesParams), span)
 		case "resources/read":
-			p := &mcp.ReadResourceParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.ReadResourceParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				onErrorResponse(w, http.StatusBadRequest, "invalid params")
 				return
 			}
-			err = m.handleResourceReadRequest(ctx, s, w, msg, p)
+			err = m.handleResourceReadRequest(ctx, s, w, msg, params.(*mcp.ReadResourceParams))
 		case "resources/templates/list":
-			p := &mcp.ListResourceTemplatesParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.ListResourceTemplatesParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				onErrorResponse(w, http.StatusBadRequest, "invalid params")
 				return
 			}
-			err = m.handleResourcesTemplatesListRequest(ctx, s, w, msg, p, span)
+			err = m.handleResourcesTemplatesListRequest(ctx, s, w, msg, params.(*mcp.ListResourceTemplatesParams), span)
 		case "resources/subscribe":
-			p := &mcp.SubscribeParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.SubscribeParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				onErrorResponse(w, http.StatusBadRequest, "invalid params")
 				return
 			}
-			err = m.handleResourcesSubscribeRequest(ctx, s, w, msg, p, span)
+			err = m.handleResourcesSubscribeRequest(ctx, s, w, msg, params.(*mcp.SubscribeParams), span)
 		case "resources/unsubscribe":
-			p := &mcp.UnsubscribeParams{}
-			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, p)
+			params = &mcp.UnsubscribeParams{}
+			span, err = parseParamsAndMaybeStartSpan(ctx, m, msg, params, r.Header)
 			if err != nil {
 				errType = metrics.MCPErrorInvalidParam
 				onErrorResponse(w, http.StatusBadRequest, "invalid params")
 				return
 			}
-			err = m.handleResourcesUnsubscribeRequest(ctx, s, w, msg, p, span)
+			err = m.handleResourcesUnsubscribeRequest(ctx, s, w, msg, params.(*mcp.UnsubscribeParams), span)
 		case "notifications/cancelled":
 			// The responsibility of cancelling the operation on server side is optional, so we just ignore it for now.
 			// https://modelcontextprotocol.io/specification/2025-06-18/basic/utilities/cancellation#behavior-requirements
@@ -371,8 +372,7 @@ func errorType(err error) metrics.MCPErrorType {
 
 // handleInitializeRequest handles the "initialize" JSON-RPC method.
 func (m *MCPProxy) handleInitializeRequest(ctx context.Context, w http.ResponseWriter, req *jsonrpc.Request, p *mcp.InitializeParams, route, subject string, span tracing.MCPSpan) error {
-	m.metrics.RecordClientCapabilities(ctx, p.Capabilities)
-
+	m.metrics.RecordClientCapabilities(ctx, p.Capabilities, p)
 	s, err := m.newSession(ctx, p, route, subject, span)
 	if err != nil {
 		m.l.Error("failed to create new session", slog.String("error", err.Error()))
@@ -789,19 +789,23 @@ func (m *MCPProxy) recordResponse(ctx context.Context, rawMsg jsonrpc.Message) {
 		case "notifications/resources/list_changed":
 		case "notifications/resources/updated":
 		case "notifications/progress":
-			m.metrics.RecordProgress(ctx)
+			params := &mcp.ProgressNotificationParams{}
+			if err := json.Unmarshal(msg.Params, &params); err != nil {
+				m.l.Error("Failed to unmarshal params", slog.String("method", msg.Method), slog.String("error", err.Error()))
+			}
+			m.metrics.RecordProgress(ctx, params)
 		case "notifications/message":
 		case "notifications/tools/list_changed":
 		case "roots/list":
 		case "sampling/createMessage":
 		case "elicitation/create":
 		default:
 			knownMethod = false
-			m.metrics.RecordMethodErrorCount(ctx)
+			m.metrics.RecordMethodErrorCount(ctx, nil)
 			m.l.Warn("Unsupported MCP request method from server", slog.String("method", msg.Method))
 		}
 		if knownMethod {
-			m.metrics.RecordMethodCount(ctx, msg.Method)
+			m.metrics.RecordMethodCount(ctx, msg.Method, nil)
 		}
 	default:
 		m.l.Warn("unexpected message type in MCP response", slog.Any("message", msg))
@@ -1223,7 +1227,7 @@ func sendToAllBackendsAndAggregateResponsesImpl[responseType any](ctx context.Co
 }
 
 // parseParamsAndMaybeStartSpan parses the params from the JSON-RPC request and starts a tracing span if params is non-nil.
-func parseParamsAndMaybeStartSpan[paramType mcp.Params](ctx context.Context, m *MCPProxy, req *jsonrpc.Request, p paramType) (tracing.MCPSpan, error) {
+func parseParamsAndMaybeStartSpan[paramType mcp.Params](ctx context.Context, m *MCPProxy, req *jsonrpc.Request, p paramType, headers http.Header) (tracing.MCPSpan, error) {
 	if req.Params == nil {
 		return nil, nil
 	}
@@ -1233,7 +1237,7 @@ func parseParamsAndMaybeStartSpan[paramType mcp.Params](ctx context.Context, m *
 		return nil, err
 	}
 
-	span := m.tracer.StartSpanAndInjectMeta(ctx, req, p)
+	span := m.tracer.StartSpanAndInjectMeta(ctx, req, p, headers)
 	return span, nil
 }