Rate limiting observability (metrics and tracing) PR 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`:**

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,13 +243,15 @@ 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(),
 		}, nil
 	}
 
 	l.telemetry.recordAllowed(ctx, checks)
+	recordRateLimitSpanOutcome(ctx, rateLimitDecisionAllowed, rateLimitRejectedByNone)
 	return &Decision{Allowed: true}, nil
 }
 

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"
 )
@@ -115,6 +118,56 @@ func TestRateLimitHandler_ToolCallRejected(t *testing.T) {
 	assert.Equal(t, float64(42), resp["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()
 

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

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