feat(token_rate_limit): configurable estimation strategies - #1008
feat(token_rate_limit): configurable estimation strategies#1008asaadbalum wants to merge 1 commit into
Conversation
|
Unsigned commits: 41daf46. Please sign your commits. |
41daf46 to
6ee2cd5
Compare
6ee2cd5 to
20fbdaa
Compare
praxis-bot
left a comment
There was a problem hiding this comment.
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
| let estimate = ctx | ||
| .get_metadata(META_ESTIMATE) | ||
| .and_then(|v| v.parse::<u64>().ok()) | ||
| .unwrap_or(0); |
There was a problem hiding this comment.
[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.
There was a problem hiding this comment.
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), | ||
| } |
There was a problem hiding this comment.
[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 ignoredThis 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).
There was a problem hiding this comment.
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.
07a0beb to
bb10c28
Compare
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>
bb10c28 to
d637bbf
Compare
Summary
Replace the fixed
reserved_tokensfield with pluggable per-rule estimationstrategies. Each rule selects a strategy that computes the estimated token
cost from request metadata before forwarding.
Built-in strategies:
fixedmax_tokensmax_tokensrequest fieldinput_plus_max_tokensmax_tokensmodel_scaledmax_tokens× model multiplierConfig shape:
Body-dependent strategies conditionally enable request body buffering
(
ReadOnly+StreamBuffer), deferring reservation fromon_requesttoon_request_body. Fixed-only configs preserve the existingon_requestflowwith zero overhead change.
reserved_tokensremains supported as shorthand for the fixed strategywith full backward compatibility — all 118 pre-existing tests pass unchanged.
Related issue
Closes #882
Validation
on_requestdeferral, body extraction, fallback chainsreserved_tokens: Ncompiles toFixed { estimate: N }make lintequivalent: clippy (all features) + fmt cleanChecklist
token-rate-limit-filterfeature gate (experimental)backend.rs,ledger.rs,token_bucket_ledger.rsuntouched)