Skip to content

feat(token-rate-limit): add shared Valkey quota enforcement - #790

Draft
nerdalert wants to merge 4 commits into
praxis-proxy:mainfrom
nerdalert:feat/token-rate-limit-valkey
Draft

feat(token-rate-limit): add shared Valkey quota enforcement#790
nerdalert wants to merge 4 commits into
praxis-proxy:mainfrom
nerdalert:feat/token-rate-limit-valkey

Conversation

@nerdalert

@nerdalert nerdalert commented Aug 20, 2026

Copy link
Copy Markdown
Member

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:

  Authentication
        |
        | trusted identity.user_id
        v
  token_rate_limit
        |
        | admitted request
        v
  Routing
    - intelligent_route
    - static routing
    - another load balancer
        |
        v
  Inference backend
        |
        | actual token usage
        v
  token_rate_limit reconciliation

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


  Praxis core
    Establishes trusted authenticated identity metadata
    identity.user_id
          |
          v
  Praxis AI
    Reserves and reconciles token quota
    Returns 429 before routing when quota is exhausted
          |
          v
  Grid + Praxis AI provider selection
    Selects an eligible provider only for admitted requests
          |
          v
  Provider gateway and inference stack
    Processes the admitted request

The ownership boundaries are intentional:

  • Praxis core authenticates the caller and publishes trusted request-scoped identity metadata.
  • Praxis AI owns quota admission, reservation, reconciliation, and quota responses.
  • Grid determines provider eligibility and routing state.
  • intelligent_route selects a provider from the current in-memory routing snapshot.
  • The inference stack executes the admitted request and reports token usage.

Quota enforcement remains outside Grid and does not add a quota lookup to Grid reconciliation or provider selection.

Related Work

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:

  Authenticated request
          |
          v
  Read trusted principal metadata
          |
          v
  Read and validate configured model
          |
          v
  Select matching quota rule
          |
          v
  Atomically reserve estimated tokens
         / \
        /   \
   admitted  capacity exhausted
      |              |
      v              v
   routing         HTTP 429
      |          no provider selected
      v          no backend contacted
   response
      |
      v
  Read actual token usage
      |
      v
  Reconcile reservation
    - refund unused estimate
    - charge overage
    - conservatively retain estimate when usage is unavailable

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:

  alice/Qwen/Qwen3-0.6B

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:

  alice/model-a
  alice/model-b
  bob/model-a

The implementation is suitable for non-Grid deployments as long as the pipeline provides:

  • trusted authenticated-principal metadata;
  • the configured model header;
  • token-usage metadata for final reconciliation.

Sliding-Window Budgets

Each rule can define multiple token budgets. Admission reserves capacity atomically across every configured window.

Example:

  rules:
    - name: standard-users
      match:
        metadata:
          identity.user_id: alice
      estimation:
        strategy: fixed
        tokens: 500
      token_budgets:
        - window: 1m
          capacity: 3000
        - window: 1h
          capacity: 50000

The request is admitted only when every applicable window has sufficient capacity.

The implementation uses one-second accounting buckets with explicit bounds on:

  • window duration;
  • number of budgets;
  • number of rules;
  • active keys;
  • key length;
  • active reservations.

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:

  • If actual usage is lower, unused reserved capacity is refunded.
  • If actual usage is higher, the overage is charged.
  • If usage cannot be determined, the original estimate remains charged.
  • Repeated settlement attempts are idempotent.
  • Expired or abandoned reservations are handled conservatively.

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:

  • development;
  • focused testing;
  • single-process deployments;
  • environments where per-instance quotas are intentional.

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.

  Consumer Gateway A ─┐
                      ├── shared Valkey quota ledger
  Consumer Gateway B ─┘

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:

  • atomic Lua-based admission across all configured windows;
  • cached scripts;
  • a reconnecting connection manager;
  • bounded one-second accounting buckets;
  • idempotent reservation reconciliation;
  • fail-closed behavior when the backend is unavailable.

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:

  • Grid;
  • Kubernetes;
  • routing overlays;
  • intelligent_route;
  • provider selection groups;
  • provider metrics;
  • EPP;
  • Prometheus;
  • a specific inference backend.

A non-Grid deployment can use:

  basic_auth
    -> token_rate_limit
    -> static router
    -> load_balancer
    -> token_count

A Grid-aware deployment can use:

  basic_auth
    -> token_rate_limit
    -> intelligent_route
    -> load_balancer
    -> token_count

Pipeline validation ensures that, when these filters are present:

  • token_rate_limit runs before intelligent_route;
  • token-usage processing occurs after quota admission.

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:

  • configured limit;
  • remaining capacity;
  • reset time;
  • Retry-After.

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

  - type: token_rate_limit
    config:
      key:
        principal:
          source: metadata
          name: identity.user_id
          onMissing: reject
        model:
          source: header
          name: x-model
          onMissing: reject
          allowedModels:
            - Qwen/Qwen3-0.6B
      reservationTimeout: 2m
      limits:
        maxKeys: 10000
        maxKeyLength: 256
        maxActiveReservations: 50000
      rules:
        - name: alice
          match:
            metadata:
              identity.user_id: alice
          estimation:
            strategy: fixed
            tokens: 15
          token_budgets:
            - window: 1m
              capacity: 60
        - name: default
          estimation:
            strategy: fixed
            tokens: 15
          token_budgets:
            - window: 1m
              capacity: 30
      backend:
        kind: valkey
        url: ${TOKEN_RATE_LIMIT_VALKEY_URL}
        namespace: praxis-ai

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:

  • passwords;
  • authorization headers;
  • API keys;
  • Valkey credentials;
  • prompts;
  • completions;
  • raw request bodies.

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:

  • cached Lua scripts;
  • a reconnecting connection manager;
  • bounded rule and window counts;
  • bounded key length and cardinality;
  • bounded one-second bucket count;
  • one atomic multi-window reservation operation;
  • asynchronous final reconciliation.

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:

  • admitted and denied requests;
  • denial reasons;
  • estimated, actual, refunded, and overage tokens;
  • created, reconciled, and orphaned reservations;
  • backend operation failures;
  • cleanup activity;
  • active keys;
  • active reservations.

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:

  • Grid scoring;
  • provider admission;
  • routing overlays;
  • provider selection;
  • session affinity;
  • existing static routing;
  • inference backend behavior.

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:

  • authenticated principal and model quota keys;
  • fixed token estimates;
  • one or more sliding-window budgets;
  • atomic reservations;
  • actual-usage reconciliation;
  • in-memory accounting;
  • shared Valkey accounting;
  • fail-closed backend behavior;
  • pre-routing 429 responses;
  • bounded operational metrics.

The current scope does not yet provide:

  • separate prompt, completion, and reasoning-token budgets;
  • JWT or OIDC authentication;
  • a quota-management API;
  • dynamic quota-policy reload;
  • administrative inspection of individual active quotas;
  • centralized policy distribution;
  • direct Grid ownership of quotas.

Those capabilities can build on the same authenticated-identity and backend contracts without coupling quota enforcement to provider routing.

Validation

Validation covers:

  • filter configuration parsing and strict unknown-field rejection;
  • missing and invalid principal handling;
  • missing and unknown model handling;
  • key-length and state-cardinality limits;
  • rule matching and default-rule behavior;
  • fixed estimation;
  • multiple atomic budget windows;
  • quota exhaustion;
  • sliding-window recovery;
  • reservation refunds and overages;
  • conservative missing-usage settlement;
  • idempotent reconciliation;
  • reservation expiry;
  • in-memory backend behavior;
  • Valkey backend conformance;
  • concurrent atomic admission;
  • shared quota across gateway processes;
  • restart persistence with Valkey;
  • fail-closed Valkey behavior;
  • rejection before provider contact;
  • compatibility with round-robin provider routing.

The distributed demonstration proves:

  • two gateway processes sharing one Valkey quota;
  • one authenticated principal and model quota across both gateways;
  • admitted requests distributed across three provider gateways;
  • quota exhaustion observed consistently from either consumer gateway;
  • HTTP 429 before provider selection;
  • no backend request for denied traffic;
  • persistence across a consumer-gateway restart;
  • recovery as usage leaves the sliding window;
  • HTTP 503 fail-closed behavior during Valkey failure.

Checklist

  • I reviewed every changed line and can explain the change.
  • New behavior includes focused configuration and filter tests.
  • Generated filter documentation is updated.
  • Distributed backend behavior has runtime validation.
  • Security-sensitive keying uses authenticated metadata.
  • State, keys, rules, windows, and reservations are bounded.
  • Request-path dependencies and performance implications are documented.
  • Commits are signed and include a Signed-off-by trailer.

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.

@praxis-bot praxis-bot left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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-scoped active-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-index that is never queried during admission (line 198 reads KEYS[8]), so those tokens vanish from capacity accounting.
  • Over time, active_sum grows without bound, causing settled_sum + active_sum + estimate > capacity to 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() })?;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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);
}

Comment thread server/src/pipelines.rs
return Err(format!("listener '{}': token_count must follow token_rate_limit", listener.name).into());
}
Ok(())
}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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>
@nerdalert
nerdalert force-pushed the feat/token-rate-limit-valkey branch from 9a9ecf4 to ad4ccf7 Compare August 22, 2026 05:31
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants