diff --git a/pkg/ratelimit/limiter.go b/pkg/ratelimit/limiter.go index d9787ec4e5..f15bebc39b 100644 --- a/pkg/ratelimit/limiter.go +++ b/pkg/ratelimit/limiter.go @@ -159,6 +159,11 @@ type limitCheck struct { bucket *bucket.TokenBucket scope string operationType string + + // toolName is set only for tool-scoped checks. Cardinality is bounded by the + // configured per-tool buckets, so it is safe as a metric attribute; server-scoped + // checks leave it empty and the attribute is omitted rather than emitted blank. + toolName string } func (c limitCheck) rejectionIdentifier() string { @@ -186,6 +191,10 @@ func (l *limiter) recordFailOpen(ctx context.Context) { // Tokens are only consumed if ALL buckets have sufficient capacity, preventing // a rejected per-tool or per-user call from draining other budgets. func (l *limiter) Allow(ctx context.Context, toolName, userID string) (*Decision, error) { + // Attribution for the caller and tool goes on the span, not the metric: userID is + // unbounded, so it would multiply the decisions counter by the size of the user base. + recordRateLimitSpanAttribution(ctx, toolName, userID) + // Collect applicable buckets in priority order. var checks []limitCheck if l.serverBucket != nil { @@ -201,6 +210,7 @@ func (l *limiter) Allow(ctx context.Context, toolName, userID string) (*Decision bucket: tb, scope: rateLimitScopeShared, operationType: rateLimitOperationTool, + toolName: toolName, }) } } @@ -238,6 +248,7 @@ func (l *limiter) Allow(ctx context.Context, toolName, userID string) (*Decision ), scope: rateLimitScopePerUser, operationType: rateLimitOperationTool, + toolName: toolName, }) } } diff --git a/pkg/ratelimit/observability.go b/pkg/ratelimit/observability.go index 97f6f1c4da..b076ba2236 100644 --- a/pkg/ratelimit/observability.go +++ b/pkg/ratelimit/observability.go @@ -121,13 +121,19 @@ func (t *rateLimitTelemetry) recordRejected(ctx context.Context, check limitChec } func (t *rateLimitTelemetry) recordDecision(ctx context.Context, decision string, check limitCheck) { - t.decisions.Add(ctx, 1, metric.WithAttributes( + attrs := []attribute.KeyValue{ attribute.String("namespace", t.namespace), attribute.String("server", t.serverName), attribute.String("decision", decision), attribute.String("scope", check.scope), attribute.String("operation_type", check.operationType), - )) + } + // Only tool-scoped checks carry a tool name. Emitting it blank for server-scoped + // decisions would add a label to every series for no information. + if check.toolName != "" { + attrs = append(attrs, attribute.String("tool_name", check.toolName)) + } + t.decisions.Add(ctx, 1, metric.WithAttributes(attrs...)) } func (t *rateLimitTelemetry) recordRedisError(ctx context.Context, err error) { @@ -161,6 +167,25 @@ func (t *rateLimitTelemetry) recordCheckLatency(ctx context.Context, duration ti )) } +// recordRateLimitSpanAttribution records who the request belonged to and which tool it +// targeted. This lives on the span rather than on toolhive_rate_limit_decisions_total +// because user identity is unbounded: as a metric attribute it would create one time +// series per caller per bucket. On a span it costs nothing and answers "which user is +// being throttled", which the counter alone cannot. +func recordRateLimitSpanAttribution(ctx context.Context, toolName, userID string) { + attrs := make([]attribute.KeyValue, 0, 2) + if toolName != "" { + attrs = append(attrs, attribute.String("rate_limit.tool_name", toolName)) + } + if userID != "" { + attrs = append(attrs, attribute.String("rate_limit.user_id", userID)) + } + if len(attrs) == 0 { + return + } + trace.SpanFromContext(ctx).SetAttributes(attrs...) +} + func recordRateLimitSpanOutcome(ctx context.Context, decision, rejectedBy string, failOpen bool) { trace.SpanFromContext(ctx).SetAttributes( attribute.String("rate_limit.decision", decision), diff --git a/pkg/ratelimit/observability_test.go b/pkg/ratelimit/observability_test.go index 420fa99957..a028f96edd 100644 --- a/pkg/ratelimit/observability_test.go +++ b/pkg/ratelimit/observability_test.go @@ -709,3 +709,93 @@ func stringAttributeMap(attributes attribute.Set) map[string]string { } return result } + +func TestRateLimitMetrics_ToolAttribution(t *testing.T) { + t.Parallel() + client, _ := newTestClient(t) + reader, meterProvider := newRateLimitMeterProvider() + + limiter, err := newLimiter(client, "test-ns", "test-server", &v1beta1.RateLimitConfig{ + Shared: &v1beta1.RateLimitBucket{ + MaxTokens: 10, + RefillPeriod: metav1.Duration{Duration: time.Minute}, + }, + Tools: []v1beta1.ToolRateLimitConfig{ + { + Name: "search", + PerUser: &v1beta1.RateLimitBucket{ + MaxTokens: 1, + RefillPeriod: metav1.Duration{Duration: time.Minute}, + }, + }, + }, + }, meterProvider) + require.NoError(t, err) + + first, err := limiter.Allow(t.Context(), "search", "user-1") + require.NoError(t, err) + require.True(t, first.Allowed) + + second, err := limiter.Allow(t.Context(), "search", "user-1") + require.NoError(t, err) + require.False(t, second.Allowed, "the per-user tool bucket holds a single token") + + metrics := collectRateLimitMetrics(t, reader) + decisions := requireRateLimitMetric(t, metrics, "toolhive_rate_limit_decisions") + + // The rejection is attributable to the tool that caused it. + assert.Equal(t, int64(1), counterValueWithAttributes(t, decisions, map[string]string{ + "decision": rateLimitDecisionRejected, + "scope": rateLimitScopePerUser, + "operation_type": rateLimitOperationTool, + "tool_name": "search", + })) + + // Server-scoped decisions carry no tool_name at all, rather than an empty one. + sum, ok := decisions.Data.(metricdata.Sum[int64]) + require.True(t, ok) + var sawServerScoped bool + for _, point := range sum.DataPoints { + operationType, found := point.Attributes.Value(attribute.Key("operation_type")) + if !found || operationType.AsString() != rateLimitOperationServer { + continue + } + sawServerScoped = true + _, hasTool := point.Attributes.Value(attribute.Key("tool_name")) + assert.False(t, hasTool, "server-scoped decisions must not carry tool_name") + } + assert.True(t, sawServerScoped, "expected at least one server-scoped decision") +} + +func TestRateLimitMetrics_UnconfiguredToolCarriesNoAttribute(t *testing.T) { + t.Parallel() + client, _ := newTestClient(t) + reader, meterProvider := newRateLimitMeterProvider() + + // Only a server-wide bucket is configured: no tool has a bucket of its own. + limiter, err := newLimiter(client, "test-ns", "test-server", &v1beta1.RateLimitConfig{ + Shared: &v1beta1.RateLimitBucket{ + MaxTokens: 10, + RefillPeriod: metav1.Duration{Duration: time.Minute}, + }, + }, meterProvider) + require.NoError(t, err) + + // A caller-supplied tool name that matches no configured bucket must never reach + // the metric: tool_name is only ever set from a successful bucket-map lookup, so + // its value set is the operator's configuration, not caller input. This is the + // property that keeps the attribute bounded in both cardinality and value length. + allowed, err := limiter.Allow(t.Context(), strings.Repeat("a", 4096), "user-1") + require.NoError(t, err) + require.True(t, allowed.Allowed) + + metrics := collectRateLimitMetrics(t, reader) + decisions := requireRateLimitMetric(t, metrics, "toolhive_rate_limit_decisions") + sum, ok := decisions.Data.(metricdata.Sum[int64]) + require.True(t, ok) + require.NotEmpty(t, sum.DataPoints) + for _, point := range sum.DataPoints { + _, hasTool := point.Attributes.Value(attribute.Key("tool_name")) + assert.False(t, hasTool, "an unconfigured tool name must not become a label value") + } +}