diff --git a/docs/middleware.md b/docs/middleware.md index 93b4b6201b..c5253177d9 100644 --- a/docs/middleware.md +++ b/docs/middleware.md @@ -63,7 +63,7 @@ Example config files: ### Known limitations - **Stale Modern headers.** A mutating webhook patches only the JSON body, so after it renames a tool the `Mcp-Method`/`Mcp-Name` headers forwarded to the backend still describe the original name. `ValidateHeaderConsistency` (`pkg/mcp/revision.go:512`) requires the header and body to agree and returns a `RequestHeaderMismatchError` (`CodeHeaderMismatch`) if they don't, but today that check is only wired into vMCP (`pkg/vmcp/server/classification.go:72`), which never runs the mutating webhook — so nothing in ToolHive catches this yet. A spec-conformant Modern (2026-07-28) backend will reject the mismatched request itself, so this fails closed rather than silently mis-authorizing. Practical guidance: a mutating webhook must not rename tools on the Modern path. -- **Controls outside the parser see the pre-mutation request.** The tool-call filter (`pkg/mcp/tool_filter.go`) reads the raw request body directly, and the rate limiter runs before the mutating webhook in the chain (see the ordering rules below) — both decide against the request as received, not as mutated. A mutating webhook can therefore rename a call into a tool that `--tools` filtering excluded, and the rate limiter debits the bucket for the requested tool rather than the executed one. This is a known gap, not a regression: the tool filter's position ahead of the MCP parser is deliberate (it needs the raw request), as already noted in the ordering rules below. +- **Controls evaluated before mutation see the pre-mutation request.** The tool-call filter (`pkg/mcp/tool_filter.go`) reads the raw request body before parsing, and the rate limiter runs after parsing but before the mutating webhook (see the ordering rules below). Both decide against the request as received, not as mutated. A mutating webhook can therefore rename a call into a tool that `--tools` filtering excluded, and the rate limiter debits the bucket for the requested tool rather than the executed one. This is a known gap, not a regression: filtering must happen before excluded calls proceed, and rate limiting intentionally rejects excess traffic before making webhook round trips. ## Architecture Diagram @@ -1055,11 +1055,11 @@ The middleware chain execution order is critical and controlled by the order in 5. **Tool Filter Middleware** (if enabled) - Filters available tools in list responses 6. **Tool Call Filter Middleware** (if enabled) - Filters tool call requests 7. **MCP Parser Middleware** (always present) - Parses JSON-RPC MCP requests -8. **Rate Limit Middleware** (if configured) - Enforces per-identity/tool limits using the parsed request -9. **Mutating Webhook Middleware** (if configured) - Patches the parsed MCP request and republishes the parse -10. **Validating Webhook Middleware** (if configured) - Approves or denies the (possibly mutated) request -11. **Usage Metrics Middleware** (if enabled) - Tracks tool call counts -12. **Telemetry Middleware** (if enabled) - OpenTelemetry instrumentation +8. **Telemetry Middleware** (if enabled) - Starts the request span; finalizes MCP identity after the inner chain returns +9. **Rate Limit Middleware** (if configured) - Enforces per-identity/tool limits using the parsed request +10. **Mutating Webhook Middleware** (if configured) - Patches the parsed MCP request and republishes the parse +11. **Validating Webhook Middleware** (if configured) - Approves or denies the (possibly mutated) request +12. **Usage Metrics Middleware** (if enabled) - Tracks tool call counts 13. **Authorization Middleware** (if enabled) - Cedar policy evaluation 14. **Header Forward Middleware** (if configured for remote servers) - Injects custom headers 15. **Recovery Middleware** (always present) - Catches panics @@ -1072,8 +1072,10 @@ The middleware chain execution order is critical and controlled by the order in - Token Exchange must come after Upstream Swap if both are used (can further transform the upstream IdP token) - Tool filters should come before MCP Parser to operate on raw requests - MCP Parser must come before Authorization (provides structured MCP data) +- Telemetry must come before Rate Limiting so the limiter can annotate the active request span; telemetry finalizes MCP names, attributes, and metrics after downstream processing from the shared parsed-request holder +- Rate Limiting must come before webhooks so rejected traffic does not invoke outbound webhook requests - Mutating webhooks must come before Validating webhooks and before Authorization, so policy evaluation sees the patched request -- Middleware that rewrites the request body must republish the parsed request via `mcp.RepublishParsedMCPRequest` and refresh `r.ContentLength` — `ParsingMiddleware` deliberately parses only once, so every later consumer (authorization, audit, telemetry, usage metrics) reads the cached parse rather than re-reading the body +- Middleware that rewrites the request body must republish the parsed request via `mcp.RepublishParsedMCPRequest` and refresh `r.ContentLength` — `ParsingMiddleware` deliberately parses only once, so downstream consumers read the refreshed context while outer audit and telemetry wrappers read the refreshed `mcp.ParsedRequestHolder` after the inner chain returns - Header Forward executes close to the backend handler (innermost position) - Recovery is always last in config, making it the innermost wrapper (the chain wraps in reverse config order, so the first entry is the outermost and runs first) - Body-size limit and Origin validation stay OUTSIDE audit: oversized bodies must be rejected before audit buffers request data, and origin validation is a pre-auth DNS-rebind guard. Their rejections (413/403) are the only ones not audited. @@ -1115,4 +1117,4 @@ export LOG_LEVEL=debug thv run --transport sse --name my-server my-image:latest ``` -This will show detailed information about each middleware component's execution and data flow. \ No newline at end of file +This will show detailed information about each middleware component's execution and data flow. diff --git a/docs/observability.md b/docs/observability.md index 5dc1543cbc..839ea6a142 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -340,6 +340,33 @@ The `mcp.resource.uri` attribute is set only for the following methods: `resources/read`, `resources/subscribe`, `resources/unsubscribe`, `notifications/resources/updated`. +### Rate Limit Attributes + +Redis-backed rate limit checks annotate the existing request span; they do not +create a separate span. Normal allowed and rejected outcomes set all three +attributes below. + +| Attribute | Type | Description | +|-----------|------|-------------| +| `rate_limit.decision` | string | `"allowed"` or `"rejected"` | +| `rate_limit.rejected_by` | string | `"none"` for allowed requests, otherwise the bucket that rejected the request | +| `rate_limit.fail_open` | bool | `false` for normal allowed and rejected outcomes | + +The bounded `rate_limit.rejected_by` values are: + +| Value | Limiting bucket | +|-------|-----------------| +| `shared_server` | Server-wide shared limit | +| `shared_tool` | Tool-specific shared limit | +| `per_user_server` | Server-wide per-user limit | +| `per_user_tool` | Tool-specific per-user limit | + +When no configured bucket applies to a tool call, the span records +`rate_limit.decision="allowed"`, `rate_limit.rejected_by="none"`, and +`rate_limit.fail_open=false`. Redis check failures do not receive these normal +outcome attributes. If multiple rate limit checks use the same request span, +the latest normal outcome replaces earlier values. + ### Tool, Prompt, and Resource Attributes **For `tools/call`:** diff --git a/pkg/mcp/parser.go b/pkg/mcp/parser.go index 1dabb9a732..2a24290abd 100644 --- a/pkg/mcp/parser.go +++ b/pkg/mcp/parser.go @@ -153,10 +153,10 @@ func ParsingMiddleware(next http.Handler) http.Handler { // the caller is responsible for terminating the request (e.g. writing an // error response) instead of proceeding with a stale or absent parse. // -// This only refreshes consumers that read the parse from the request context or -// from a ParsedRequestHolder. Middleware that inspects the raw body from OUTSIDE -// ParsingMiddleware — the tool-call filter and the rate limiter — has already -// decided against the pre-rewrite body and is not corrected by republishing. +// This refreshes downstream context consumers and outer wrappers that inspect +// ParsedRequestHolder after the inner chain returns. It cannot change decisions +// already made before the rewrite: tool-call filtering and rate limiting still +// evaluate the request as received. // // The caller must also refresh r.ContentLength when it replaces r.Body, or the // reverse proxy will reject the forwarded request. diff --git a/pkg/ratelimit/limiter.go b/pkg/ratelimit/limiter.go index 14fca1ce87..d4f6648ee4 100644 --- a/pkg/ratelimit/limiter.go +++ b/pkg/ratelimit/limiter.go @@ -142,13 +142,17 @@ type bucketSpec struct { refillPeriod time.Duration } -// limitCheck keeps a bucket paired with its metric dimensions. +// limitCheck keeps a bucket paired with its observability dimensions. type limitCheck struct { bucket *bucket.TokenBucket scope string operationType string } +func (c limitCheck) rejectionIdentifier() string { + return c.scope + "_" + c.operationType +} + // limiter is the concrete implementation of Limiter. type limiter struct { client redis.Cmdable @@ -221,6 +225,7 @@ func (l *limiter) Allow(ctx context.Context, toolName, userID string) (*Decision } if len(checks) == 0 { + recordRateLimitSpanOutcome(ctx, rateLimitDecisionAllowed, rateLimitRejectedByNone) return &Decision{Allowed: true}, nil } @@ -238,6 +243,7 @@ func (l *limiter) Allow(ctx context.Context, toolName, userID string) (*Decision } if rejectedIdx >= 0 { l.telemetry.recordRejected(ctx, checks[rejectedIdx]) + recordRateLimitSpanOutcome(ctx, rateLimitDecisionRejected, checks[rejectedIdx].rejectionIdentifier()) return &Decision{ Allowed: false, RetryAfter: buckets[rejectedIdx].RetryAfter(), @@ -245,6 +251,7 @@ func (l *limiter) Allow(ctx context.Context, toolName, userID string) (*Decision } l.telemetry.recordAllowed(ctx, checks) + recordRateLimitSpanOutcome(ctx, rateLimitDecisionAllowed, rateLimitRejectedByNone) return &Decision{Allowed: true}, nil } diff --git a/pkg/ratelimit/middleware_test.go b/pkg/ratelimit/middleware_test.go index c0727290f6..f4e17d7b33 100644 --- a/pkg/ratelimit/middleware_test.go +++ b/pkg/ratelimit/middleware_test.go @@ -10,18 +10,21 @@ import ( "io" "net/http" "net/http/httptest" + "strings" "testing" "time" "github.com/alicebob/miniredis/v2" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" "go.uber.org/mock/gomock" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" v1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/telemetry" transporttypes "github.com/stacklok/toolhive/pkg/transport/types" transportmocks "github.com/stacklok/toolhive/pkg/transport/types/mocks" ) @@ -207,6 +210,56 @@ func TestRateLimitedBodyMarshalFallback(t *testing.T) { assert.NotContains(t, string(got), `"id"`) } +func TestRateLimitHandler_AnnotatesTelemetryRequestSpan(t *testing.T) { + t.Parallel() + client, _ := newTestClient(t) + limiter, err := newLimiter( + client, + "test-ns", + "test-server", + newSpanTestRateLimitConfig(t, rateLimitScopeShared, rateLimitOperationTool), + nil, + ) + require.NoError(t, err) + + decision, err := limiter.Allow(t.Context(), "search", "") + require.NoError(t, err) + require.True(t, decision.Allowed) + + tracerProvider, recorder := newRateLimitTracerProvider(t) + meterProvider := sdkmetric.NewMeterProvider() + t.Cleanup(func() { + require.NoError(t, meterProvider.Shutdown(context.Background())) + }) + telemetryMiddleware := telemetry.NewHTTPMiddleware( + telemetry.Config{}, + tracerProvider, + meterProvider, + "test-server", + "streamable-http", + ) + + handler := mcp.ParsingMiddleware(telemetryMiddleware(rateLimitHandler(limiter)(http.HandlerFunc( + func(http.ResponseWriter, *http.Request) { + t.Fatal("next handler should not be called when rate limited") + }, + )))) + req := httptest.NewRequest( + http.MethodPost, + "/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"search"}}`), + ) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + handler.ServeHTTP(w, req) + + assert.Equal(t, http.StatusTooManyRequests, w.Code) + spans := recorder.Ended() + require.Len(t, spans, 1) + requireRateLimitSpanAttributes(t, spans[0], "rejected", "shared_tool") +} + func TestRateLimitHandler_RedisErrorFailOpen(t *testing.T) { t.Parallel() diff --git a/pkg/ratelimit/observability.go b/pkg/ratelimit/observability.go index 447ed64215..7194602f9c 100644 --- a/pkg/ratelimit/observability.go +++ b/pkg/ratelimit/observability.go @@ -15,6 +15,7 @@ import ( "github.com/redis/go-redis/v9" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/metric" + "go.opentelemetry.io/otel/trace" "github.com/stacklok/toolhive/pkg/telemetry" ) @@ -24,6 +25,7 @@ const ( rateLimitDecisionAllowed = "allowed" rateLimitDecisionRejected = "rejected" + rateLimitRejectedByNone = "none" rateLimitScopeShared = "shared" rateLimitScopePerUser = "per_user" @@ -139,6 +141,14 @@ func (t *rateLimitTelemetry) recordCheckLatency(ctx context.Context, duration ti )) } +func recordRateLimitSpanOutcome(ctx context.Context, decision, rejectedBy string) { + trace.SpanFromContext(ctx).SetAttributes( + attribute.String("rate_limit.decision", decision), + attribute.String("rate_limit.rejected_by", rejectedBy), + attribute.Bool("rate_limit.fail_open", false), + ) +} + func classifyRedisError(err error) string { if redis.IsAuthError(err) { return redisErrorTypeAuth diff --git a/pkg/ratelimit/observability_test.go b/pkg/ratelimit/observability_test.go index 8eb2ba58c8..4f340284ba 100644 --- a/pkg/ratelimit/observability_test.go +++ b/pkg/ratelimit/observability_test.go @@ -15,6 +15,8 @@ import ( "go.opentelemetry.io/otel/attribute" sdkmetric "go.opentelemetry.io/otel/sdk/metric" "go.opentelemetry.io/otel/sdk/metric/metricdata" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" v1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" @@ -212,6 +214,134 @@ func TestRateLimitMetrics_NilMeterProviderIsNoOp(t *testing.T) { assert.True(t, decision.Allowed) } +func TestRateLimitSpanAttributes_NormalOutcomes(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + scope string + operationType string + toolName string + userID string + rejectedBy string + }{ + { + name: "shared server", + scope: rateLimitScopeShared, + operationType: rateLimitOperationServer, + rejectedBy: "shared_server", + }, + { + name: "shared tool", + scope: rateLimitScopeShared, + operationType: rateLimitOperationTool, + toolName: "search", + rejectedBy: "shared_tool", + }, + { + name: "per-user server", + scope: rateLimitScopePerUser, + operationType: rateLimitOperationServer, + userID: "alice", + rejectedBy: "per_user_server", + }, + { + name: "per-user tool", + scope: rateLimitScopePerUser, + operationType: rateLimitOperationTool, + toolName: "search", + userID: "alice", + rejectedBy: "per_user_tool", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + client, _ := newTestClient(t) + limiter, err := newLimiter( + client, + "test-ns", + "test-server", + newSpanTestRateLimitConfig(t, tt.scope, tt.operationType), + nil, + ) + require.NoError(t, err) + + tracerProvider, recorder := newRateLimitTracerProvider(t) + tracer := tracerProvider.Tracer("rate-limit-test") + + allowedCtx, allowedSpan := tracer.Start(t.Context(), "request") + decision, err := limiter.Allow(allowedCtx, tt.toolName, tt.userID) + require.NoError(t, err) + require.True(t, decision.Allowed) + allowedSpan.End() + + rejectedCtx, rejectedSpan := tracer.Start(t.Context(), "request") + decision, err = limiter.Allow(rejectedCtx, tt.toolName, tt.userID) + require.NoError(t, err) + require.False(t, decision.Allowed) + rejectedSpan.End() + + spans := recorder.Ended() + require.Len(t, spans, 2, "the limiter must annotate ambient spans without creating another span") + requireRateLimitSpanAttributes(t, spans[0], "allowed", "none") + requireRateLimitSpanAttributes(t, spans[1], "rejected", tt.rejectedBy) + }) + } +} + +func TestRateLimitSpanAttributes_NoApplicableBucketIsAllowed(t *testing.T) { + t.Parallel() + client, _ := newTestClient(t) + limiter, err := newLimiter( + client, + "test-ns", + "test-server", + newSpanTestRateLimitConfig(t, rateLimitScopeShared, rateLimitOperationTool), + nil, + ) + require.NoError(t, err) + + tracerProvider, recorder := newRateLimitTracerProvider(t) + ctx, span := tracerProvider.Tracer("rate-limit-test").Start(t.Context(), "request") + decision, err := limiter.Allow(ctx, "other-tool", "") + require.NoError(t, err) + require.True(t, decision.Allowed) + span.End() + + spans := recorder.Ended() + require.Len(t, spans, 1) + requireRateLimitSpanAttributes(t, spans[0], "allowed", "none") +} + +func TestRateLimitSpanAttributes_RedisErrorLeavesOutcomeUnset(t *testing.T) { + t.Parallel() + client, redisServer := newTestClient(t) + limiter, err := newLimiter( + client, + "test-ns", + "test-server", + newSpanTestRateLimitConfig(t, rateLimitScopeShared, rateLimitOperationServer), + nil, + ) + require.NoError(t, err) + redisServer.Close() + + tracerProvider, recorder := newRateLimitTracerProvider(t) + ctx, span := tracerProvider.Tracer("rate-limit-test").Start(t.Context(), "request") + _, err = limiter.Allow(ctx, "", "") + require.Error(t, err) + span.End() + + spans := recorder.Ended() + require.Len(t, spans, 1) + attributes := spanAttributeMap(spans[0]) + assert.NotContains(t, attributes, "rate_limit.decision") + assert.NotContains(t, attributes, "rate_limit.rejected_by") + assert.NotContains(t, attributes, "rate_limit.fail_open") +} + func TestClassifyRedisError(t *testing.T) { t.Parallel() @@ -277,6 +407,66 @@ func newRateLimitMeterProvider() (*sdkmetric.ManualReader, *sdkmetric.MeterProvi return reader, sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) } +func newRateLimitTracerProvider(t *testing.T) (*sdktrace.TracerProvider, *tracetest.SpanRecorder) { + t.Helper() + recorder := tracetest.NewSpanRecorder() + provider := sdktrace.NewTracerProvider( + sdktrace.WithSpanProcessor(recorder), + sdktrace.WithSampler(sdktrace.AlwaysSample()), + ) + t.Cleanup(func() { + require.NoError(t, provider.Shutdown(context.Background())) + }) + return provider, recorder +} + +func newSpanTestRateLimitConfig(t *testing.T, scope, operationType string) *v1beta1.RateLimitConfig { + t.Helper() + bucket := &v1beta1.RateLimitBucket{ + MaxTokens: 1, + RefillPeriod: metav1.Duration{Duration: time.Minute}, + } + + switch { + case scope == rateLimitScopeShared && operationType == rateLimitOperationServer: + return &v1beta1.RateLimitConfig{Shared: bucket} + case scope == rateLimitScopeShared && operationType == rateLimitOperationTool: + return &v1beta1.RateLimitConfig{ + Tools: []v1beta1.ToolRateLimitConfig{{Name: "search", Shared: bucket}}, + } + case scope == rateLimitScopePerUser && operationType == rateLimitOperationServer: + return &v1beta1.RateLimitConfig{PerUser: bucket} + case scope == rateLimitScopePerUser && operationType == rateLimitOperationTool: + return &v1beta1.RateLimitConfig{ + Tools: []v1beta1.ToolRateLimitConfig{{Name: "search", PerUser: bucket}}, + } + default: + t.Fatalf("unsupported rate limit span test dimensions: %s/%s", scope, operationType) + return nil + } +} + +func requireRateLimitSpanAttributes( + t *testing.T, + span sdktrace.ReadOnlySpan, + decision string, + rejectedBy string, +) { + t.Helper() + attributes := spanAttributeMap(span) + assert.Equal(t, decision, attributes["rate_limit.decision"]) + assert.Equal(t, rejectedBy, attributes["rate_limit.rejected_by"]) + assert.Equal(t, false, attributes["rate_limit.fail_open"]) +} + +func spanAttributeMap(span sdktrace.ReadOnlySpan) map[string]any { + attributes := make(map[string]any, len(span.Attributes())) + for _, attr := range span.Attributes() { + attributes[string(attr.Key)] = attr.Value.AsInterface() + } + return attributes +} + func collectRateLimitMetrics(t *testing.T, reader *sdkmetric.ManualReader) metricdata.ResourceMetrics { t.Helper() var metrics metricdata.ResourceMetrics diff --git a/pkg/runner/middleware.go b/pkg/runner/middleware.go index da116f1b18..d09302d24d 100644 --- a/pkg/runner/middleware.go +++ b/pkg/runner/middleware.go @@ -163,17 +163,36 @@ func PopulateMiddlewareConfigs(config *RunConfig) error { } middlewareConfigs = append(middlewareConfigs, *mcpParserConfig) + // Telemetry middleware (if enabled). + // Positioned after MCP parsing and before rate limiting so the limiter can + // annotate an active request span. MCP identity is finalized after the inner + // chain returns, using the parse republished by any mutating webhook. + if config.TelemetryConfig != nil { + telemetryParams := telemetry.FactoryMiddlewareParams{ + Config: config.TelemetryConfig, + ServerName: config.Name, + Transport: config.Transport.String(), + } + telemetryConfig, err := types.NewMiddlewareConfig(telemetry.MiddlewareType, telemetryParams) + if err != nil { + return fmt.Errorf("failed to create telemetry middleware config: %w", err) + } + middlewareConfigs = append(middlewareConfigs, *telemetryConfig) + } + // Rate limit middleware (if configured) - // Positioned after MCP parser (needs tool name from context). - // Will also need user identity from auth when per-user limits are added (#4550). + // Positioned after MCP parser (needs tool name from context) and telemetry + // (needs the active request span for rate-limit attributes). middlewareConfigs, err = addRateLimitMiddleware(middlewareConfigs, config) if err != nil { return err } // Mutating Webhooks middleware (if configured). - // Must run BEFORE validating webhooks: - // Audit -> ... -> MCP Parser -> [Mutating Webhooks] -> [Validating Webhooks] -> Authz + // Must run before validating webhooks and authorization. Parsing, telemetry, + // and rate limiting remain outside both webhook layers; republishing through + // the shared parsed-request holder lets telemetry observe the final request + // after the inner chain returns. middlewareConfigs, err = addMutatingWebhookMiddleware(middlewareConfigs, config) if err != nil { return err @@ -195,20 +214,6 @@ func PopulateMiddlewareConfigs(config *RunConfig) error { return err } - // Telemetry middleware (if enabled) - if config.TelemetryConfig != nil { - telemetryParams := telemetry.FactoryMiddlewareParams{ - Config: config.TelemetryConfig, - ServerName: config.Name, - Transport: config.Transport.String(), - } - telemetryConfig, err := types.NewMiddlewareConfig(telemetry.MiddlewareType, telemetryParams) - if err != nil { - return fmt.Errorf("failed to create telemetry middleware config: %w", err) - } - middlewareConfigs = append(middlewareConfigs, *telemetryConfig) - } - // Authorization middleware (if enabled) if config.AuthzConfig != nil { authzCfgData, err := injectUpstreamProviderIfNeeded(config.AuthzConfig, config.EmbeddedAuthServerConfig) diff --git a/pkg/runner/middleware_test.go b/pkg/runner/middleware_test.go index 2a29e7c4d6..2fa1bce7be 100644 --- a/pkg/runner/middleware_test.go +++ b/pkg/runner/middleware_test.go @@ -1317,6 +1317,50 @@ func TestPopulateMiddlewareConfigs_RateLimit(t *testing.T) { } } +func TestPopulateMiddlewareConfigs_RateLimitTelemetryOrdering(t *testing.T) { + t.Parallel() + config := &RunConfig{ + Name: "test-server", + TelemetryConfig: &telemetry.Config{}, + MutatingWebhooks: []webhook.Config{{ + Name: "mutating-hook", + URL: "http://example.com/mutate", + Timeout: webhook.DefaultTimeout, + }}, + ValidatingWebhooks: []webhook.Config{{ + Name: "validating-hook", + URL: "http://example.com/validate", + Timeout: webhook.DefaultTimeout, + }}, + RateLimitNamespace: "default", + RateLimitConfig: &v1beta1.RateLimitConfig{ + Shared: &v1beta1.RateLimitBucket{ + MaxTokens: 5, + RefillPeriod: metav1.Duration{Duration: time.Minute}, + }, + }, + ScalingConfig: &ScalingConfig{ + SessionRedis: &SessionRedisConfig{Address: "redis:6379"}, + }, + } + + require.NoError(t, PopulateMiddlewareConfigs(config)) + + parserIdx := indexOfMiddleware(t, config.MiddlewareConfigs, mcp.ParserMiddlewareType) + telemetryIdx := indexOfMiddleware(t, config.MiddlewareConfigs, telemetry.MiddlewareType) + rateLimitIdx := indexOfMiddleware(t, config.MiddlewareConfigs, ratelimit.MiddlewareType) + mutatingIdx := indexOfMiddleware(t, config.MiddlewareConfigs, mutating.MiddlewareType) + validatingIdx := indexOfMiddleware(t, config.MiddlewareConfigs, validating.MiddlewareType) + assert.Less(t, parserIdx, telemetryIdx, + "MCP parsing must run before telemetry creates the request span") + assert.Less(t, telemetryIdx, rateLimitIdx, + "telemetry must create the request span before rate limiting annotates it") + assert.Less(t, rateLimitIdx, mutatingIdx, + "rate limiting must reject excess traffic before invoking mutating webhooks") + assert.Less(t, mutatingIdx, validatingIdx, + "mutating webhooks must republish the final request before validation") +} + func TestPopulateMiddlewareConfigs_FullCoverage(t *testing.T) { t.Parallel() diff --git a/pkg/telemetry/integration_test.go b/pkg/telemetry/integration_test.go index 269ed2417d..4e60caad87 100644 --- a/pkg/telemetry/integration_test.go +++ b/pkg/telemetry/integration_test.go @@ -478,6 +478,115 @@ func TestTelemetryIntegration_ToolSpecificMetrics(t *testing.T) { assert.NoError(t, err) } +func TestTelemetryIntegration_UsesRepublishedRequestForMCPIdentity(t *testing.T) { + t.Parallel() + + const ( + requestedTool = "requested_search" + executedTool = "executed_search" + ) + + for _, withExistingHolder := range []bool{false, true} { + name := "telemetry holder" + if withExistingHolder { + name = "outer audit holder" + } + t.Run(name, func(t *testing.T) { + t.Parallel() + + recorder := tracetest.NewSpanRecorder() + tracerProvider := sdktrace.NewTracerProvider( + sdktrace.WithSpanProcessor(recorder), + sdktrace.WithSampler(sdktrace.AlwaysSample()), + ) + t.Cleanup(func() { + require.NoError(t, tracerProvider.Shutdown(context.Background())) + }) + + metricsReader := sdkmetric.NewManualReader() + meterProvider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(metricsReader)) + t.Cleanup(func() { + require.NoError(t, meterProvider.Shutdown(context.Background())) + }) + + middleware := NewHTTPMiddleware( + Config{ServiceName: "test-service", ServiceVersion: "1.0.0"}, + tracerProvider, + meterProvider, + "github", + "streamable-http", + ) + + mutatedBody := []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"executed_search","arguments":{"query":"mutated"}}}`) + inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + republished, err := mcp.RepublishParsedMCPRequest(r, mutatedBody) + require.NoError(t, err) + require.Equal(t, executedTool, mcp.GetParsedMCPRequest(republished.Context()).ResourceID) + w.WriteHeader(http.StatusOK) + }) + handler := mcp.ParsingMiddleware(middleware(inner)) + + requestBody := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"requested_search","arguments":{"query":"original"}}}` + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(requestBody)) + req.Header.Set("Content-Type", "application/json") + + var outerHolder *mcp.ParsedRequestHolder + if withExistingHolder { + outerHolder = &mcp.ParsedRequestHolder{} + req = req.WithContext(mcp.WithParsedRequestHolder(req.Context(), outerHolder)) + } + + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + require.Equal(t, http.StatusOK, rec.Code) + + spans := recorder.Ended() + require.Len(t, spans, 1) + assert.Equal(t, "tools/call "+executedTool, spans[0].Name()) + + spanAttrs := make(map[string]any, len(spans[0].Attributes())) + for _, attr := range spans[0].Attributes() { + spanAttrs[string(attr.Key)] = attr.Value.AsInterface() + } + assert.Equal(t, executedTool, spanAttrs["gen_ai.tool.name"]) + assert.Contains(t, spanAttrs["gen_ai.tool.call.arguments"], "query=mutated") + assert.NotContains(t, spanAttrs["gen_ai.tool.call.arguments"], "query=original") + + var rm metricdata.ResourceMetrics + require.NoError(t, metricsReader.Collect(context.Background(), &rm)) + var requestResourceID, toolCallName string + for _, scopeMetrics := range rm.ScopeMetrics { + for _, metric := range scopeMetrics.Metrics { + sum, ok := metric.Data.(metricdata.Sum[int64]) + if !ok { + continue + } + for _, dataPoint := range sum.DataPoints { + for _, attr := range dataPoint.Attributes.ToSlice() { + switch { + case metric.Name == metricRequestCounter && attr.Key == "mcp_resource_id": + requestResourceID = attr.Value.AsString() + case metric.Name == "toolhive_mcp_tool_calls" && attr.Key == "tool": + toolCallName = attr.Value.AsString() + } + } + } + } + } + assert.Equal(t, executedTool, requestResourceID) + assert.Equal(t, executedTool, toolCallName) + assert.NotEqual(t, requestedTool, requestResourceID) + assert.NotEqual(t, requestedTool, toolCallName) + + if outerHolder != nil { + require.NotNil(t, outerHolder.Parsed) + assert.Equal(t, executedTool, outerHolder.Parsed.ResourceID, + "telemetry must reuse the holder installed by an outer audit middleware") + } + }) + } +} + func TestTelemetryIntegration_MultipleRequests(t *testing.T) { t.Parallel() diff --git a/pkg/telemetry/middleware.go b/pkg/telemetry/middleware.go index e155212556..7b086371db 100644 --- a/pkg/telemetry/middleware.go +++ b/pkg/telemetry/middleware.go @@ -165,6 +165,19 @@ func (m *HTTPMiddleware) Handler(next http.Handler) http.Handler { } // Normal HTTP request handling + // Reuse an outer holder (for example, audit's) so webhook mutations are + // visible to every wrapper. When there is no outer holder, install one for + // telemetry itself and seed it with the parse published by the parser. + parsedHolder, ok := mcpparser.ParsedRequestHolderFromContext(ctx) + if !ok { + parsedHolder = &mcpparser.ParsedRequestHolder{ + Parsed: mcpparser.GetParsedMCPRequest(ctx), + } + ctx = mcpparser.WithParsedRequestHolder(ctx, parsedHolder) + } else if parsedHolder.Parsed == nil { + parsedHolder.Parsed = mcpparser.GetParsedMCPRequest(ctx) + } + // Extract trace context from incoming request headers ctx = otel.GetTextMapPropagator().Extract(ctx, propagation.HeaderCarrier(r.Header)) @@ -187,12 +200,14 @@ func (m *HTTPMiddleware) Handler(next http.Handler) http.Handler { attribute.String("transport", m.transport), )) - // Create span name based on MCP method if available, otherwise use HTTP method + path - spanName := m.createSpanName(ctx) - if spanName == "" { - spanName = fmt.Sprintf("%s %s", r.Method, r.URL.Path) - } - ctx, span := m.tracer.Start(ctx, spanName, trace.WithSpanKind(trace.SpanKindServer)) + // Start the request span before rate limiting. MCP identity is populated + // after the inner chain returns because a mutating webhook may republish a + // different method, resource, or argument set while the request is in flight. + ctx, span := m.tracer.Start( + ctx, + fmt.Sprintf("%s %s", r.Method, r.URL.Path), + trace.WithSpanKind(trace.SpanKindServer), + ) defer span.End() // Create a response writer wrapper to capture response details @@ -205,9 +220,6 @@ func (m *HTTPMiddleware) Handler(next http.Handler) http.Handler { // Add HTTP attributes m.addHTTPAttributes(span, r) - // Add MCP attributes if parsed data is available - m.addMCPAttributes(ctx, span, r) - // Add environment variables as attributes m.addEnvironmentAttributes(span) @@ -217,13 +229,31 @@ func (m *HTTPMiddleware) Handler(next http.Handler) http.Handler { // Call the next handler with the instrumented context next.ServeHTTP(rw, r.WithContext(ctx)) - // Record completion metrics and finalize span + // A mutating webhook publishes its replacement parse through the shared + // holder. Use that final parse for all MCP identity exported by telemetry. + finalCtx := contextWithFinalParsedMCPRequest(ctx, parsedHolder) + if spanName := m.createSpanName(finalCtx); spanName != "" { + span.SetName(spanName) + } + m.addMCPAttributes(finalCtx, span, r) + + // Record completion metrics and finalize span. duration := time.Since(startTime) m.finalizeSpan(span, rw, duration) - m.recordMetrics(ctx, r, rw, duration) + m.recordMetrics(finalCtx, r, rw, duration) }) } +func contextWithFinalParsedMCPRequest( + ctx context.Context, + holder *mcpparser.ParsedRequestHolder, +) context.Context { + if holder == nil || holder.Parsed == nil || holder.Parsed == mcpparser.GetParsedMCPRequest(ctx) { + return ctx + } + return context.WithValue(ctx, mcpparser.MCPRequestContextKey, holder.Parsed) +} + func (*HTTPMiddleware) createSpanName(ctx context.Context) string { parsedMCP := mcpparser.GetParsedMCPRequest(ctx) if parsedMCP == nil || parsedMCP.Method == "" { diff --git a/pkg/vmcp/ratelimit/decorator_test.go b/pkg/vmcp/ratelimit/decorator_test.go index fb9bbe1337..3a43ceb38c 100644 --- a/pkg/vmcp/ratelimit/decorator_test.go +++ b/pkg/vmcp/ratelimit/decorator_test.go @@ -9,9 +9,15 @@ import ( "testing" "time" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + v1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/pkg/auth" baseratelimit "github.com/stacklok/toolhive/pkg/ratelimit" "github.com/stacklok/toolhive/pkg/vmcp" @@ -135,6 +141,61 @@ func TestCallToolRateLimitedDoesNotDelegate(t *testing.T) { assert.Equal(t, 5*time.Second, limited.RetryAfter) } +func TestCallToolRateLimitedAnnotatesAmbientSpan(t *testing.T) { + t.Parallel() + redisServer := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: redisServer.Addr()}) + t.Cleanup(func() { + require.NoError(t, client.Close()) + }) + limiter, err := baseratelimit.NewLimiter(client, "test-ns", "test-vmcp", &v1beta1.RateLimitConfig{ + Tools: []v1beta1.ToolRateLimitConfig{ + { + Name: "backend_a_echo", + Shared: &v1beta1.RateLimitBucket{ + MaxTokens: 1, + RefillPeriod: metav1.Duration{Duration: time.Minute}, + }, + }, + }, + }) + require.NoError(t, err) + + decision, err := limiter.Allow(t.Context(), "backend_a_echo", "") + require.NoError(t, err) + require.True(t, decision.Allowed) + + recorder := tracetest.NewSpanRecorder() + tracerProvider := sdktrace.NewTracerProvider( + sdktrace.WithSpanProcessor(recorder), + sdktrace.WithSampler(sdktrace.AlwaysSample()), + ) + t.Cleanup(func() { + require.NoError(t, tracerProvider.Shutdown(context.Background())) + }) + ctx, span := tracerProvider.Tracer("vmcp-rate-limit-test").Start(t.Context(), "request") + inner := &recordingCore{} + decorated := NewDecorator(inner, limiter) + + result, err := decorated.CallTool(ctx, nil, "backend_a_echo", nil, nil) + span.End() + + require.Error(t, err) + assert.Nil(t, result) + assert.False(t, inner.called) + var limited *baseratelimit.RateLimitedError + require.ErrorAs(t, err, &limited) + spans := recorder.Ended() + require.Len(t, spans, 1, "the decorator must preserve the ambient span without creating another span") + attributes := make(map[string]any, len(spans[0].Attributes())) + for _, attr := range spans[0].Attributes() { + attributes[string(attr.Key)] = attr.Value.AsInterface() + } + assert.Equal(t, "rejected", attributes["rate_limit.decision"]) + assert.Equal(t, "shared_tool", attributes["rate_limit.rejected_by"]) + assert.Equal(t, false, attributes["rate_limit.fail_open"]) +} + func TestCallToolLimiterErrorFailsOpen(t *testing.T) { t.Parallel() diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go index 4e08a7592e..ba19076d34 100644 --- a/pkg/vmcp/server/server.go +++ b/pkg/vmcp/server/server.go @@ -560,8 +560,9 @@ func New( // This enables embedding the vmcp server inside another HTTP server or framework. // // The returned handler includes all routes (health, metrics, well-known, MCP) -// and the full middleware chain (recovery, body limit, header validation, -// audit, auth, MCP parsing, telemetry). +// and the full HTTP middleware chain (recovery, body limit, header validation, +// auth, audit, MCP parsing, telemetry). Rate limiting decorates the core VMCP +// and runs inside the MCP SDK handler. // // Each call builds a fresh handler. The method is safe to call multiple times. // All returned handlers share the same underlying MCPServer and SessionManager, @@ -624,6 +625,8 @@ func (s *Server) Handler(_ context.Context) (http.Handler, error) { // Code wraps: audit → auth → MCP-parsing → telemetry → classification // Execution order: recovery → body-limit → header-val → audit → auth → // MCP-parsing → telemetry → classification → handler + // Rate limiting is a core VMCP decorator rather than an HTTP middleware and + // runs during tool dispatch inside the handler. // // Upstream token refresh failures are detected inside AuthMiddleware itself: // GetAllUpstreamCredentials returns a non-empty failed-provider slice when diff --git a/pkg/webhook/mutating/middleware.go b/pkg/webhook/mutating/middleware.go index c1c3d18535..bffbd100f4 100644 --- a/pkg/webhook/mutating/middleware.go +++ b/pkg/webhook/mutating/middleware.go @@ -130,8 +130,9 @@ func createMutatingHandler(executors []clientExecutor, serverName, transport str } // A mutating webhook rewrote the body, so the parse cached by ParsingMiddleware - // now describes a request the backend will not execute. Refresh it before any - // downstream consumer (authorization, audit, telemetry) reads it. + // now describes a request the backend will not execute. Refresh it for downstream + // consumers and publish it through the holder so outer audit and telemetry + // wrappers observe the final request after the inner chain returns. if !bytes.Equal(bodyBytes, mutatedBody) { republished, err := mcp.RepublishParsedMCPRequest(r, mutatedBody) if err != nil {