diff --git a/internal/mcpproxy/mcpproxy.go b/internal/mcpproxy/mcpproxy.go index 473db52b5d..52f04940da 100644 --- a/internal/mcpproxy/mcpproxy.go +++ b/internal/mcpproxy/mcpproxy.go @@ -238,8 +238,7 @@ func (m *mcpRequestContext) newSession(ctx context.Context, p *mcp.InitializePar if initErr != nil { m.l.Error("failed to create MCP session", slog.String("backend", backend.Name), slog.String("error", initErr.Error())) // If one backend fails, don't fail the overall connection. Create a session to the rest of the backends, as they - // may provide the needed methods. - // TODO: should we record a metric for this? + // may provide the needed methods. initializeSession has recorded the failure against the backend. return } m.metrics.WithBackend(backend.Name).RecordInitializationDuration(ctx, backendStartAt, p) @@ -483,7 +482,20 @@ type initializeResult struct { result *mcp.InitializeResult } -func (m *mcpRequestContext) initializeSession(ctx context.Context, routeName filterapi.MCPRouteName, backend filterapi.MCPBackend, p *mcp.InitializeParams, startAt time.Time) (*initializeResult, error) { +func (m *mcpRequestContext) initializeSession(ctx context.Context, routeName filterapi.MCPRouteName, backend filterapi.MCPBackend, p *mcp.InitializeParams, startAt time.Time) (result *initializeResult, err error) { + // Record a failure against the backend and the phase it happened in, the same way a failed + // request is recorded, so a backend that answers initialize but rejects notifications/initialized + // is distinguishable from one that never initialized. The duration is measured from the + // request's startAt, not from the start of this backend's initialization. + backendMetrics := m.metrics.WithBackend(backend.Name) + phase := "initialize" + defer func() { + if err != nil { + backendMetrics.RecordMethodErrorCount(ctx, phase, p, metrics.MCPStatusError) + backendMetrics.RecordRequestErrorDuration(ctx, startAt, errorType(err), p) + } + }() + // Send the initialize request to the MCP backend listener. reqID := mustJSONRPCRequestID() var ( @@ -588,13 +600,13 @@ func (m *mcpRequestContext) initializeSession(ctx context.Context, routeName fil if m.l.Enabled(ctx, slog.LevelDebug) { m.l.Debug("MCP session initialized", slog.Any("capabilities", initResult.Capabilities)) } - backendMetrics := m.metrics.WithBackend(backend.Name) backendMetrics.RecordServerCapabilities(ctx, initResult.Capabilities, p) backendMetrics.RecordMethodCount(ctx, "initialize", p) backendMetrics.RecordRequestDuration(ctx, startAt, p) } // Need to invoke "notifications/initialized" to complete the initialization. + phase = "notifications/initialized" { // Send the notifications/initialized request to the MCP backend listener. mcpReq := &jsonrpc.Request{Method: "notifications/initialized", Params: emptyJSONRPCMessage} @@ -611,7 +623,7 @@ func (m *mcpRequestContext) initializeSession(ctx context.Context, routeName fil body, _ := io.ReadAll(resp.Body) return nil, fmt.Errorf("MCP notifications/initialized request failed with status code %d, body=%s", resp.StatusCode, string(body)) } - m.metrics.WithBackend(backend.Name).RecordMethodCount(ctx, "notifications/initialized", p) + backendMetrics.RecordMethodCount(ctx, "notifications/initialized", p) } if m.l.Enabled(ctx, slog.LevelDebug) { m.l.Debug("sent MCP notifications/initialized", slog.String("backend", backend.Name), slog.String("session_id", sessionID)) diff --git a/internal/mcpproxy/mcpproxy_test.go b/internal/mcpproxy/mcpproxy_test.go index a1cbee864d..034967e910 100644 --- a/internal/mcpproxy/mcpproxy_test.go +++ b/internal/mcpproxy/mcpproxy_test.go @@ -19,9 +19,13 @@ import ( "github.com/modelcontextprotocol/go-sdk/jsonrpc" "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/attribute" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" "github.com/envoyproxy/ai-gateway/internal/filterapi" "github.com/envoyproxy/ai-gateway/internal/internalapi" + "github.com/envoyproxy/ai-gateway/internal/metrics" + "github.com/envoyproxy/ai-gateway/internal/testing/testotel" "github.com/envoyproxy/ai-gateway/internal/tracing/tracingapi" ) @@ -312,6 +316,65 @@ func TestNewSession_Success(t *testing.T) { require.NotEmpty(t, s.clientGatewaySessionID()) } +// TestNewSession_PartialBackendFailure_RecordsMetrics covers a composite session where one backend +// fails to initialize and the other succeeds: the session must stay usable through the surviving +// backend, and the failed backend must be visible in the backend-scoped metrics. +func TestNewSession_PartialBackendFailure_RecordsMetrics(t *testing.T) { + var callCount perBackendCallCount + backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + backend := r.Header.Get(internalapi.MCPBackendHeader) + if backend == "backend2" { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte("initialization failed")) + return + } + if callCount.inc(backend)%2 == 1 { + w.Header().Set(sessionIDHeader, "test-session-123") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(validInitializeResponse)) + } else { + w.WriteHeader(http.StatusAccepted) + } + })) + defer backendServer.Close() + + mr := sdkmetric.NewManualReader() + proxy := newTestMCPProxyWithOTEL(mr, noopTracer) + proxy.backendListenerAddr = backendServer.URL + + s, err := proxy.newSession(t.Context(), &mcp.InitializeParams{}, "test-route", "", nil, time.Now()) + require.NoError(t, err) + require.NotNil(t, s) + require.Contains(t, s.perBackendSessions, filterapi.MCPBackendName("backend1")) + require.NotContains(t, s.perBackendSessions, filterapi.MCPBackendName("backend2")) + + // The failed backend is counted as a failed initialize and its failure duration is recorded, + // the same way a failed request is. + failedInitialize := testotel.GetCounterValue(t, mr, "mcp.method.count", attribute.NewSet( + attribute.String("mcp.backend", "backend2"), + attribute.String("mcp.method.name", "initialize"), + attribute.String("status", string(metrics.MCPStatusError)), + )) + require.Equal(t, float64(1), failedInitialize) + failedCount, _ := testotel.GetHistogramValues(t, mr, "mcp.request.duration", attribute.NewSet( + attribute.String("mcp.backend", "backend2"), + attribute.String("error.type", string(metrics.MCPErrorInternal)), + )) + require.Equal(t, uint64(1), failedCount) + + // The surviving backend keeps its success metrics only. + okCount, _ := testotel.GetHistogramValues(t, mr, "mcp.initialization.duration", attribute.NewSet( + attribute.String("mcp.backend", "backend1"), + )) + require.Equal(t, uint64(1), okCount) + okInitialize := testotel.GetCounterValue(t, mr, "mcp.method.count", attribute.NewSet( + attribute.String("mcp.backend", "backend1"), + attribute.String("mcp.method.name", "initialize"), + attribute.String("status", string(metrics.MCPStatusSuccess)), + )) + require.Equal(t, float64(1), okInitialize) +} + func TestNewSession_NoBackend(t *testing.T) { proxy := newTestMCPProxy() @@ -489,7 +552,8 @@ func TestInitializeSession_NotificationsInitializedFailure(t *testing.T) { })) defer backendServer.Close() - proxy := newTestMCPProxy() + mr := sdkmetric.NewManualReader() + proxy := newTestMCPProxyWithOTEL(mr, noopTracer) proxy.backendListenerAddr = backendServer.URL sessionID, err := proxy.initializeSession(t.Context(), "route1", filterapi.MCPBackend{Name: "test-backend"}, &mcp.InitializeParams{}, time.Now()) @@ -497,6 +561,20 @@ func TestInitializeSession_NotificationsInitializedFailure(t *testing.T) { require.Error(t, err) require.Empty(t, sessionID) require.Contains(t, err.Error(), "notifications/initialized request failed") + + // The failure is attributed to notifications/initialized, not to initialize, which succeeded. + failedNotification := testotel.GetCounterValue(t, mr, "mcp.method.count", attribute.NewSet( + attribute.String("mcp.backend", "test-backend"), + attribute.String("mcp.method.name", "notifications/initialized"), + attribute.String("status", string(metrics.MCPStatusError)), + )) + require.Equal(t, float64(1), failedNotification) + okInitialize := testotel.GetCounterValue(t, mr, "mcp.method.count", attribute.NewSet( + attribute.String("mcp.backend", "test-backend"), + attribute.String("mcp.method.name", "initialize"), + attribute.String("status", string(metrics.MCPStatusSuccess)), + )) + require.Equal(t, float64(1), okInitialize) } func TestInvokeJSONRPCRequest_Success(t *testing.T) {