Skip to content

feat(token_rate_limit): configurable estimation strategies - #1008

Open
asaadbalum wants to merge 1 commit into
praxis-proxy:mainfrom
asaadbalum:feat/issue-882-estimation-strategies
Open

feat(token_rate_limit): configurable estimation strategies#1008
asaadbalum wants to merge 1 commit into
praxis-proxy:mainfrom
asaadbalum:feat/issue-882-estimation-strategies

Conversation

@asaadbalum

@asaadbalum asaadbalum commented Sep 8, 2026

Copy link
Copy Markdown

Summary

Replace the fixed reserved_tokens field with pluggable per-rule estimation
strategies. Each rule selects a strategy that computes the estimated token
cost from request metadata before forwarding.

Built-in strategies:

Strategy Basis Use case
fixed constant per request Uniform cost (current behavior, backward-compatible)
max_tokens max_tokens request field Simple output upper bound
input_plus_max_tokens Content-Length + max_tokens Conservative full-cost estimate
model_scaled max_tokens × model multiplier Model-aware budgets

Config shape:

rules:
  - name: team-alpha
    algorithm: sliding_window
    window: 1h
    capacity: 100000
    estimation:
      strategy: max_tokens
      multiplier: 1.2
      fallback_estimate: 500

Body-dependent strategies conditionally enable request body buffering
(ReadOnly + StreamBuffer), deferring reservation from on_request to
on_request_body. Fixed-only configs preserve the existing on_request flow
with zero overhead change.

reserved_tokens remains supported as shorthand for the fixed strategy
with full backward compatibility — all 118 pre-existing tests pass unchanged.

Related issue

Closes #882

Validation

  • Unit tests: 157 pass (118 existing + 39 new)
  • Config parsing: all 4 strategies, mutual exclusion, parameter validation
  • Body-dependent flow: on_request deferral, body extraction, fallback chains
  • Edge cases: empty body, malformed JSON, missing Content-Length, mixed rules
  • model_scaled: per-model multipliers, default fallback, x-model header, outer multiplier composition
  • Reconciliation: META_ESTIMATE stashing, overestimate/underestimate settlement
  • Backward compat: reserved_tokens: N compiles to Fixed { estimate: N }
  • make lint equivalent: clippy (all features) + fmt clean
  • Full workspace build clean

Checklist

  • Conventional commit with DCO sign-off
  • No new dependencies
  • Behind existing token-rate-limit-filter feature gate (experimental)
  • Doc comments updated (deferred list corrected, YAML example added)
  • No backend changes (backend.rs, ledger.rs, token_bucket_ledger.rs untouched)

@asaadbalum
asaadbalum requested review from a team and leseb September 8, 2026 15:02
@praxis-bot-app

praxis-bot-app Bot commented Sep 8, 2026

Copy link
Copy Markdown

Unsigned commits: 41daf46. Please sign your commits.

@asaadbalum
asaadbalum requested a review from szedan-rh September 8, 2026 15:06
@asaadbalum
asaadbalum force-pushed the feat/issue-882-estimation-strategies branch from 41daf46 to 6ee2cd5 Compare September 8, 2026 15:06
@asaadbalum asaadbalum changed the title feat(token_rate_limit): configurable estimation strategies (M3) feat(token_rate_limit): configurable estimation strategies Sep 8, 2026
@asaadbalum
asaadbalum force-pushed the feat/issue-882-estimation-strategies branch from 6ee2cd5 to 20fbdaa Compare September 8, 2026 15:28

@praxis-bot praxis-bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

praxis-bot review: feat(token_rate_limit): configurable estimation strategies

Well-structured feature. The compiled estimation pattern, body-deferral mechanism, backward-compatible reserved_tokens path, and comprehensive test coverage (39 new tests, 118 existing unchanged) are all solid. Two findings below.

Findings: 0 Critical, 0 Large, 2 Medium

Comment thread filters/src/token_rate_limit/mod.rs Outdated
let estimate = ctx
.get_metadata(META_ESTIMATE)
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(0);

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

[Medium] META_ESTIMATE defaults to 0 when missing, which produces incorrect settlement math.

If this fallback is ever reached (metadata corruption, future code path that skips handle_reserve_outcome), the backend will settle with estimate=0 against whatever was actually reserved at admission. For a request that reserved 500 tokens with actual=200, the settlement computes overage = 200 - 0 = 200 and adds 200 tokens to the window on top of the 500 already reserved, double-counting usage (700 charged instead of 200).

Since handle_reserve_outcome always sets META_ESTIMATE, this path should be unreachable today. But the silent unwrap_or(0) is a latent correctness hazard.

Requested change: return None from reconciliation_context when META_ESTIMATE is missing (skip reconciliation, leaving the reservation as-is), or at minimum log a warning before defaulting:

let estimate = match ctx.get_metadata(META_ESTIMATE).and_then(|v| v.parse::<u64>().ok()) {
    Some(est) => est,
    None => {
        tracing::warn!("token_rate_limit: META_ESTIMATE missing at reconciliation, skipping");
        return None;
    },
};

Skipping reconciliation is safer than settling with a known-wrong estimate -- the reservation's original charge stays in place, which is conservative.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Accepted. Defaulting a missing META_ESTIMATE to 0 would settle actual - 0 on top of the original reservation and over-charge the window. reconciliation_context now returns None (reservation stands) and logs a warning. Covered by missing_meta_estimate_skips_reconciliation_instead_of_settling_at_zero.

})
},
EstimationStrategy::ModelScaled => compile_model_scaled_estimation(rule_name, est, multiplier, capacity),
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

[Medium] Strategy-irrelevant config fields are silently accepted.

EstimationConfig is a flat struct with all strategy knobs (model_multipliers, default_multiplier, bytes_per_token), but compile_estimation_config only reads the fields relevant to the chosen strategy. Fields that don't apply are silently ignored:

estimation:
  strategy: fixed
  fallback_estimate: 100
  model_multipliers:        # silently ignored
    gpt-4: 2.0
  bytes_per_token: 3.5      # silently ignored

This violates the project's deny_unknown_fields philosophy and can mask config errors (e.g., an operator migrating from model_scaled to fixed who forgets to remove model_multipliers).

Requested change: add strategy-specific field presence validation in each match arm. For example:

EstimationStrategy::Fixed => {
    if est.model_multipliers.is_some() || est.default_multiplier.is_some() || est.bytes_per_token.is_some() {
        return Err(format!(
            "token_rate_limit: rule '{rule_name}': fixed strategy does not accept \
             model_multipliers, default_multiplier, or bytes_per_token"
        ).into());
    }
    compile_fixed_estimation(rule_name, est.fallback_estimate, multiplier, capacity)
},

Similar guards for MaxTokens (reject model_multipliers, default_multiplier, bytes_per_token) and InputPlusMaxTokens (reject model_multipliers, default_multiplier).

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Accepted. Silently ignoring strategy-irrelevant knobs fights deny_unknown_fields and can hide a leftover model_multipliers/bytes_per_token after a strategy change. compile now rejects those fields per strategy. Covered by from_config_rejects_strategy_irrelevant_estimation_fields.

@asaadbalum
asaadbalum force-pushed the feat/issue-882-estimation-strategies branch from 07a0beb to bb10c28 Compare September 9, 2026 14:04
Replace the fixed reserved_tokens field with pluggable per-rule estimation
strategies. Each rule selects a strategy that computes the estimated token
cost from request metadata before forwarding.

Built-in strategies:
- fixed: constant per request (backward-compatible with reserved_tokens)
- max_tokens: extracted from request body max_tokens field
- input_plus_max_tokens: Content-Length-based input estimate plus max_tokens
- model_scaled: max_tokens scaled by per-model multiplier

Config shape:
  estimation:
    strategy: max_tokens
    multiplier: 1.2          # optional safety margin (default: 1.0)
    fallback_estimate: 500   # used when max_tokens absent

Body-dependent strategies conditionally enable request body buffering
(ReadOnly + StreamBuffer), deferring reservation from on_request to
on_request_body. Fixed-only configs preserve the existing on_request flow
with no body access overhead.

reserved_tokens remains supported as shorthand for the fixed strategy
with full backward compatibility -- all 118 existing tests pass unchanged.

Ref: ai#882
Signed-off-by: Asaad Balum <asaad.balum@gmail.com>
@asaadbalum
asaadbalum force-pushed the feat/issue-882-estimation-strategies branch from bb10c28 to d637bbf Compare September 9, 2026 14:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Configurable estimation strategies for request-time cost prediction

2 participants