Open
Conversation
… based in token-cost-based-ratelimit
… based in llm-cost-based-ratelimit
Contributor
WalkthroughAdvanced-ratelimit namespaced per-instance metadata and supports key-extraction fallbacks; api-key-auth writes resolved ApplicationID into shared metadata; llm-cost-based-ratelimit and token-based-ratelimit gained a Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~40 minutes Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
policies/llm-cost-based-ratelimit/llm_cost_based_ratelimit_test.go (1)
537-571: Add an explicit test for the “absent consumerBased” path.This test currently verifies only explicit
consumerBased: false, while the test description also claims behavior for the absent case.Proposed test addition
+func TestTransformToRatelimitParams_ConsumerBased_Absent(t *testing.T) { + params := map[string]interface{}{ + "budgetLimits": []interface{}{ + map[string]interface{}{"amount": float64(10), "duration": "1h"}, + }, + } + + result := transformToRatelimitParams(params) + quotas, ok := result["quotas"].([]interface{}) + if !ok || len(quotas) != 1 { + t.Fatal("Expected 1 quota") + } + + quota := quotas[0].(map[string]interface{}) + keyExtraction, ok := quota["keyExtraction"].([]interface{}) + if !ok || len(keyExtraction) != 1 { + t.Fatalf("Expected only routename key extraction, got %v", keyExtraction) + } +}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@policies/llm-cost-based-ratelimit/llm_cost_based_ratelimit_test.go` around lines 537 - 571, Update the tests to cover the "consumerBased absent" path by adding a new test (e.g., TestTransformToRatelimitParams_ConsumerBased_Absent) or extending TestTransformToRatelimitParams_ConsumerBased_False to call transformToRatelimitParams with params that omit the "consumerBased" key (keep the same budgetLimits shape) and assert the resulting quotas[0].keyExtraction contains exactly one entry with entry["type"] == "routename", mirroring the existing assertions so the absent-key behavior is explicitly verified.policies/token-based-ratelimit/token_based_ratelimit_test.go (1)
480-514: Consider adding a dedicated “consumerBased omitted” test.This test validates explicit
falsewell, but not the omitted-parameter default path mentioned in the comment.Proposed test addition
+func TestTransformToRatelimitParams_ConsumerBased_Absent(t *testing.T) { + params := map[string]interface{}{ + "totalTokenLimits": []interface{}{ + map[string]interface{}{"count": float64(1000), "duration": "1h"}, + }, + } + + result := transformToRatelimitParams(params, nil) + quotas, ok := result["quotas"].([]interface{}) + if !ok || len(quotas) != 1 { + t.Fatal("Expected 1 quota") + } + + quota := quotas[0].(map[string]interface{}) + keyExtraction, ok := quota["keyExtraction"].([]interface{}) + if !ok || len(keyExtraction) != 1 { + t.Fatalf("Expected only routename key extraction, got %v", keyExtraction) + } +}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@policies/token-based-ratelimit/token_based_ratelimit_test.go` around lines 480 - 514, Add a new unit test (e.g., TestTransformToRatelimitParams_ConsumerBased_Omitted) that omits the "consumerBased" key from the input params and calls transformToRatelimitParams(nil) to verify the default behavior matches the existing false case: assert there is exactly one quota, that quota.keyExtraction is a []interface{} of length 1, and that the single entry has type "routename"; mirror the same assertions used in TestTransformToRatelimitParams_ConsumerBased_False so the omitted-parameter path is explicitly covered.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@policies/llm-cost-based-ratelimit/policy-definition.yaml`:
- Around line 45-54: Update the top-level policy description to reflect that
per-consumer spending limits are optional by referencing the consumerBased
boolean flag (consumerBased defaults to false); change language that currently
states the policy "enforces per-consumer spending limits" to something like "can
enforce per-consumer spending limits when consumerBased is true, otherwise a
shared limit applies" and mention the default behavior (consumerBased: false) so
readers know per-consumer enforcement is opt-in.
---
Nitpick comments:
In `@policies/llm-cost-based-ratelimit/llm_cost_based_ratelimit_test.go`:
- Around line 537-571: Update the tests to cover the "consumerBased absent" path
by adding a new test (e.g., TestTransformToRatelimitParams_ConsumerBased_Absent)
or extending TestTransformToRatelimitParams_ConsumerBased_False to call
transformToRatelimitParams with params that omit the "consumerBased" key (keep
the same budgetLimits shape) and assert the resulting quotas[0].keyExtraction
contains exactly one entry with entry["type"] == "routename", mirroring the
existing assertions so the absent-key behavior is explicitly verified.
In `@policies/token-based-ratelimit/token_based_ratelimit_test.go`:
- Around line 480-514: Add a new unit test (e.g.,
TestTransformToRatelimitParams_ConsumerBased_Omitted) that omits the
"consumerBased" key from the input params and calls
transformToRatelimitParams(nil) to verify the default behavior matches the
existing false case: assert there is exactly one quota, that quota.keyExtraction
is a []interface{} of length 1, and that the single entry has type "routename";
mirror the same assertions used in
TestTransformToRatelimitParams_ConsumerBased_False so the omitted-parameter path
is explicitly covered.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b39ba54a-6e4f-46b2-9669-8de56ab93663
📒 Files selected for processing (11)
policies/advanced-ratelimit/policy-definition.yamlpolicies/advanced-ratelimit/ratelimit.gopolicies/api-key-auth/apikey.gopolicies/api-key-auth/apikey_test.gopolicies/api-key-auth/policy-definition.yamlpolicies/llm-cost-based-ratelimit/llm_cost_based_ratelimit.gopolicies/llm-cost-based-ratelimit/llm_cost_based_ratelimit_test.gopolicies/llm-cost-based-ratelimit/policy-definition.yamlpolicies/token-based-ratelimit/policy-definition.yamlpolicies/token-based-ratelimit/token_based_ratelimit.gopolicies/token-based-ratelimit/token_based_ratelimit_test.go
Comment on lines
+45
to
+54
| consumerBased: | ||
| type: boolean | ||
| x-wso2-policy-advanced-param: false | ||
| description: | | ||
| When true, rate limits are applied per consumer (GenAI application) identified | ||
| by the x-wso2-application-id metadata key set by the api-key-auth policy. | ||
| Each application gets its own independent cost counter. | ||
| When false (default), a single shared limit applies across all consumers. | ||
| default: false | ||
|
|
Contributor
There was a problem hiding this comment.
Align top-level policy description with conditional behavior.
With consumerBased defaulting to false, the file description at Line 6 (“enforces per-consumer spending limits”) is now too absolute. Please make it explicit that per-consumer enforcement is optional.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@policies/llm-cost-based-ratelimit/policy-definition.yaml` around lines 45 -
54, Update the top-level policy description to reflect that per-consumer
spending limits are optional by referencing the consumerBased boolean flag
(consumerBased defaults to false); change language that currently states the
policy "enforces per-consumer spending limits" to something like "can enforce
per-consumer spending limits when consumerBased is true, otherwise a shared
limit applies" and mention the default behavior (consumerBased: false) so
readers know per-consumer enforcement is opt-in.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@policies/advanced-ratelimit/ratelimit.go`:
- Around line 861-868: The parser now reads keyExtraction[*].fallback into
comp.Fallback, but the policy schema still disallows that field; update
policies/advanced-ratelimit/policy-definition.yaml to add a "fallback" property
(string) to the keyExtraction item schema (or allow it via properties instead of
being blocked by additionalProperties: false) so configurations pass validation
at policy creation/attachment; ensure the schema entry for the keyExtraction
array includes the new "fallback" property with type: string (and any desired
default/description) to match the parser's expectation.
- Around line 394-410: The instanceID currently built from idBuf/quotas only
includes quota.Name and KeyExtraction fields, causing collisions when other
policy parameters differ; extend the deterministic hash input to include all
distinguishing configuration used by advanced-ratelimit (e.g., each quota's CEL
expressions, MetadataFallback values, limit values, cost extraction
identifiers/flags, and any other per-quota settings) before computing idHash so
instanceID uniquely namespaces SharedContext.Metadata for different policy
instances; update the loop that builds idBuf (and any other place around
instanceID/idHash usage such as the similar block at lines ~434-438) to append
those additional fields in a stable order prior to sha256.Sum256 and
hex.EncodeToString.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 146e3b18-e269-4aa6-8799-d952ea2010cd
📒 Files selected for processing (6)
policies/advanced-ratelimit/ratelimit.gopolicies/advanced-ratelimit/ratelimit_integration_test.gopolicies/llm-cost-based-ratelimit/llm_cost_based_ratelimit.gopolicies/llm-cost-based-ratelimit/llm_cost_based_ratelimit_test.gopolicies/token-based-ratelimit/token_based_ratelimit.gopolicies/token-based-ratelimit/token_based_ratelimit_test.go
🚧 Files skipped from review as they are similar to previous changes (3)
- policies/llm-cost-based-ratelimit/llm_cost_based_ratelimit_test.go
- policies/llm-cost-based-ratelimit/llm_cost_based_ratelimit.go
- policies/token-based-ratelimit/token_based_ratelimit.go
Comment on lines
+394
to
+410
| // Compute a short deterministic instance ID from the quota names and key extraction | ||
| // types. This namespaces all SharedContext.Metadata keys so that multiple | ||
| // advanced-ratelimit instances in the same request chain do not overwrite each other. | ||
| var idBuf strings.Builder | ||
| for i, q := range quotas { | ||
| idBuf.WriteString(q.Name) | ||
| idBuf.WriteString(":") | ||
| for _, ke := range q.KeyExtraction { | ||
| idBuf.WriteString(ke.Type) | ||
| idBuf.WriteString(ke.Key) | ||
| } | ||
| if i < len(quotas)-1 { | ||
| idBuf.WriteString("|") | ||
| } | ||
| } | ||
| idHash := sha256.Sum256([]byte(idBuf.String())) | ||
| instanceID := hex.EncodeToString(idHash[:])[:8] |
Contributor
There was a problem hiding this comment.
Broaden the instanceID hash; it still collides for distinct policy instances.
instanceID is derived only from quota name plus keyExtraction type/key. Two advanced-ratelimit instances with the same quota names and key types/keys but different CEL expressions, metadata fallbacks, limits, or cost-extraction behavior will still share the same metadata namespace and overwrite each other’s per-request state. That brings back the collision bug for a pretty realistic API-level + route-level combination.
💡 Suggested direction
var idBuf strings.Builder
for i, q := range quotas {
- idBuf.WriteString(q.Name)
- idBuf.WriteString(":")
+ idBuf.WriteString(fmt.Sprintf("[%d]%s:", i, q.Name))
+ for _, lim := range q.Limits {
+ idBuf.WriteString(fmt.Sprintf("l=%d,d=%s,b=%d;", lim.Limit, lim.Duration, lim.Burst))
+ }
for _, ke := range q.KeyExtraction {
- idBuf.WriteString(ke.Type)
- idBuf.WriteString(ke.Key)
+ idBuf.WriteString(ke.Type)
+ idBuf.WriteString("|")
+ idBuf.WriteString(ke.Key)
+ idBuf.WriteString("|")
+ idBuf.WriteString(ke.Expression)
+ idBuf.WriteString("|")
+ idBuf.WriteString(ke.Fallback)
+ idBuf.WriteString(";")
}
+ if q.CostExtractionEnabled && q.CostExtractor != nil {
+ idBuf.WriteString("cost-extraction-enabled;")
+ }
if i < len(quotas)-1 {
idBuf.WriteString("|")
}
}Also applies to: 434-438
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@policies/advanced-ratelimit/ratelimit.go` around lines 394 - 410, The
instanceID currently built from idBuf/quotas only includes quota.Name and
KeyExtraction fields, causing collisions when other policy parameters differ;
extend the deterministic hash input to include all distinguishing configuration
used by advanced-ratelimit (e.g., each quota's CEL expressions, MetadataFallback
values, limit values, cost extraction identifiers/flags, and any other per-quota
settings) before computing idHash so instanceID uniquely namespaces
SharedContext.Metadata for different policy instances; update the loop that
builds idBuf (and any other place around instanceID/idHash usage such as the
similar block at lines ~434-438) to append those additional fields in a stable
order prior to sha256.Sum256 and hex.EncodeToString.
Comment on lines
+861
to
+868
| // Parse optional fallback value | ||
| if fallbackRaw, ok := compMap["fallback"]; ok { | ||
| if fallbackStr, ok := fallbackRaw.(string); ok { | ||
| comp.Fallback = fallbackStr | ||
| } else { | ||
| return nil, fmt.Errorf("keyExtraction[%d].fallback must be a string", i) | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
fallback is still unconfigurable through the policy schema.
The parser now accepts keyExtraction[*].fallback, but policies/advanced-ratelimit/policy-definition.yaml still rejects that field (additionalProperties: false and no fallback property). In practice, any policy using this new behavior will be rejected before GetPolicy() runs, so the feature is not actually usable yet. Please update the schema in the same PR.
Based on learnings, policy configurations are validated against policy-definition.yaml schemas at API creation time or when policies are attached.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@policies/advanced-ratelimit/ratelimit.go` around lines 861 - 868, The parser
now reads keyExtraction[*].fallback into comp.Fallback, but the policy schema
still disallows that field; update
policies/advanced-ratelimit/policy-definition.yaml to add a "fallback" property
(string) to the keyExtraction item schema (or allow it via properties instead of
being blocked by additionalProperties: false) so configurations pass validation
at policy creation/attachment; ensure the schema entry for the keyExtraction
array includes the new "fallback" property with type: string (and any desired
default/description) to match the parser's expectation.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Adds consumer-based rate limiting support to gateway-controllers policies. Currently, token and cost rate limits are shared across all GenAI applications — if App A exhausts the token budget, App B gets blocked too. This change allows each application to have its own independent counter against the same limit.
Goals
x-wso2-application-id) gets its own independent counterApproach
api-key-authAfter a successful API key authentication, write the application ID into
SharedContext.Metadata(in addition toAuthContext.PropertiesandAnalyticsMetadatawhere it was already written). This makes it accessible to downstream rate limiting policies via key extraction.token-based-ratelimitconsumerBasedboolean parameter topolicy-definition.yamlconsumerBased: true, key extraction switches fromroutenamealone toroutename + x-wso2-application-id, giving each application its own counterllm-cost-based-ratelimitconsumerBasedboolean parameter topolicy-definition.yamltoken-based-ratelimitllm_cost_delegate. The backend's response phase then read the consumer's delegate and drained the consumer counter twice while never touching the backend counter. Fixed by using separate keys:llm_cost_delegatefor backend andllm_cost_delegate_consumerfor consumer.advanced-ratelimitFixed a metadata key collision that occurs when two instances with different key extraction configs run in the same request chain (e.g. a backend request-count limit and a consumer request-count limit):
advanced-ratelimitpasses state between its request and response phases via hardcoded metadata keys (ratelimit:keys,ratelimit:result, etc.)routenameand a consumer instance writesroutename:app-id, the backend's response phase reads the consumer's key and deducts from the wrong counter — backend budget is never exhaustedinstanceIDfield toRateLimitPolicy— a short SHA-256 hash of the quota names and key extraction config. All metadata reads/writes go throughp.metaKey(base)which appends the instance ID, namespacing each instance's metadata slots. Limiter objects and the limiter cache are unaffected.Release note
Added consumer-based rate limiting support to
token-based-ratelimit,llm-cost-based-ratelimit, andadvanced-ratelimitpolicies. A newconsumerBasedparameter allows limits to be tracked independently per GenAI application rather than shared across all applications. Also fixed metadata key collisions inadvanced-ratelimitandllm-cost-based-ratelimitthat caused incorrect counter updates when backend and consumer instances ran in the same request chain.Documentation
N/A
Summary by CodeRabbit