feat(token-rate-limit): add shared Valkey quota enforcement - #790
feat(token-rate-limit): add shared Valkey quota enforcement#790nerdalert wants to merge 4 commits into
Conversation
praxis-bot
left a comment
There was a problem hiding this comment.
PR Review
Summary: Adds a token_rate_limit filter with in-memory and Valkey backends for authenticated, model-scoped token quota enforcement. Also adds producer-defined selection groups, a picker module for round-robin/random/deterministic selection, and pipeline ordering validation.
Overall: The design is well thought out and the in-memory ledger is solid. However, the Valkey Lua scripts have a key-scoping bug in the expiry cleanup path that will cause the active-tokens counter to inflate monotonically, eventually denying all requests. The on_response_body sync path also propagates reconciliation queue errors in a way that could disrupt response delivery.
| Severity | Count |
|---|---|
| Critical | 0 |
| Large | 1 |
| Medium | 2 |
| redis.call('HINCRBY', physical .. ':settled', bucket, amount) | ||
| redis.call('ZADD', physical .. ':settled-index', bucket, bucket) | ||
| redis.call('HDEL', active_key, reservation) | ||
| redis.call('DECRBY', physical .. ':active-tokens', amount) |
There was a problem hiding this comment.
[Large] The expiry cleanup path constructs key names by appending :active-tokens and :settled-index to the per-key physical prefix, but the admission and reservation paths use the rule-scoped KEYS[9] ({rule_prefix}:active-tokens) and KEYS[8] ({rule_prefix}:settled-index). These are different Redis keys.
Consequences:
KEYS[9](rule-scopedactive-tokens) is incremented on every reservation (line 211) and decremented on normal reconciliation (reconcile script line 231), but never decremented during expiry cleanup. Every expired reservation permanently inflates this counter.- The settled tokens from expired reservations are written to a stray per-key
settled-indexthat is never queried during admission (line 198 readsKEYS[8]), so those tokens vanish from capacity accounting. - Over time,
active_sumgrows without bound, causingsettled_sum + active_sum + estimate > capacityto deny all requests even when there are no actual active reservations.
Change lines 158-161 to use the rule-scoped keys that admission actually reads:
redis.call('HINCRBY', settled, bucket, amount)
redis.call('ZADD', settled_index, bucket, bucket)
redis.call('HDEL', active_key, reservation)
redis.call('DECRBY', active_tokens_key, amount)settled (line 136) is KEYS[2] which is per-key -- that is correct for the settled hash. But settled_index and active_tokens_key must be the rule-scoped variables assigned at the top of the script (lines 138-139), not constructed from physical.
| estimate: rule.estimate, | ||
| now_ms: self.now_ms(), | ||
| }) | ||
| .map_err(|error| -> FilterError { error.into() })?; |
There was a problem hiding this comment.
[Medium] on_response_body is a sync callback that propagates enqueue_reconcile errors via ?, returning a FilterError. If the Valkey reconciliation channel is full (bounded at 1024), this fails the response body path for a request that was already admitted and served. Reconciliation is a post-response accounting concern; its failure should not disrupt the response being delivered to the client.
Change this to log the error and continue rather than propagating it:
if let Err(error) = rule.backend.enqueue_reconcile(ReconcileRequest { ... }) {
tracing::error!(%error, "token-rate-limit: failed to enqueue reconciliation");
counter!("praxis_ai_token_rate_limit_backend_errors_total", "backend" => "valkey", "operation" => "enqueue_reconcile").increment(1);
}| return Err(format!("listener '{}': token_count must follow token_rate_limit", listener.name).into()); | ||
| } | ||
| Ok(()) | ||
| } |
There was a problem hiding this comment.
[Medium] validate_token_rate_limit_order has no unit tests. The existing test module covers provider boundary validation, misaligned clusters, and open security filters, but nothing exercises the three code paths here: (1) token_rate_limit after intelligent_route is rejected, (2) token_count before token_rate_limit is rejected, (3) correct ordering is accepted. Add tests paralleling the provider boundary tests to cover these constraints.
Signed-off-by: Brent Salisbury <bsalisbu@redhat.com>
Signed-off-by: Brent Salisbury <bsalisbu@redhat.com>
Keep the token quota branch buildable against the corrected Praxis identity\nmetadata API and current AI filter lifecycle. Preserve Basic Auth, token counting,\nprovider routing, and the existing bounded-route coverage while updating the\ncompatibility paths required by the stacked dependency. Signed-off-by: Brent Salisbury <bsalisbu@redhat.com>
9a9ecf4 to
ad4ccf7
Compare
Use the published Praxis dependency in normal CI. The compatibility workflow\ncontinues to apply its checked-out Praxis source explicitly, avoiding duplicate\npatch tables and personal Git sources in supply-chain validation. Signed-off-by: Brent Salisbury <bsalisbu@redhat.com>
Summary
Dicalimer 🔴 This is not for merge, but a fairly cleaned up effort at a groundwork for quotas used for this demo: Distributed token quota demonstration: grid-distributed-token-rate-limit:
https://github.com/praxis-proxy/experimental/tree/main/demos/grid-distributed-token-rate-limit
This PR adds authenticated, model-scoped token quota enforcement to Praxis AI through the new token_rate_limit filter.
The filter admits or rejects a request before provider routing begins. It reserves an estimated token amount against one or more sliding-window budgets, then reconciles that reservation with the actual token usage reported after the response completes.
The quota mechanism is independent of Grid routing:
It can therefore be used with Grid provider selection, static upstream routing, or another routing implementation. Grid does not store quotas, contact the quota backend, or participate in quota enforcement.
Project Stack
The ownership boundaries are intentional:
Quota enforcement remains outside Grid and does not add a quota lookup to Grid reconciliation or provider selection.
Related Work
AI token rate-limiting proposal: docs/proposals/00121_token-rate-limiting.md
https://github.com/praxis-proxy/ai/blob/main/docs/proposals/00121_token-rate-limiting.md
AI token-rate-limiting implementation: docs(proposal): add implementation to token rate-limiting #658
docs(proposal): add implementation to token rate-limiting #658
Praxis authenticated-principal metadata establishes the trusted identity.user_id contract consumed by this filter.
AI provider-selection foundation: feat(routing): add provider load balancing, feat(routing): add provider load balancing #731
feat(routing): add provider load balancing #731
Grid provider-selection groups: feat(routing): publish provider selection modes, feat(routing): publish provider selection modes grid#65
feat(routing): publish provider selection modes grid#65
Distributed token quota demonstration: grid-distributed-token-rate-limit
https://github.com/praxis-proxy/experimental/tree/main/demos/grid-distributed-token-rate-limit
AI PR #731 supplies the provider-selection and load-balancing foundation. The distributed quota behavior is demonstrated separately by the experimental demo.
Behavior
The filter performs the following request lifecycle:
Admission happens before routing. A rejected request therefore cannot select a provider or consume provider capacity.
Quota Keys
The current quota key is composed from two bounded inputs:
authenticated principal + configured model
For example:
The principal must come from trusted request metadata established by an authentication filter. The limiter does not trust a client-provided identity header and does not parse credentials itself.
The model is read from a configured request header and must match the configured model allowlist.
This provides independent quota accounting for combinations such as:
The implementation is suitable for non-Grid deployments as long as the pipeline provides:
Sliding-Window Budgets
Each rule can define multiple token budgets. Admission reserves capacity atomically across every configured window.
Example:
The request is admitted only when every applicable window has sufficient capacity.
The implementation uses one-second accounting buckets with explicit bounds on:
Sliding-window recovery occurs naturally as prior usage ages out of the configured window. It does not require clearing the backend or restarting a gateway.
Reservation and Reconciliation
Admission uses a fixed token estimate so the decision can be made before inference begins.
After the response completes, the reservation is reconciled with actual token usage:
This avoids admitting unlimited concurrent requests based only on completed usage.
Backends
In-Memory
The in-memory backend provides exact local sliding-window accounting for one gateway process.
It is appropriate for:
Its state is not shared across gateway replicas and does not survive process restart.
Valkey
The Valkey backend provides shared quota enforcement across multiple gateway processes.
Both gateways reserve against the same principal-and-model budget. A request through one gateway therefore reduces the capacity visible through the other gateway.
The backend uses:
Valkey URLs can be supplied through an environment-variable reference so credentials do not need to appear directly in the filter configuration.
rediss://is supported when the deployment provides the required trust roots and validates the server certificate. A password or private cluster network alone does not encrypt Valkey traffic.Routing Independence
The filter does not depend on:
A non-Grid deployment can use:
A Grid-aware deployment can use:
Pipeline validation ensures that, when these filters are present:
Provider selection can change between requests without changing the quota key. With the Valkey backend, quota also remains consistent when requests arrive through different gateway replicas.
Responses
When quota capacity is exhausted, the filter returns HTTP 429 before routing.
The response includes bounded rate-limit information such as:
Denied requests do not contain provider attribution because provider selection never occurred.
Backend failures use a distinct fail-closed service response rather than admitting requests without accounting.
Missing identity, missing model, unknown model, oversized keys, missing rules, and state-capacity failures are handled explicitly rather than silently bypassing quota enforcement.
Configuration Example
Security Considerations
The filter consumes an authenticated identity; it does not establish one.
identity.user_id must be written only by a trusted authentication filter after successful verification. Clients must not be able to set or override this metadata directly.
The implementation does not place the following in quota metadata, metric labels, error bodies, or logs:
Principal and model values are bounded before allocating state. Metrics use bounded result, operation, token-kind, and backend labels rather than raw user or model identifiers.
Valkey credentials should be supplied through a Secret-backed environment variable. Production deployments should use encrypted and authenticated transport.
Performance and Hot-Path Behavior
The filter adds quota admission to the request path, so its backend choice defines the operational tradeoff.
In-Memory Backend
The local backend performs bounded in-process accounting and synchronization. It does not make network, Kubernetes, filesystem, Grid, or metrics-service calls.
Valkey Backend
The shared backend performs a Valkey operation during admission so multiple gateways can enforce one quota atomically. This is an intentional distributed consistency boundary.
The implementation limits its cost through:
No Grid call, Kubernetes lookup, provider-health query, Prometheus scrape, EPP request, or overlay parsing is introduced by quota enforcement.
Quota denial occurs before provider routing and inference, avoiding unnecessary downstream work.
Observability
The filter records bounded operational metrics for:
Raw principals and complete quota keys are not used as Prometheus labels.
The distributed demo additionally presents request-level quota decisions, selected providers for admitted requests, pre-provider 429 behavior, shared enforcement across gateway replicas, and sliding-window recovery.
Compatibility
The filter is additive and must be explicitly included in a pipeline.
Deployments that do not configure token_rate_limit retain their existing behavior.
The implementation does not change:
The filter can be introduced independently of Grid. A Grid-aware deployment requires the provider-selection work only for the demonstrated round-robin routing behavior, not for quota enforcement itself.
Current Scope
This PR implements:
The current scope does not yet provide:
Those capabilities can build on the same authenticated-identity and backend contracts without coupling quota enforcement to provider routing.
Validation
Validation covers:
The distributed demonstration proves:
Checklist
Breaking Changes
No breaking changes are intended.
The filter is opt-in. Existing listeners and filter pipelines remain unchanged unless token_rate_limit is explicitly configured.