From 89e1acfd10606e912ce55264c30ab7695941acbe Mon Sep 17 00:00:00 2001 From: Jordi Gil Date: Fri, 28 Aug 2026 21:45:48 -0400 Subject: [PATCH 1/5] feat(token_rate_limit): add M1/M2/M6 core filter (experimental) Implements the agreed M1/M2/M6 core of the token rate limiting proposal (00121_token-rate-limiting.md in praxis-proxy/enhancements, tracked by epic ai#121): sliding-window and token-bucket algorithms, per-rule request/response matching, and a pluggable backend (in-memory by default, Valkey for shared state across replicas). Gated behind the token-rate-limit-filter Cargo feature (disabled by default, mirroring the azure-ad/gcp-adc/http-callout pattern) so nothing changes for consumers who don't opt in while the remaining epic milestones (M3+) are still being designed. Signed-off-by: Jordi Gil --- Cargo.lock | 176 +- Cargo.toml | 1 + filters/Cargo.toml | 11 + filters/src/lib.rs | 4 + filters/src/register.rs | 15 +- filters/src/token_rate_limit/backend.rs | 2028 +++++++++++++++++ filters/src/token_rate_limit/config.rs | 375 +++ filters/src/token_rate_limit/ledger.rs | 767 +++++++ filters/src/token_rate_limit/mod.rs | 952 ++++++++ filters/src/token_rate_limit/tests.rs | 1157 ++++++++++ .../token_rate_limit/token_bucket_ledger.rs | 779 +++++++ filters/src/token_usage/mod.rs | 12 + server/Cargo.toml | 2 + tests/integration/Cargo.toml | 3 + 14 files changed, 6231 insertions(+), 51 deletions(-) create mode 100644 filters/src/token_rate_limit/backend.rs create mode 100644 filters/src/token_rate_limit/config.rs create mode 100644 filters/src/token_rate_limit/ledger.rs create mode 100644 filters/src/token_rate_limit/mod.rs create mode 100644 filters/src/token_rate_limit/tests.rs create mode 100644 filters/src/token_rate_limit/token_bucket_ledger.rs diff --git a/Cargo.lock b/Cargo.lock index 97014bbaf5..eabe1b438d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3,15 +3,10 @@ version = 4 [[package]] -name = "ahash" -version = "0.7.8" +name = "adler2" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "891477e0c6a8957309ee5c45a6368af3ae14bb510732d2684ffa19af310920f9" -dependencies = [ - "getrandom 0.2.17", - "once_cell", - "version_check", -] +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" [[package]] name = "ahash" @@ -108,7 +103,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -119,7 +114,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -137,6 +132,12 @@ dependencies = [ "rustversion", ] +[[package]] +name = "arcstr" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d" + [[package]] name = "arrayvec" version = "0.7.8" @@ -220,6 +221,17 @@ dependencies = [ "serde_json", ] +[[package]] +name = "async-lock" +version = "3.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "290f7f2596bd5b78a9fec8088ccd89180d7f9f55b94b0576823bbbdc72ee8311" +dependencies = [ + "event-listener", + "event-listener-strategy", + "pin-project-lite", +] + [[package]] name = "async-trait" version = "0.1.92" @@ -643,12 +655,12 @@ checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" [[package]] name = "chacha20" -version = "0.10.2" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" dependencies = [ "cfg-if", - "cpufeatures 0.3.1", + "cpufeatures 0.3.0", "rand_core 0.10.1", ] @@ -727,12 +739,16 @@ checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" [[package]] name = "combine" -version = "4.6.8" +version = "4.6.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e" +checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" dependencies = [ "bytes", + "futures-core", "memchr", + "pin-project-lite", + "tokio", + "tokio-util", ] [[package]] @@ -798,9 +814,9 @@ dependencies = [ [[package]] name = "cpufeatures" -version = "0.3.1" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ca28b0ae3115b884660db4118d803791fd6756b6e88f39c0f3f7859060d7566" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" dependencies = [ "libc", ] @@ -822,9 +838,9 @@ checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" [[package]] name = "crc32fast" -version = "1.5.1" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8498c871161e1742aaa9d52551b2d6ebdd4c3d45a3be423e3728f33b955be550" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" dependencies = [ "cfg-if", ] @@ -1120,7 +1136,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -1143,6 +1159,16 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener", + "pin-project-lite", +] + [[package]] name = "evmap" version = "11.0.0" @@ -1185,12 +1211,13 @@ checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" [[package]] name = "flate2" -version = "1.1.10" +version = "1.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" dependencies = [ "crc32fast", "libz-ng-sys", + "miniz_oxide", ] [[package]] @@ -1210,7 +1237,7 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf5efcf77a4da27927d3ab0509dec5b0954bb3bc59da5a1de9e52642ebd4cdf9" dependencies = [ - "ahash 0.8.12", + "ahash", "num_cpus", "parking_lot", "seize", @@ -1451,9 +1478,6 @@ name = "hashbrown" version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" -dependencies = [ - "ahash 0.7.8", -] [[package]] name = "hashbrown" @@ -1486,6 +1510,11 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", +] [[package]] name = "hashlink" @@ -2278,6 +2307,16 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.2.2" @@ -2358,7 +2397,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -2688,6 +2727,7 @@ dependencies = [ "praxis-proxy-core", "praxis-proxy-filter", "quixotic-plecostomus-core", + "redis", "reqwest", "serde", "serde_json", @@ -2695,6 +2735,7 @@ dependencies = [ "serde_json_path", "sha2 0.11.0", "tempfile", + "thiserror 2.0.20", "tokio", "tokio-util", "tracing", @@ -3006,7 +3047,7 @@ version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03da047801ff44bb6a4d407d4860c05fd70bb81714e6b2f3812603d5b145b042" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "itertools", "log", "multimap", @@ -3076,7 +3117,7 @@ version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cb5cd191fad9142d096efc97c64715db73db4adb60c38ebdf22f5b54cd6ea712" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "prost", "prost-build", "prost-types", @@ -3218,7 +3259,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3227,7 +3268,7 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a820a1a3e59644b38bc03606b638b41f1c5a474215efcec37f94b15871a7e1c" dependencies = [ - "ahash 0.8.12", + "ahash", "async-trait", "blake2", "bstr", @@ -3247,7 +3288,7 @@ dependencies = [ "quixotic-plecostomus-http", "quixotic-plecostomus-lru", "quixotic-plecostomus-timeout", - "rand 0.8.8", + "rand 0.8.7", "regex", "rmp", "rmp-serde", @@ -3262,7 +3303,7 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d7d6d9e04414ba6819c51ec210eddf17865c505431a492c86c6356338c2b4f60" dependencies = [ - "ahash 0.8.12", + "ahash", "async-trait", "brotli", "bstr", @@ -3293,7 +3334,7 @@ dependencies = [ "quixotic-plecostomus-runtime", "quixotic-plecostomus-rustls", "quixotic-plecostomus-timeout", - "rand 0.8.8", + "rand 0.8.7", "regex", "serde", "serde_yaml", @@ -3350,9 +3391,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec55d3bcbd56e52b3424bc62739a237d43c14df8ed04aaf294ae187b60da348a" dependencies = [ "arrayvec", - "hashbrown 0.12.3", + "hashbrown 0.17.1", "parking_lot", - "rand 0.8.8", + "rand 0.8.7", ] [[package]] @@ -3389,7 +3430,7 @@ dependencies = [ "quixotic-plecostomus-core", "quixotic-plecostomus-error", "quixotic-plecostomus-http", - "rand 0.8.8", + "rand 0.8.7", "regex", "tokio", ] @@ -3402,7 +3443,7 @@ checksum = "51621ff504cfde18409889d6d5b8773e0b0de0e4defa610006a36c27adc06a10" dependencies = [ "log", "once_cell", - "rand 0.8.8", + "rand 0.8.7", "serde", "thread_local", "tokio", @@ -3461,9 +3502,9 @@ checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" [[package]] name = "rand" -version = "0.8.8" +version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e058c7de0b26af77780c769414d6257830bb240f3c38477dbc2c16e5f54d6d4c" +checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" dependencies = [ "libc", "rand_chacha 0.3.1", @@ -3585,6 +3626,29 @@ dependencies = [ "yasna", ] +[[package]] +name = "redis" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e37a4ca5c6ca42aa3e6df2fd32b987a65d32a4c2159a6f3fe0fd1df306a2658f" +dependencies = [ + "arcstr", + "async-lock", + "bytes", + "cfg-if", + "combine", + "futures-util", + "itoa", + "percent-encoding", + "pin-project-lite", + "ryu", + "socket2", + "tokio", + "tokio-util", + "url", + "xxhash-rust", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -3807,7 +3871,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3888,7 +3952,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4192,7 +4256,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214" dependencies = [ "cfg-if", - "cpufeatures 0.3.1", + "cpufeatures 0.3.0", "digest 0.11.3", ] @@ -4214,7 +4278,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" dependencies = [ "cfg-if", - "cpufeatures 0.3.1", + "cpufeatures 0.3.0", "digest 0.11.3", ] @@ -4243,6 +4307,12 @@ dependencies = [ "libc", ] +[[package]] +name = "simd-adler32" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" + [[package]] name = "simd_cesu8" version = "1.2.0" @@ -4287,7 +4357,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -4573,10 +4643,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -5162,9 +5232,9 @@ dependencies = [ [[package]] name = "uuid" -version = "1.26.0" +version = "1.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5772d71c9be8a8a6ac2117d949c5b224c1b72241bb611d9a3012edcf8af7812" +checksum = "f053576934f05a761a402421fbbe3d425d9366f75f978806a037b3ca481abecc" dependencies = [ "getrandom 0.4.3", "js-sys", @@ -5363,7 +5433,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -5688,6 +5758,12 @@ dependencies = [ "yaml_serde", ] +[[package]] +name = "xxhash-rust" +version = "0.8.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6" + [[package]] name = "yaml_serde" version = "0.10.7" diff --git a/Cargo.toml b/Cargo.toml index cd4dbf106d..9273ca9034 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,6 +43,7 @@ notify = "8.2.0" quote = "1.0.47" rand = "0.10.2" rcgen = "0.14.9" +redis = { version = "1.5.0", default-features = false, features = ["tokio-comp"] } regex = "1.13.1" reqwest = { version = "0.13.4", default-features = false, features = ["rustls", "json", "stream"] } rmcp = { version = "3.1.4", default-features = false, features = ["client", "transport-streamable-http-client-reqwest", "reqwest"] } diff --git a/filters/Cargo.toml b/filters/Cargo.toml index 29608d0a3a..9be2fa9879 100644 --- a/filters/Cargo.toml +++ b/filters/Cargo.toml @@ -25,6 +25,15 @@ azure-ad-filter = ["experimental"] # authentication filter. Work in progress; the config surface may change. # Activates the `experimental` marker. gcp-adc-filter = ["experimental"] +# Experimental: the token_rate_limit filter. Its parent proposal +# (00121_token-rate-limiting in praxis-proxy/enhancements) is not yet +# `accepted`, and open design questions remain (bucket-keying strategy, +# HA/clustered-Valkey failure modes, how this relates to Kuadrant's +# TokenRateLimitPolicy) -- see ai#796. Activates the `experimental` marker. +# Also pulls in `redis` and `thiserror`, which only this filter uses -- +# optional so the rest of the crate (and its consumers) don't compile the +# Valkey client and its transitive deps unless this feature is on. +token-rate-limit-filter = ["experimental", "dep:redis", "dep:thiserror"] # Marker feature activated transitively by any experimental feature, mirroring # praxis core's server crate. Lets consumers gate on "anything experimental". experimental = [] @@ -51,6 +60,7 @@ pingora-core.workspace = true praxis-ai-apis = { workspace = true } praxis-core = { workspace = true } praxis-filter = { workspace = true } +redis = { workspace = true, optional = true } reqwest.workspace = true serde = { workspace = true } serde_json = { workspace = true } @@ -58,6 +68,7 @@ serde_json_canonicalizer = { workspace = true } serde_json_path = { workspace = true } serde_yaml = { workspace = true } sha2 = { workspace = true } +thiserror = { workspace = true, optional = true } tokio = { workspace = true, features = ["rt", "sync", "time"] } tokio-util = { workspace = true } tracing = { workspace = true } diff --git a/filters/src/lib.rs b/filters/src/lib.rs index 0813ee463c..ae98cd3e05 100644 --- a/filters/src/lib.rs +++ b/filters/src/lib.rs @@ -24,6 +24,8 @@ pub mod prompt_enrich; mod register; pub mod routing; mod time_to_first_token; +#[cfg(feature = "token-rate-limit-filter")] +mod token_rate_limit; mod token_usage; pub use agentic::{a2a::A2aFilter, mcp::McpFilter}; @@ -40,6 +42,8 @@ pub use prompt_enrich::PromptEnrichFilter; pub use register::{build_ai_registry, register_ai_filters}; pub use routing::{CredentialInjectFilter, IntelligentRouteFilter, ProviderRouteFilter}; pub use time_to_first_token::TimeToFirstTokenFilter; +#[cfg(feature = "token-rate-limit-filter")] +pub use token_rate_limit::TokenRateLimitFilter; pub use token_usage::{TokenCountFilter, TokenUsageHeadersFilter}; // ----------------------------------------------------------------------------- diff --git a/filters/src/register.rs b/filters/src/register.rs index e38577da50..ceb8f0bdf6 100644 --- a/filters/src/register.rs +++ b/filters/src/register.rs @@ -12,6 +12,8 @@ use crate::AzureAdFilter; use crate::GcpAdcFilter; #[cfg(feature = "http-callout-filter")] use crate::HttpCalloutFilter; +#[cfg(feature = "token-rate-limit-filter")] +use crate::TokenRateLimitFilter; use crate::{ A2aFilter, AiGuardrailsFilter, CredentialInjectFilter, IntelligentRouteFilter, McpFilter, ModelToHeaderFilter, PromptEnrichFilter, ProviderRouteFilter, Sigv4SignFilter, TimeToFirstTokenFilter, TokenCountFilter, @@ -115,6 +117,15 @@ fn register_general_ai_filters(registry: &mut FilterRegistry) { @register registry, http "prompt_enrich" => PromptEnrichFilter::from_config ); + praxis_filter::register_filters!( + @register registry, + http "time_to_first_token" => TimeToFirstTokenFilter::from_config + ); + register_token_filters(registry); +} + +/// Register token counting/usage/rate-limiting filters. +fn register_token_filters(registry: &mut FilterRegistry) { praxis_filter::register_filters!( @register registry, http "token_count" => TokenCountFilter::from_config @@ -123,9 +134,10 @@ fn register_general_ai_filters(registry: &mut FilterRegistry) { @register registry, http "token_usage_headers" => TokenUsageHeadersFilter::from_config ); + #[cfg(feature = "token-rate-limit-filter")] praxis_filter::register_filters!( @register registry, - http "time_to_first_token" => TimeToFirstTokenFilter::from_config + http "token_rate_limit" => TokenRateLimitFilter::from_config ); } @@ -431,6 +443,7 @@ mod tests { assert_experimental_registration(&names, "http_callout", cfg!(feature = "http-callout-filter")); assert_experimental_registration(&names, "azure_ad", cfg!(feature = "azure-ad-filter")); assert_experimental_registration(&names, "gcp_adc", cfg!(feature = "gcp-adc-filter")); + assert_experimental_registration(&names, "token_rate_limit", cfg!(feature = "token-rate-limit-filter")); } #[test] diff --git a/filters/src/token_rate_limit/backend.rs b/filters/src/token_rate_limit/backend.rs new file mode 100644 index 0000000000..b26cd445af --- /dev/null +++ b/filters/src/token_rate_limit/backend.rs @@ -0,0 +1,2028 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Praxis Contributors + +//! Pluggable token-rate-limit state backends. +//! +//! Adapted, unmodified in logic, from the `token_rate_limit::backend` module +//! on nerdalert's `poc/distributed-token-rate-limit-demo` spike branch +//! (). +//! `reserve`/`reconcile` are key-agnostic (`ReserveRequest`/`ReconcileRequest` +//! carry a plain `String` key); this filter supplies a single fixed key +//! (see `FALLBACK_KEY`) instead of the source branch's principal+model +//! composite key, so no logic here was changed to adopt it. + +use std::{ + sync::{Arc, Mutex, OnceLock}, + time::Duration, +}; + +use async_trait::async_trait; +use metrics::counter; +use redis::aio::MultiplexedConnection; +use sha2::{Digest as _, Sha256}; +use tokio::sync::mpsc; + +use super::{ + ledger::{Budget, Decision, Ledger, Settlement}, + token_bucket_ledger::{self, TokenBucketLedger}, +}; + +/// Bound on every Valkey network operation (connect or `EVAL`), so an +/// unreachable-but-not-yet-timed-out-at-the-OS-level backend still fails +/// closed quickly rather than hanging the request indefinitely. +const VALKEY_TIMEOUT: Duration = Duration::from_millis(500); + +/// Request to admit an estimated token cost against a key's budget. +#[derive(Debug, Clone)] +pub(super) struct ReserveRequest { + /// Opaque budget key (this milestone always uses `FALLBACK_KEY`). + pub(super) key: String, + /// Estimated token cost to reserve if admitted. + pub(super) estimate: u64, + /// Caller's current time, in milliseconds. + pub(super) now_ms: u64, +} + +/// Request to settle a prior reservation against actual usage. +#[derive(Debug, Clone)] +pub(super) struct ReconcileRequest { + /// Same key the original [`ReserveRequest`] used. + pub(super) key: String, + /// Reservation ID returned by [`BackendReserve::Admitted`]. + pub(super) reservation_id: u64, + /// Actual token usage, if known; `None` charges at `estimate`. + pub(super) actual: Option, + /// The original reservation's estimate (for backends, like Valkey, + /// that reconcile out-of-band and need it for a default charge). + pub(super) estimate: u64, + /// Caller's current time, in milliseconds. + pub(super) now_ms: u64, +} + +/// Result of a [`TokenRateLimitStateBackend::reserve`] call. +#[derive(Debug, Clone)] +pub(super) enum BackendReserve { + /// Request may proceed with this reservation. + Admitted { + /// Opaque ID used for later reconciliation. + reservation_id: u64, + /// Estimate actually reserved. + estimate: u64, + }, + /// Request must be rejected before routing. + Denied { + /// Conservative delay before another admission attempt. + retry_after_ms: u64, + }, +} + +/// Result of a [`TokenRateLimitStateBackend::reconcile`] call. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum BackendSettlement { + /// Actual usage was applied exactly once. + Applied { + /// Actual tokens charged. + actual: u64, + /// Estimate returned to the budget. + refund: u64, + /// Usage above the estimate. + overage: u64, + }, + /// The reservation was already reconciled or conservatively expired. + Noop, +} + +/// Failure modes shared by every [`TokenRateLimitStateBackend`] impl. +#[derive(Debug, thiserror::Error)] +pub(super) enum BackendError { + /// The backend could not be reached or timed out. + #[error("shared quota backend unavailable: {0}")] + Unavailable(String), + /// The backend responded, but not in the expected shape. + #[error("shared quota backend returned an invalid response")] + InvalidResponse, +} + +/// Where sliding-window admission state lives: in-process or shared. +#[async_trait] +pub(super) trait TokenRateLimitStateBackend: Send + Sync { + /// Attempt to admit `request.estimate` against `request.key`'s budget. + async fn reserve(&self, request: ReserveRequest) -> Result; + + /// Settle a prior reservation against actual usage, awaiting + /// completion. Backends that reconcile out-of-band (e.g. Valkey via + /// [`Self::enqueue_reconcile`]) still implement this for their own + /// background worker to call. + async fn reconcile(&self, request: ReconcileRequest) -> Result; + + /// Settle a prior reservation without blocking the caller. + /// + /// For in-process state this may just reconcile synchronously (cheap, + /// no I/O); for a networked backend this enqueues the work onto a + /// background worker instead, so the response is never held up on a + /// reconciliation round-trip. + fn enqueue_reconcile(&self, request: ReconcileRequest) -> Result<(), BackendError>; + + /// The smallest configured budget capacity, for rate-limit headers. + fn limit(&self) -> u64; + + /// Attempt an in-process, synchronous settlement (no I/O, no async + /// dispatch) for a prior reservation. + /// + /// Returns `None` for backends whose state isn't local (e.g. a + /// networked Valkey backend) -- callers should fall back to + /// [`Self::enqueue_reconcile`] in that case. Every in-process backend + /// (regardless of algorithm) implements this itself rather than + /// exposing its concrete state type, so the filter never needs to + /// know which algorithm produced it. + fn reconcile_sync(&self, _request: &ReconcileRequest) -> Option { + None + } + + /// Reclaim idle/orphaned in-process state and report current gauges. + /// + /// Returns `None` for backends with no local state to reap (e.g. + /// Valkey, where expiry is handled by the Lua reserve script + /// itself) -- callers should skip gauge reporting entirely in that + /// case rather than reporting misleading zeros. + fn cleanup(&self, _now_ms: u64, _max_keys_to_scan: usize) -> Option { + None + } +} + +/// In-process state snapshot after a [`TokenRateLimitStateBackend::cleanup`] +/// pass, backend-agnostic so the filter can report gauges without knowing +/// which algorithm produced them. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub(super) struct CleanupReport { + /// Reservations reaped this pass because they exceeded + /// `reservation_timeout` without being reconciled. + pub(super) orphaned: usize, + /// Reservations still awaiting reconciliation. + pub(super) active_reservations: usize, + /// Distinct budget keys currently retained. + pub(super) active_keys: usize, +} + +/// In-process sliding-window state: one gateway instance, one budget. +pub(super) struct InMemoryTokenRateLimitBackend { + /// The underlying exact sliding-window ledger. + ledger: Arc, +} + +impl InMemoryTokenRateLimitBackend { + /// Wrap an already-constructed [`Ledger`] as a backend. + pub(super) fn new(ledger: Ledger) -> Self { + Self { + ledger: Arc::new(ledger), + } + } +} + +#[async_trait] +impl TokenRateLimitStateBackend for InMemoryTokenRateLimitBackend { + async fn reserve(&self, request: ReserveRequest) -> Result { + Ok( + match self.ledger.reserve(&request.key, request.estimate, request.now_ms) { + Decision::Admitted(reservation) => BackendReserve::Admitted { + reservation_id: reservation.id, + estimate: reservation.estimate, + }, + Decision::Denied { retry_after_ms, .. } => BackendReserve::Denied { retry_after_ms }, + }, + ) + } + + async fn reconcile(&self, request: ReconcileRequest) -> Result { + Ok( + match self + .ledger + .reconcile(request.reservation_id, request.actual, request.now_ms) + { + Settlement::Applied { + actual, + refund, + overage, + } => BackendSettlement::Applied { + actual, + refund, + overage, + }, + Settlement::Noop => BackendSettlement::Noop, + }, + ) + } + + fn enqueue_reconcile(&self, request: ReconcileRequest) -> Result<(), BackendError> { + let _ = self + .ledger + .reconcile(request.reservation_id, request.actual, request.now_ms); + Ok(()) + } + + fn limit(&self) -> u64 { + self.ledger.limit() + } + + fn reconcile_sync(&self, request: &ReconcileRequest) -> Option { + Some( + match self + .ledger + .reconcile(request.reservation_id, request.actual, request.now_ms) + { + Settlement::Applied { + actual, + refund, + overage, + } => BackendSettlement::Applied { + actual, + refund, + overage, + }, + Settlement::Noop => BackendSettlement::Noop, + }, + ) + } + + fn cleanup(&self, now_ms: u64, max_keys_to_scan: usize) -> Option { + Some(CleanupReport { + orphaned: self.ledger.cleanup(now_ms, max_keys_to_scan), + active_reservations: self.ledger.active_count(), + active_keys: self.ledger.key_count(), + }) + } +} + +/// In-process token-bucket state: one gateway instance, one budget, +/// continuously refilled rather than admitted against a trailing window. +pub(super) struct InMemoryTokenBucketBackend { + /// The underlying exact token-bucket ledger. + ledger: Arc, +} + +impl InMemoryTokenBucketBackend { + /// Wrap an already-constructed [`TokenBucketLedger`] as a backend. + pub(super) fn new(ledger: TokenBucketLedger) -> Self { + Self { + ledger: Arc::new(ledger), + } + } + + /// Shared reconcile path for `reconcile`/`enqueue_reconcile`/`reconcile_sync`. + fn reconcile_ledger(&self, request: &ReconcileRequest) -> BackendSettlement { + match self + .ledger + .reconcile(request.reservation_id, request.actual, request.now_ms) + { + token_bucket_ledger::Settlement::Applied { + actual, + refund, + overage, + } => BackendSettlement::Applied { + actual, + refund, + overage, + }, + token_bucket_ledger::Settlement::Noop => BackendSettlement::Noop, + } + } +} + +#[async_trait] +impl TokenRateLimitStateBackend for InMemoryTokenBucketBackend { + async fn reserve(&self, request: ReserveRequest) -> Result { + Ok( + match self.ledger.reserve(&request.key, request.estimate, request.now_ms) { + token_bucket_ledger::Decision::Admitted(reservation) => BackendReserve::Admitted { + reservation_id: reservation.id, + estimate: reservation.estimate, + }, + token_bucket_ledger::Decision::Denied { retry_after_ms } => BackendReserve::Denied { retry_after_ms }, + }, + ) + } + + async fn reconcile(&self, request: ReconcileRequest) -> Result { + Ok(self.reconcile_ledger(&request)) + } + + fn enqueue_reconcile(&self, request: ReconcileRequest) -> Result<(), BackendError> { + let _ = self.reconcile_ledger(&request); + Ok(()) + } + + fn limit(&self) -> u64 { + self.ledger.limit() + } + + fn reconcile_sync(&self, request: &ReconcileRequest) -> Option { + Some(self.reconcile_ledger(request)) + } + + fn cleanup(&self, now_ms: u64, max_keys_to_scan: usize) -> Option { + Some(CleanupReport { + orphaned: self.ledger.cleanup(now_ms, max_keys_to_scan), + active_reservations: self.ledger.active_count(), + active_keys: self.ledger.key_count(), + }) + } +} + +/// Atomically admit a reservation against every configured budget for one +/// key, or deny it -- the Valkey/Lua analog of [`Ledger::reserve`]. +/// +/// `KEYS`: `[1]` physical key, `[2]` settled zset, `[3]` active hash, +/// `[4]` namespace keys zset, `[5]` namespace active-count string, +/// `[6]` namespace reservation-id sequence, `[7]` namespace active-index +/// zset (global reservation-expiry tracking). `ARGV`: reservation +/// timeout (ms), max keys, max active reservations, estimate, budget +/// count, then `(window_ms, capacity)` pairs. Returns +/// `[1, id, estimate]` on admission or `[0, retry_after_ms]` on denial. +const RESERVE_SCRIPT: &str = " +local now = redis.call('TIME') +local now_ms = tonumber(now[1]) * 1000 + math.floor(tonumber(now[2]) / 1000) +local timeout_ms = tonumber(ARGV[1]) +local max_keys = tonumber(ARGV[2]) +local max_active = tonumber(ARGV[3]) +local estimate = tonumber(ARGV[4]) +local budget_count = tonumber(ARGV[5]) +local settled = KEYS[2] +local active = KEYS[3] + +local active_total = tonumber(redis.call('GET', KEYS[5]) or '0') +local expired_global = redis.call('ZRANGE', KEYS[7], '-inf', now_ms, 'BYSCORE') +for i = 1, #expired_global do + local member = expired_global[i] + local split = string.find(member, '|') + if split then + local physical = string.sub(member, 1, split - 1) + local reservation = string.sub(member, split + 1) + local active_key = physical .. ':active' + local value = redis.call('HGET', active_key, reservation) + if value then + local value_split = string.find(value, '|') + local amount = tonumber(string.sub(value, 1, value_split - 1)) + local reserved_at = tonumber(string.sub(value, value_split + 1)) + redis.call('ZADD', physical .. ':settled', reserved_at, 'expired:' .. reservation .. ':' .. amount) + redis.call('HDEL', active_key, reservation) + active_total = math.max(0, active_total - 1) + end + end + redis.call('ZREM', KEYS[7], member) +end +redis.call('SET', KEYS[5], active_total) + +local max_window = 0 +for i = 1, budget_count do + local window = tonumber(ARGV[5 + (i * 2) - 1]) + if window > max_window then max_window = window end + redis.call('ZREMRANGEBYSCORE', settled, '-inf', now_ms - window) +end + +local expired = {} +local active_values = redis.call('HGETALL', active) +for i = 1, #active_values, 2 do + local id = active_values[i] + local value = active_values[i + 1] + local sep = string.find(value, '|') + local reserved_at = tonumber(string.sub(value, sep + 1)) + if now_ms - reserved_at >= timeout_ms then + local amount = tonumber(string.sub(value, 1, sep - 1)) + redis.call('ZADD', settled, reserved_at, 'expired:' .. id .. ':' .. amount) + redis.call('HDEL', active, id) + active_total = math.max(0, active_total - 1) + end +end +redis.call('SET', KEYS[5], active_total) + +redis.call('ZREMRANGEBYSCORE', KEYS[4], '-inf', now_ms) +local key_exists = redis.call('ZSCORE', KEYS[4], KEYS[1]) ~= false +if not key_exists and redis.call('ZCARD', KEYS[4]) >= max_keys then + return {0, max_window} +end +if active_total >= max_active then + return {0, max_window} +end + +for i = 1, budget_count do + local window = tonumber(ARGV[5 + (i * 2) - 1]) + local capacity = tonumber(ARGV[5 + (i * 2)]) + local settled_sum = 0 + local entries = redis.call('ZRANGE', settled, now_ms - window, '+inf', 'BYSCORE', 'WITHSCORES') + for j = 1, #entries, 2 do + local member = entries[j] + local amount = string.match(member, ':(%d+)$') + if amount then settled_sum = settled_sum + tonumber(amount) end + end + local active_values = redis.call('HGETALL', active) + local active_sum = 0 + for j = 1, #active_values, 2 do + local sep = string.find(active_values[j + 1], '|') + active_sum = active_sum + tonumber(string.sub(active_values[j + 1], 1, sep - 1)) + end + if settled_sum + active_sum + estimate > capacity then + return {0, max_window} + end +end + +local id = redis.call('INCR', KEYS[6]) +redis.call('HSET', active, id, estimate .. '|' .. now_ms) +redis.call('INCR', KEYS[5]) +redis.call('ZADD', KEYS[7], now_ms + timeout_ms, KEYS[1] .. '|' .. id) +local ttl = math.max(max_window + timeout_ms, 1000) +redis.call('ZADD', KEYS[4], now_ms + ttl, KEYS[1]) +redis.call('PEXPIRE', settled, ttl) +redis.call('PEXPIRE', active, ttl) +redis.call('PEXPIRE', KEYS[1], ttl) +return {1, id, estimate} +"; + +/// Atomically settle a prior reservation against actual usage -- the +/// Valkey/Lua analog of [`Ledger::reconcile`]. +/// +/// `KEYS`: same layout as [`RESERVE_SCRIPT`]. `ARGV`: `[1]` reservation +/// ID, `[2]` actual usage. Returns `[0]` if the reservation was already +/// reconciled/expired (no-op), or `[1, actual, refund, overage]`. +const RECONCILE_SCRIPT: &str = " +local value = redis.call('HGET', KEYS[3], ARGV[1]) +if not value then return {0} end +local sep = string.find(value, '|') +local estimate = tonumber(string.sub(value, 1, sep - 1)) +local actual = tonumber(ARGV[2]) +redis.call('HDEL', KEYS[3], ARGV[1]) +local active_total = math.max(0, tonumber(redis.call('GET', KEYS[5]) or '0') - 1) +redis.call('SET', KEYS[5], active_total) +redis.call('ZREM', KEYS[7], KEYS[1] .. '|' .. ARGV[1]) +local now = redis.call('TIME') +local now_ms = tonumber(now[1]) * 1000 + math.floor(tonumber(now[2]) / 1000) +redis.call('ZADD', KEYS[2], now_ms, 'settled:' .. ARGV[1] .. ':' .. actual) +return {1, actual, math.max(0, estimate - actual), math.max(0, actual - estimate)} +"; + +/// Drain `receiver`, reconciling each request against `worker`'s backend +/// with bounded retries, off the request/response path entirely. +/// +/// Generic over any [`TokenRateLimitStateBackend`] (sliding-window, +/// token-bucket, or any future Valkey-backed algorithm) -- the retry/ +/// audit behavior is identical regardless of which algorithm's Lua +/// script `worker.reconcile` ultimately calls. +/// +/// A dropped/failed reconciliation after retries is intentionally *not* +/// escalated back to the request that triggered it (that response has +/// already been sent) -- it's counted and logged so operators can audit +/// it, and the reservation still expires and gets conservatively charged +/// via `reservation_timeout` regardless. +async fn run_reconcile_worker(worker: B, mut receiver: mpsc::Receiver) +where + B: TokenRateLimitStateBackend + 'static, +{ + while let Some(request) = receiver.recv().await { + let mut attempts = 0; + loop { + match worker.reconcile(request.clone()).await { + Ok(_) => { + counter!("praxis_ai_token_rate_limit_backend_reconciliation_total", "backend" => "valkey", "result" => "completed") + .increment(1); + break; + }, + Err(error) if attempts < 2 => { + attempts += 1; + tracing::warn!(attempts, %error, "token-rate-limit reconciliation retry"); + tokio::time::sleep(Duration::from_millis(25 * attempts)).await; + }, + Err(error) => { + counter!("praxis_ai_token_rate_limit_backend_errors_total", "backend" => "valkey", "operation" => "reconcile") + .increment(1); + tracing::error!(%error, "token-rate-limit reconciliation abandoned after retries"); + break; + }, + } + } + } +} + +/// Shared Valkey connection handling for every Valkey-backed algorithm: +/// reusing one cached multiplexed connection across calls and running +/// `EVAL`s against it, both bounded by [`VALKEY_TIMEOUT`] -- enforced by +/// the `redis` crate itself (see [`Self::connection`]), not by wrapping +/// calls in our own `tokio::time::timeout` -- so an unreachable/wedged +/// backend fails closed quickly instead of hanging the request. +/// +/// Built once per filter instance (not per rule) and `Clone`d into every +/// Valkey-backed rule's [`ValkeyBackendConfig`]/[`ValkeyTokenBucketConfig`] +/// -- cloning is cheap (`redis::Client` is a plain [`Clone`] wrapper +/// around connection info, and `connection` below is `Arc`-shared), so +/// every rule ends up sharing the same one multiplexed connection +/// instead of opening a redundant one per rule pointed at the same URL. +#[derive(Clone)] +pub(super) struct ValkeyEval { + /// Lazy Valkey/Redis client, used only to (re-)establish + /// `connection` below. + client: redis::Client, + + /// Cached multiplexed connection, established on first use and + /// reused by every subsequent call (a `MultiplexedConnection` is a + /// cheap-to-clone handle onto one shared pipelined connection, not + /// a dedicated socket per clone) -- paying a fresh TCP/TLS handshake + /// on every request would defeat the point of a "multiplexed" + /// connection and add unnecessary latency to the request path. + /// Cleared by [`Self::invalidate`] after a failed command, so the + /// next call re-establishes it rather than reusing a + /// wedged/reset one indefinitely. + connection: Arc>>, +} + +impl ValkeyEval { + /// Open a (lazy, not-yet-connected) Valkey client. + /// + /// # Errors + /// + /// Returns [`BackendError::Unavailable`] if `url` isn't a well-formed + /// Valkey/Redis connection URL. + pub(super) fn new(url: String) -> Result { + let client = redis::Client::open(url).map_err(|e| BackendError::Unavailable(e.to_string()))?; + Ok(Self { + client, + connection: Arc::new(tokio::sync::Mutex::new(None)), + }) + } + + /// Return the cached multiplexed connection, establishing (and + /// caching) a fresh one on first use or after a prior failure + /// invalidated it. The `redis` crate bounds both the connection + /// attempt itself and every command later sent over it to + /// [`VALKEY_TIMEOUT`] (via [`Self::connection_config`]), so an + /// unreachable or wedged Valkey fails closed rather than hanging + /// the request indefinitely -- we don't additionally wrap this in + /// our own `tokio::time::timeout`, which would just race the + /// crate's own enforcement of the same bound. + async fn connection(&self) -> Result { + let mut cached = self.connection.lock().await; + if let Some(connection) = cached.as_ref() { + let connection = connection.clone(); + drop(cached); + return Ok(connection); + } + let connection = self + .client + .get_multiplexed_async_connection_with_config(&Self::connection_config()) + .await + .map_err(|error| map_valkey_error("connection", &error))?; + *cached = Some(connection.clone()); + drop(cached); + Ok(connection) + } + + /// [`redis::AsyncConnectionConfig`] binding both the connection + /// attempt and every command's response to [`VALKEY_TIMEOUT`]. + /// Without this, `redis` still applies its own defaults (500ms + /// response, 1s connect, as of `redis` 1.6) -- close to, but not + /// exactly, this crate's own documented bound, and liable to drift + /// further from it silently on a future `redis` upgrade. + fn connection_config() -> redis::AsyncConnectionConfig { + redis::AsyncConnectionConfig::new() + .set_connection_timeout(Some(VALKEY_TIMEOUT)) + .set_response_timeout(Some(VALKEY_TIMEOUT)) + } + + /// Drop the cached connection so the next call re-establishes it. + async fn invalidate(&self) { + *self.connection.lock().await = None; + } + + /// Run one `EVAL script KEYS... ARGV...` command against the cached + /// connection (see [`Self::connection`]), invalidating it on any + /// failure -- including a [`VALKEY_TIMEOUT`] response timeout, + /// enforced by `redis` itself, see [`Self::connection_config`] -- + /// so a subsequent call doesn't keep retrying a wedged/reset one. + async fn eval( + &self, + script: &str, + keys: &[String; N], + args: &[String], + ) -> Result, BackendError> { + let mut command = redis::cmd("EVAL"); + command.arg(script).arg(keys.len()); + for key in keys { + command.arg(key); + } + for arg in args { + command.arg(arg); + } + let mut connection = self.connection().await?; + let result: redis::RedisResult> = command.query_async(&mut connection).await; + match result { + Ok(value) => Ok(value), + Err(error) => { + self.invalidate().await; + Err(map_valkey_error("command", &error)) + }, + } + } +} + +/// Wrap a [`redis::RedisError`] as a [`BackendError::Unavailable`], +/// tagged with which phase (`"connection"` or `"command"`) it came +/// from -- `redis`'s own error text (e.g. plain `"timed out"` for a +/// response-timeout) doesn't say which on its own. +fn map_valkey_error(phase: &'static str, error: &redis::RedisError) -> BackendError { + BackendError::Unavailable(format!("Valkey {phase}: {error}")) +} + +/// Shared background-reconciliation scaffolding for every Valkey-backed +/// algorithm: a queue plus a spawn-at-most-once guard for +/// [`run_reconcile_worker`]. +struct ReconcileWorker { + /// Sending half of the reconciliation queue; cloned into the worker. + tx: mpsc::Sender, + /// Receiving half, taken exactly once by [`Self::start`]. + rx: Mutex>>, + /// Ensures the background worker is spawned at most once. + started: OnceLock<()>, +} + +impl ReconcileWorker { + /// A live worker: holds a real receiver, ready for [`Self::start`]. + fn new() -> Self { + let (tx, rx) = mpsc::channel(1024); + Self { + tx, + rx: Mutex::new(Some(rx)), + started: OnceLock::new(), + } + } + + /// A throwaway (never-sent-to, never-started) worker -- used only + /// when cloning a backend to hand the *real* background worker its + /// own handle to `reserve`/`reconcile`, without that clone holding + /// the real sender (which would keep the channel open forever) or + /// being able to spawn a second worker. + fn detached() -> Self { + let (tx, _rx) = mpsc::channel(1); + Self { + tx, + rx: Mutex::new(None), + started: OnceLock::new(), + } + } + + /// Enqueue a reconciliation request for the background worker. + /// + /// # Errors + /// + /// Returns [`BackendError::Unavailable`] if the queue is full or the + /// worker has stopped. + fn enqueue(&self, request: ReconcileRequest) -> Result<(), BackendError> { + self.tx + .try_send(request) + .map_err(|error| BackendError::Unavailable(format!("reconciliation queue is full or stopped: {error}"))) + } + + /// Lazily spawn [`run_reconcile_worker`] on `runtime`, at most once. + /// `make_worker` builds the backend clone the worker itself will + /// call `reconcile` against (see [`Self::detached`]). + fn start(&self, runtime: &tokio::runtime::Handle, make_worker: impl FnOnce() -> B) + where + B: TokenRateLimitStateBackend + 'static, + { + self.started.get_or_init(|| { + let Some(receiver) = self.rx.lock().ok().and_then(|mut guard| guard.take()) else { + return; + }; + runtime.spawn(run_reconcile_worker(make_worker(), receiver)); + }); + } +} + +/// Valkey/Redis-backed sliding-window state, shared across every gateway +/// instance/replica pointed at the same `url`/`namespace`. +/// +/// Admission (`reserve`) is synchronous with the request (an EVAL round- +/// trip); reconciliation (`enqueue_reconcile`) is deferred to a +/// background worker so it never adds latency to the response path. +pub(super) struct ValkeyTokenRateLimitBackend { + /// Shared connection/EVAL handling, see [`ValkeyEval`]. + valkey: ValkeyEval, + /// Key namespace prefix, see [`ValkeyBackendConfig::namespace`]. + namespace: String, + /// Rule identifier, see [`ValkeyBackendConfig::rule`]. + rule: String, + /// Sliding-window budgets enforced atomically per key. + budgets: Vec, + /// See [`ValkeyBackendConfig::reservation_timeout_ms`]. + reservation_timeout_ms: u64, + /// See [`ValkeyBackendConfig::max_keys`]. + max_keys: usize, + /// See [`ValkeyBackendConfig::max_active_reservations`]. + max_active_reservations: usize, + /// Smallest configured budget capacity, for rate-limit headers. + limit: u64, + /// Shared background-reconciliation scaffolding, see [`ReconcileWorker`]. + worker: ReconcileWorker, +} + +/// Construction parameters for [`ValkeyTokenRateLimitBackend`]. +pub(super) struct ValkeyBackendConfig { + /// Filter-level Valkey connection, shared (`Clone`d) across every + /// Valkey-backed rule -- see [`ValkeyEval`]'s doc comment. + pub(super) valkey: ValkeyEval, + /// Key namespace prefix, isolating this rule's state from any other + /// rule/deployment sharing the same Valkey instance. + pub(super) namespace: String, + /// Rule identifier, folded into the per-key hash alongside `namespace`. + pub(super) rule: String, + /// Sliding-window budgets enforced atomically per key. + pub(super) budgets: Vec, + /// Time after which an ambiguous (never-reconciled) reservation is + /// charged at its estimate, mirroring the in-memory ledger's own + /// field of the same name. + pub(super) reservation_timeout_ms: u64, + /// Maximum distinct keys retained per namespace. + pub(super) max_keys: usize, + /// Maximum reservations awaiting reconciliation across all keys in + /// this namespace. + pub(super) max_active_reservations: usize, +} + +impl ValkeyTokenRateLimitBackend { + /// Build this rule's backend from an already-open, filter-shared + /// [`ValkeyEval`] connection. + pub(super) fn new(config: ValkeyBackendConfig) -> Self { + let limit = config.budgets.iter().map(|budget| budget.capacity).min().unwrap_or(0); + Self { + valkey: config.valkey, + namespace: config.namespace, + rule: config.rule, + budgets: config.budgets, + reservation_timeout_ms: config.reservation_timeout_ms, + max_keys: config.max_keys, + max_active_reservations: config.max_active_reservations, + limit, + worker: ReconcileWorker::new(), + } + } + + /// Clone this backend's connection/config, but with a detached + /// [`ReconcileWorker`] -- used only to hand the background worker its + /// own handle to `reserve`/`reconcile` (see [`ReconcileWorker::detached`]). + fn clone_without_sender(&self) -> Self { + Self { + valkey: self.valkey.clone(), + namespace: self.namespace.clone(), + rule: self.rule.clone(), + budgets: self.budgets.clone(), + reservation_timeout_ms: self.reservation_timeout_ms, + max_keys: self.max_keys, + max_active_reservations: self.max_active_reservations, + limit: self.limit, + worker: ReconcileWorker::detached(), + } + } + + /// Lazily spawn the background reconciliation worker on the calling + /// Tokio runtime, at most once per backend instance. + /// + /// # Errors + /// + /// Returns [`BackendError::Unavailable`] if called outside a Tokio + /// runtime context. + fn start_worker(&self) -> Result<(), BackendError> { + let runtime = tokio::runtime::Handle::try_current() + .map_err(|_error| BackendError::Unavailable("Valkey reconciliation requires a Tokio runtime".into()))?; + self.worker.start(&runtime, || self.clone_without_sender()); + Ok(()) + } + + /// Deterministic per-key Valkey key names for this rule/namespace. + fn key_parts(&self, key: &str) -> [String; 7] { + let mut digest = Sha256::new(); + digest.update(self.namespace.as_bytes()); + digest.update([0]); + digest.update(self.rule.as_bytes()); + digest.update([0]); + digest.update(key.as_bytes()); + let hash = digest + .finalize() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + let prefix = format!("{}:v1:{}", self.namespace, hash); + [ + prefix.clone(), + format!("{prefix}:settled"), + format!("{prefix}:active"), + format!("{}:keys", self.namespace), + format!("{}:active-count", self.namespace), + format!("{}:reservation-seq", self.namespace), + format!("{}:active-index", self.namespace), + ] + } +} + +/// Arguments for [`RESERVE_SCRIPT`]: timeout/bounds, then one +/// `(window_ms, capacity)` pair per configured budget. +fn reserve_args( + reservation_timeout_ms: u64, + max_keys: usize, + max_active_reservations: usize, + request: &ReserveRequest, + budgets: &[Budget], +) -> Vec { + let mut args = vec![ + reservation_timeout_ms.to_string(), + max_keys.to_string(), + max_active_reservations.to_string(), + request.estimate.to_string(), + budgets.len().to_string(), + ]; + for budget in budgets { + args.push(budget.window_ms.to_string()); + args.push(budget.capacity.to_string()); + } + args +} + +#[async_trait] +impl TokenRateLimitStateBackend for ValkeyTokenRateLimitBackend { + async fn reserve(&self, request: ReserveRequest) -> Result { + let keys = self.key_parts(&request.key); + let args = reserve_args( + self.reservation_timeout_ms, + self.max_keys, + self.max_active_reservations, + &request, + &self.budgets, + ); + let response = self.valkey.eval(RESERVE_SCRIPT, &keys, &args).await?; + match response.as_slice() { + [1, id, estimate] => Ok(BackendReserve::Admitted { + reservation_id: u64::try_from(*id).map_err(|_error| BackendError::InvalidResponse)?, + estimate: u64::try_from(*estimate).map_err(|_error| BackendError::InvalidResponse)?, + }), + [0, retry_after] => Ok(BackendReserve::Denied { + retry_after_ms: u64::try_from(*retry_after).map_err(|_error| BackendError::InvalidResponse)?, + }), + _ => Err(BackendError::InvalidResponse), + } + } + + async fn reconcile(&self, request: ReconcileRequest) -> Result { + let keys = self.key_parts(&request.key); + let actual = request.actual.unwrap_or(request.estimate); + let args = [request.reservation_id.to_string(), actual.to_string()]; + let response = self.valkey.eval(RECONCILE_SCRIPT, &keys, &args).await?; + match response.as_slice() { + [0] => Ok(BackendSettlement::Noop), + [1, actual, refund, overage] => Ok(BackendSettlement::Applied { + actual: u64::try_from(*actual).map_err(|_error| BackendError::InvalidResponse)?, + refund: u64::try_from(*refund).map_err(|_error| BackendError::InvalidResponse)?, + overage: u64::try_from(*overage).map_err(|_error| BackendError::InvalidResponse)?, + }), + _ => Err(BackendError::InvalidResponse), + } + } + + fn enqueue_reconcile(&self, request: ReconcileRequest) -> Result<(), BackendError> { + self.start_worker()?; + self.worker.enqueue(request) + } + + fn limit(&self) -> u64 { + self.limit + } +} + +// ----------------------------------------------------------------------------- +// Valkey-backed token bucket +// ----------------------------------------------------------------------------- + +/// Atomically admit a reservation against one key's token bucket, or deny +/// it -- the Valkey/Lua analog of [`TokenBucketLedger::reserve`]. +/// +/// `KEYS`: `[1]` physical state hash (`tokens`, `last_refill_ms`), `[2]` +/// active hash, `[3]` namespace keys zset, `[4]` namespace active-count +/// string, `[5]` namespace reservation-id sequence, `[6]` namespace +/// active-index zset. Deliberately namespaced with a `:tb:` segment +/// distinct from [`RESERVE_SCRIPT`]'s sliding-window keys (see +/// [`ValkeyTokenBucketBackend::key_parts`]) so a `token_bucket` rule and +/// a `sliding_window` rule can safely share one `namespace:` without +/// either algorithm's bookkeeping corrupting the other's. `ARGV`: `[1]` +/// capacity, `[2]` `refill_rate` (tokens/sec), `[3]` reservation timeout +/// (ms), `[4]` max keys, `[5]` max active reservations, `[6]` estimate. +/// Returns `[1, id, estimate]` on admission or `[0, retry_after_ms]` on +/// denial. +pub(super) const TOKEN_BUCKET_RESERVE_SCRIPT: &str = " +local now = redis.call('TIME') +local now_ms = tonumber(now[1]) * 1000 + math.floor(tonumber(now[2]) / 1000) +local capacity = tonumber(ARGV[1]) +local refill_rate = tonumber(ARGV[2]) +local timeout_ms = tonumber(ARGV[3]) +local max_keys = tonumber(ARGV[4]) +local max_active = tonumber(ARGV[5]) +local estimate = tonumber(ARGV[6]) + +local active_total = tonumber(redis.call('GET', KEYS[4]) or '0') + +local expired_global = redis.call('ZRANGE', KEYS[6], '-inf', now_ms, 'BYSCORE') +for i = 1, #expired_global do + local member = expired_global[i] + local split = string.find(member, '|') + if split then + local physical = string.sub(member, 1, split - 1) + local reservation = string.sub(member, split + 1) + local active_key = physical .. ':active' + local value = redis.call('HGET', active_key, reservation) + if value then + redis.call('HDEL', active_key, reservation) + active_total = math.max(0, active_total - 1) + -- Tokens for an abandoned reservation stay charged (already + -- decremented at reserve time under the immediate-decrement + -- design): no credit-back happens on expiry, only on reconcile. + end + end + redis.call('ZREM', KEYS[6], member) +end +redis.call('SET', KEYS[4], active_total) + +local state = redis.call('HMGET', KEYS[1], 'tokens', 'last_refill_ms') +local tokens = tonumber(state[1]) +local last_refill_ms = tonumber(state[2]) +if tokens == nil then + tokens = capacity + last_refill_ms = now_ms +end +local elapsed_ms = math.max(0, now_ms - last_refill_ms) +tokens = math.min(capacity, tokens + (elapsed_ms / 1000.0) * refill_rate) + +local ttl = math.max(math.ceil((capacity / refill_rate) * 1000) + timeout_ms, 1000) +redis.call('ZREMRANGEBYSCORE', KEYS[3], '-inf', now_ms) +local key_exists = redis.call('EXISTS', KEYS[1]) == 1 +if not key_exists and redis.call('ZCARD', KEYS[3]) >= max_keys then + redis.call('HSET', KEYS[1], 'tokens', tokens, 'last_refill_ms', now_ms) + return {0, 1} +end +if active_total >= max_active then + redis.call('HSET', KEYS[1], 'tokens', tokens, 'last_refill_ms', now_ms) + return {0, 1} +end +if tokens < estimate then + local deficit = estimate - tokens + local retry_after_ms = math.max(1, math.ceil((deficit / refill_rate) * 1000)) + redis.call('HSET', KEYS[1], 'tokens', tokens, 'last_refill_ms', now_ms) + redis.call('PEXPIRE', KEYS[1], ttl) + return {0, retry_after_ms} +end + +tokens = tokens - estimate +local id = redis.call('INCR', KEYS[5]) +redis.call('HSET', KEYS[2], id, estimate .. '|' .. now_ms) +redis.call('INCR', KEYS[4]) +redis.call('ZADD', KEYS[6], now_ms + timeout_ms, KEYS[1] .. '|' .. id) +redis.call('HSET', KEYS[1], 'tokens', tokens, 'last_refill_ms', now_ms) +redis.call('ZADD', KEYS[3], now_ms + ttl, KEYS[1]) +redis.call('PEXPIRE', KEYS[1], ttl) +redis.call('PEXPIRE', KEYS[2], ttl) +return {1, id, estimate} +"; + +/// Atomically settle a prior token-bucket reservation against actual +/// usage -- the Valkey/Lua analog of [`TokenBucketLedger::reconcile`]. +/// +/// `KEYS`: same layout as [`TOKEN_BUCKET_RESERVE_SCRIPT`]. `ARGV`: `[1]` +/// reservation ID, `[2]` actual usage, `[3]` capacity, `[4]` `refill_rate`. +/// Returns `[0]` if the reservation was already reconciled/expired +/// (no-op), or `[1, actual, refund, overage]`. +const TOKEN_BUCKET_RECONCILE_SCRIPT: &str = " +local value = redis.call('HGET', KEYS[2], ARGV[1]) +if not value then return {0} end +local sep = string.find(value, '|') +local estimate = tonumber(string.sub(value, 1, sep - 1)) +local actual = tonumber(ARGV[2]) +local capacity = tonumber(ARGV[3]) +local refill_rate = tonumber(ARGV[4]) +redis.call('HDEL', KEYS[2], ARGV[1]) +local active_total = math.max(0, tonumber(redis.call('GET', KEYS[4]) or '0') - 1) +redis.call('SET', KEYS[4], active_total) +redis.call('ZREM', KEYS[6], KEYS[1] .. '|' .. ARGV[1]) + +local now = redis.call('TIME') +local now_ms = tonumber(now[1]) * 1000 + math.floor(tonumber(now[2]) / 1000) +local state = redis.call('HMGET', KEYS[1], 'tokens', 'last_refill_ms') +local tokens = tonumber(state[1]) +local last_refill_ms = tonumber(state[2]) +if tokens == nil then + tokens = capacity + last_refill_ms = now_ms +end +local elapsed_ms = math.max(0, now_ms - last_refill_ms) +tokens = math.min(capacity, tokens + (elapsed_ms / 1000.0) * refill_rate) + +local refund = math.max(0, estimate - actual) +local overage = math.max(0, actual - estimate) +if refund > 0 then + tokens = math.min(capacity, tokens + refund) +elseif overage > 0 then + tokens = math.max(0, tokens - overage) +end +redis.call('HSET', KEYS[1], 'tokens', tokens, 'last_refill_ms', now_ms) +return {1, actual, refund, overage} +"; + +/// Valkey/Redis-backed token-bucket state, shared across every gateway +/// instance/replica pointed at the same `url`/`namespace`. +pub(super) struct ValkeyTokenBucketBackend { + /// Shared connection/EVAL handling, see [`ValkeyEval`]. + valkey: ValkeyEval, + /// Key namespace prefix, see [`ValkeyBackendConfig::namespace`]. + namespace: String, + /// Rule identifier, see [`ValkeyBackendConfig::rule`]. + rule: String, + /// Maximum tokens held at once. + capacity: u64, + /// Tokens refilled per second, up to `capacity`. + refill_rate: f64, + /// See [`ValkeyBackendConfig::reservation_timeout_ms`]. + reservation_timeout_ms: u64, + /// See [`ValkeyBackendConfig::max_keys`]. + max_keys: usize, + /// See [`ValkeyBackendConfig::max_active_reservations`]. + max_active_reservations: usize, + /// Shared background-reconciliation scaffolding, see [`ReconcileWorker`]. + worker: ReconcileWorker, +} + +/// Construction parameters for [`ValkeyTokenBucketBackend`]. +pub(super) struct ValkeyTokenBucketConfig { + /// Filter-level Valkey connection, shared (`Clone`d) across every + /// Valkey-backed rule -- see [`ValkeyEval`]'s doc comment. + pub(super) valkey: ValkeyEval, + /// Key namespace prefix, isolating this rule's state from any other + /// rule/deployment sharing the same Valkey instance. + pub(super) namespace: String, + /// Rule identifier, folded into the per-key hash alongside `namespace`. + pub(super) rule: String, + /// Maximum tokens held at once. + pub(super) capacity: u64, + /// Tokens refilled per second, up to `capacity`. + pub(super) refill_rate: f64, + /// Time after which an ambiguous (never-reconciled) reservation + /// stops being tracked as active (it's already charged). + pub(super) reservation_timeout_ms: u64, + /// Maximum distinct keys retained per namespace/algorithm. + pub(super) max_keys: usize, + /// Maximum reservations awaiting reconciliation across all keys in + /// this namespace/algorithm. + pub(super) max_active_reservations: usize, +} + +impl ValkeyTokenBucketBackend { + /// Build this rule's backend from an already-open, filter-shared + /// [`ValkeyEval`] connection. + /// + /// # Errors + /// + /// Returns [`BackendError::Unavailable`] if `capacity`/`refill_rate` + /// aren't positive and finite, or if `capacity` or `capacity / + /// refill_rate` exceeds the bounds documented on + /// [`token_bucket_ledger::MAX_F64_SAFE_INTEGER`]/ + /// [`token_bucket_ledger::MAX_CAPACITY_REFILL_RATE_RATIO_SECS`]. + pub(super) fn new(config: ValkeyTokenBucketConfig) -> Result { + token_bucket_ledger::validate_capacity_and_refill_rate(config.capacity, config.refill_rate) + .map_err(BackendError::Unavailable)?; + Ok(Self { + valkey: config.valkey, + namespace: config.namespace, + rule: config.rule, + capacity: config.capacity, + refill_rate: config.refill_rate, + reservation_timeout_ms: config.reservation_timeout_ms, + max_keys: config.max_keys, + max_active_reservations: config.max_active_reservations, + worker: ReconcileWorker::new(), + }) + } + + /// See [`ValkeyTokenRateLimitBackend::clone_without_sender`]. + fn clone_without_sender(&self) -> Self { + Self { + valkey: self.valkey.clone(), + namespace: self.namespace.clone(), + rule: self.rule.clone(), + capacity: self.capacity, + refill_rate: self.refill_rate, + reservation_timeout_ms: self.reservation_timeout_ms, + max_keys: self.max_keys, + max_active_reservations: self.max_active_reservations, + worker: ReconcileWorker::detached(), + } + } + + /// See [`ValkeyTokenRateLimitBackend::start_worker`]. + /// + /// # Errors + /// + /// Returns [`BackendError::Unavailable`] if called outside a Tokio + /// runtime context. + fn start_worker(&self) -> Result<(), BackendError> { + let runtime = tokio::runtime::Handle::try_current() + .map_err(|_error| BackendError::Unavailable("Valkey reconciliation requires a Tokio runtime".into()))?; + self.worker.start(&runtime, || self.clone_without_sender()); + Ok(()) + } + + /// Deterministic per-key Valkey key names for this rule/namespace, + /// under a `:tb:` segment distinct from the sliding-window backend's + /// [`ValkeyTokenRateLimitBackend::key_parts`] -- see + /// [`TOKEN_BUCKET_RESERVE_SCRIPT`]'s doc comment for why the two + /// algorithms must never share bookkeeping keys. + fn key_parts(&self, key: &str) -> [String; 6] { + let mut digest = Sha256::new(); + digest.update(self.namespace.as_bytes()); + digest.update([0]); + digest.update(b"token_bucket"); + digest.update([0]); + digest.update(self.rule.as_bytes()); + digest.update([0]); + digest.update(key.as_bytes()); + let hash = digest + .finalize() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + let prefix = format!("{}:v1:tb:{}", self.namespace, hash); + [ + prefix.clone(), + format!("{prefix}:active"), + format!("{}:tb:keys", self.namespace), + format!("{}:tb:active-count", self.namespace), + format!("{}:tb:reservation-seq", self.namespace), + format!("{}:tb:active-index", self.namespace), + ] + } +} + +#[async_trait] +impl TokenRateLimitStateBackend for ValkeyTokenBucketBackend { + async fn reserve(&self, request: ReserveRequest) -> Result { + let keys = self.key_parts(&request.key); + let args = [ + self.capacity.to_string(), + self.refill_rate.to_string(), + self.reservation_timeout_ms.to_string(), + self.max_keys.to_string(), + self.max_active_reservations.to_string(), + request.estimate.to_string(), + ]; + let response = self.valkey.eval(TOKEN_BUCKET_RESERVE_SCRIPT, &keys, &args).await?; + match response.as_slice() { + [1, id, estimate] => Ok(BackendReserve::Admitted { + reservation_id: u64::try_from(*id).map_err(|_error| BackendError::InvalidResponse)?, + estimate: u64::try_from(*estimate).map_err(|_error| BackendError::InvalidResponse)?, + }), + [0, retry_after] => Ok(BackendReserve::Denied { + retry_after_ms: u64::try_from(*retry_after).map_err(|_error| BackendError::InvalidResponse)?, + }), + _ => Err(BackendError::InvalidResponse), + } + } + + async fn reconcile(&self, request: ReconcileRequest) -> Result { + let keys = self.key_parts(&request.key); + let actual = request.actual.unwrap_or(request.estimate); + let args = [ + request.reservation_id.to_string(), + actual.to_string(), + self.capacity.to_string(), + self.refill_rate.to_string(), + ]; + let response = self.valkey.eval(TOKEN_BUCKET_RECONCILE_SCRIPT, &keys, &args).await?; + match response.as_slice() { + [0] => Ok(BackendSettlement::Noop), + [1, actual, refund, overage] => Ok(BackendSettlement::Applied { + actual: u64::try_from(*actual).map_err(|_error| BackendError::InvalidResponse)?, + refund: u64::try_from(*refund).map_err(|_error| BackendError::InvalidResponse)?, + overage: u64::try_from(*overage).map_err(|_error| BackendError::InvalidResponse)?, + }), + _ => Err(BackendError::InvalidResponse), + } + } + + fn enqueue_reconcile(&self, request: ReconcileRequest) -> Result<(), BackendError> { + self.start_worker()?; + self.worker.enqueue(request) + } + + fn limit(&self) -> u64 { + self.capacity + } +} + +#[cfg(test)] +#[expect(clippy::allow_attributes, reason = "blanket test suppressions")] +#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic, reason = "tests")] +mod tests { + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + use super::*; + use crate::token_rate_limit::ledger::LedgerConfig; + + fn memory_backend(capacity: u64) -> InMemoryTokenRateLimitBackend { + let ledger = Ledger::new(LedgerConfig { + budgets: vec![Budget { + window_ms: 60_000, + capacity, + }], + reservation_timeout_ms: 1_000, + max_keys: 8, + max_key_length: 64, + max_active_reservations: 8, + }) + .unwrap(); + InMemoryTokenRateLimitBackend::new(ledger) + } + + #[tokio::test] + async fn reconcile_sync_settles_in_process_state_without_a_network_round_trip() { + let backend = memory_backend(100); + let admitted = backend + .reserve(ReserveRequest { + key: "a".into(), + estimate: 40, + now_ms: 0, + }) + .await + .unwrap(); + let BackendReserve::Admitted { reservation_id, .. } = admitted else { + panic!("expected admission") + }; + + let settlement = backend.reconcile_sync(&ReconcileRequest { + key: "a".into(), + reservation_id, + actual: Some(10), + estimate: 40, + now_ms: 0, + }); + assert_eq!( + settlement, + Some(BackendSettlement::Applied { + actual: 10, + refund: 30, + overage: 0 + }), + "in-process backend must resolve reconcile_sync synchronously, without a Valkey-style enqueued worker" + ); + } + + #[tokio::test] + async fn cleanup_reports_active_reservations_and_keys_for_in_process_state() { + let backend = memory_backend(100); + backend + .reserve(ReserveRequest { + key: "a".into(), + estimate: 10, + now_ms: 0, + }) + .await + .unwrap(); + + let report = backend + .cleanup(0, 8) + .expect("in-process backend must report cleanup state for gauges"); + assert_eq!( + report.active_reservations, 1, + "one un-reconciled reservation should be counted" + ); + assert_eq!(report.active_keys, 1, "one distinct key should be tracked"); + assert_eq!(report.orphaned, 0, "nothing has timed out yet"); + } + + /// `reconcile`/`enqueue_reconcile` are part of the shared + /// [`TokenRateLimitStateBackend`] trait contract -- callers reach an + /// in-process backend exclusively through `reconcile_sync` today (see + /// `TokenRateLimitFilter::reconcile`'s doc comment), but the trait + /// methods themselves must still behave correctly for any future or + /// generic (`Arc`) caller that goes + /// through them instead. + /// Reserve `estimate` against `backend`, returning the resulting + /// reservation ID (panics if denied -- every caller below reserves + /// well within its backend's configured capacity). + async fn reserve_or_panic(backend: &impl TokenRateLimitStateBackend, estimate: u64) -> u64 { + let admitted = backend + .reserve(ReserveRequest { + key: "a".into(), + estimate, + now_ms: 0, + }) + .await + .unwrap(); + let BackendReserve::Admitted { reservation_id, .. } = admitted else { + panic!("expected admission") + }; + reservation_id + } + + /// The `reconcile` half of + /// `assert_trait_reconcile_methods_apply_directly`, split out to keep + /// both under clippy's function-length budget. + async fn assert_trait_reconcile_applies_directly(backend: &impl TokenRateLimitStateBackend) { + let reservation_id = reserve_or_panic(backend, 50).await; + let settlement = backend + .reconcile(ReconcileRequest { + key: "a".into(), + reservation_id, + actual: Some(10), + estimate: 50, + now_ms: 0, + }) + .await + .unwrap(); + assert_eq!( + settlement, + BackendSettlement::Applied { + actual: 10, + refund: 40, + overage: 0 + } + ); + } + + /// The `enqueue_reconcile` half -- see + /// `assert_trait_reconcile_applies_directly`. The in-process + /// implementation applies it inline rather than truly deferring it, + /// but must still succeed and take effect. + async fn assert_trait_enqueue_reconcile_applies_directly(backend: &impl TokenRateLimitStateBackend) { + let reservation_id = reserve_or_panic(backend, 40).await; + backend + .enqueue_reconcile(ReconcileRequest { + key: "a".into(), + reservation_id, + actual: Some(5), + estimate: 40, + now_ms: 0, + }) + .unwrap(); + let request = ReserveRequest { + key: "a".into(), + estimate: 35, + now_ms: 0, + }; + assert!( + matches!(backend.reserve(request).await.unwrap(), BackendReserve::Admitted { .. }), + "enqueue_reconcile must have released the 35 unused tokens from the second reservation" + ); + } + + #[tokio::test] + async fn in_memory_sliding_window_backend_trait_reconcile_methods_apply_directly() { + let backend = memory_backend(100); + assert_trait_reconcile_applies_directly(&backend).await; + assert_trait_enqueue_reconcile_applies_directly(&backend).await; + } + + /// The token-bucket analog of + /// `in_memory_sliding_window_backend_trait_reconcile_methods_apply_directly`. + #[tokio::test] + async fn in_memory_token_bucket_backend_trait_reconcile_methods_apply_directly() { + let backend = bucket_backend(100, 1.0); + assert_trait_reconcile_applies_directly(&backend).await; + assert_trait_enqueue_reconcile_applies_directly(&backend).await; + } + + /// Reconciling an unknown/already-settled reservation ID must be a + /// silent no-op (idempotent double-reconciliation), never a panic or + /// a double-credit -- for both algorithms' in-process backends, + /// through both the async `reconcile` trait method and the + /// synchronous `reconcile_sync` fast path. + #[tokio::test] + async fn in_memory_backends_reconcile_is_noop_for_an_unknown_reservation_id() { + let unknown_reservation = ReconcileRequest { + key: "a".into(), + reservation_id: 999_999, + actual: Some(1), + estimate: 1, + now_ms: 0, + }; + + let sliding = memory_backend(100); + assert_eq!( + sliding.reconcile(unknown_reservation.clone()).await.unwrap(), + BackendSettlement::Noop + ); + assert_eq!( + sliding.reconcile_sync(&unknown_reservation), + Some(BackendSettlement::Noop) + ); + + let bucket = bucket_backend(100, 1.0); + assert_eq!( + bucket.reconcile(unknown_reservation.clone()).await.unwrap(), + BackendSettlement::Noop + ); + assert_eq!( + bucket.reconcile_sync(&unknown_reservation), + Some(BackendSettlement::Noop) + ); + } + + /// [`ReconcileWorker::start`] on a [`ReconcileWorker::detached`] + /// worker must be a no-op: there's no receiver to hand a spawned + /// [`run_reconcile_worker`], so it must return without ever calling + /// `make_worker` (a real caller passes a closure that builds a live + /// backend clone there -- doing that unnecessarily would be wasted + /// work at best and a logic error at worst). + #[test] + fn reconcile_worker_start_on_a_detached_worker_never_spawns() { + let worker = ReconcileWorker::detached(); + let runtime = tokio::runtime::Builder::new_current_thread().build().unwrap(); + worker.start(runtime.handle(), || -> AlwaysFailsReconcile { + panic!("a detached ReconcileWorker has no receiver to hand a spawned worker -- make_worker must not run") + }); + } + + /// [`ValkeyEval::eval`] must invalidate its cached connection and + /// surface an error on any command failure -- not just a connection + /// failure -- so a subsequent call re-establishes a fresh connection + /// rather than reusing one Valkey has already rejected a command on. + /// Requires a live Valkey/Redis (see `TOKEN_RATE_LIMIT_VALKEY_URL` in + /// `tests.rs`); skips otherwise. + #[tokio::test] + async fn valkey_eval_invalidates_the_connection_after_a_script_error_and_recovers() { + let Ok(url) = std::env::var("TOKEN_RATE_LIMIT_VALKEY_URL") else { + tracing::warn!("skipping: TOKEN_RATE_LIMIT_VALKEY_URL not set"); + return; + }; + let valkey = ValkeyEval::new(url).unwrap(); + + // A live, successfully-connected session: establishes and caches + // the connection this test then forces `eval` to invalidate. + // `{1}` (a Lua table), not a bare `1`, since `eval` deserializes + // into `Vec` (a multi-bulk reply), same shape as the real + // reserve/reconcile scripts. + assert!( + valkey.eval("return {1}", &["k".to_owned()], &[]).await.is_ok(), + "sanity check: a trivial script must succeed against a live Valkey" + ); + + // A script Valkey's Lua interpreter rejects outright (unbalanced + // syntax) -- the command round-trips successfully at the + // connection level, but Valkey replies with an error, exercising + // the `Ok(Err(error))` arm of `eval`'s match (as opposed to + // `valkey_failure_fails_closed`'s unreachable-host connection + // failure, which exercises the earlier `connection()` error path). + let result = valkey.eval("this is not valid lua(", &["k".to_owned()], &[]).await; + assert!( + result.is_err(), + "an invalid script must surface as a BackendError, not panic or hang" + ); + + // The connection must have been invalidated and cleanly + // re-established, not left wedged, for the next legitimate call. + assert!( + valkey.eval("return {1}", &["k".to_owned()], &[]).await.is_ok(), + "eval must recover on the next call after invalidating a failed connection" + ); + } + + /// [`map_valkey_error`] must tag its message with which phase + /// (`"connection"` vs. `"command"`) the underlying [`redis::RedisError`] + /// came from -- `redis`'s own error text alone doesn't say (e.g. a + /// response-timeout's `Display` is a bare `"timed out"`), and that + /// distinction is the only thing this crate's own wrapping around + /// `redis`'s errors adds. + #[test] + fn map_valkey_error_tags_the_message_with_which_phase_failed() { + let timed_out = redis::RedisError::from(std::io::Error::from(std::io::ErrorKind::TimedOut)); + + let BackendError::Unavailable(message) = map_valkey_error("connection", &timed_out) else { + panic!("map_valkey_error must always return BackendError::Unavailable") + }; + assert_eq!(message, "Valkey connection: timed out"); + + let BackendError::Unavailable(message) = map_valkey_error("command", &timed_out) else { + panic!("map_valkey_error must always return BackendError::Unavailable") + }; + assert_eq!(message, "Valkey command: timed out"); + } + + /// A one-shot TCP proxy in front of `upstream`'s `host:port`, for + /// fault-injecting a Valkey that accepts a command and then hangs. + /// Real Valkey/Redis has no config knob for this; `DEBUG SLEEP` + /// comes closest but additionally requires `enable-debug-command` + /// server-side and isn't callable from a script at all, so this + /// proxies the real, unmodified Valkey under test instead of + /// relying on either. + /// + /// Returns the proxy's local address and a flag that, once set, + /// makes every open connection stop relaying upstream replies back + /// to the client -- the bytes are still read off the wire (so the + /// upstream Valkey itself never blocks or errors), just dropped. + async fn spawn_wedgeable_proxy(upstream: &str) -> (std::net::SocketAddr, Arc) { + let upstream = upstream + .strip_prefix("redis://") + .expect("test fixture: TOKEN_RATE_LIMIT_VALKEY_URL must be a bare redis://host:port URL") + .to_owned(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_addr = listener.local_addr().unwrap(); + let wedged = Arc::new(AtomicBool::new(false)); + + let wedged_for_task = Arc::clone(&wedged); + tokio::spawn(async move { + while let Ok((client, _)) = listener.accept().await { + tokio::spawn(relay_one_connection( + client, + upstream.clone(), + Arc::clone(&wedged_for_task), + )); + } + }); + + (proxy_addr, wedged) + } + + /// One [`spawn_wedgeable_proxy`] connection's relay loop: client + /// bytes always flow through to `upstream` unmodified; `upstream`'s + /// replies flow back to the client unless/until `wedged` is set, at + /// which point they're read off the wire (so `upstream` never + /// blocks) but silently dropped instead of relayed. + async fn relay_one_connection(mut client: tokio::net::TcpStream, upstream: String, wedged: Arc) { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + + let Ok(mut server) = tokio::net::TcpStream::connect(&upstream).await else { + return; + }; + let (mut client_read, mut client_write) = client.split(); + let (mut server_read, mut server_write) = server.split(); + tokio::join!( + async { + drop(tokio::io::copy(&mut client_read, &mut server_write).await); + }, + async { + let mut buffer = [0_u8; 4096]; + loop { + let Ok(read @ 1..) = server_read.read(&mut buffer).await else { + return; + }; + if wedged.load(Ordering::SeqCst) { + continue; // Accepted off the wire, never relayed: a silent hang. + } + let Some(bytes) = buffer.get(..read) else { + return; + }; + if client_write.write_all(bytes).await.is_err() { + return; + } + } + } + ); + } + + /// [`ValkeyEval::eval`] must fail closed at (roughly) [`VALKEY_TIMEOUT`] + /// on a command Valkey accepts but never replies to, not hang + /// indefinitely -- proving [`ValkeyEval::connection_config`] (not + /// just `redis`'s own, possibly-different, default) is what's + /// actually bounding the response wait. Requires a live + /// Valkey/Redis (see `TOKEN_RATE_LIMIT_VALKEY_URL` in `tests.rs`); + /// skips otherwise. + #[tokio::test] + async fn eval_times_out_and_invalidates_the_connection_when_wedged() { + let Ok(url) = std::env::var("TOKEN_RATE_LIMIT_VALKEY_URL") else { + tracing::warn!("skipping: TOKEN_RATE_LIMIT_VALKEY_URL not set"); + return; + }; + + let (proxy_addr, wedged) = spawn_wedgeable_proxy(&url).await; + let valkey = ValkeyEval::new(format!("redis://{proxy_addr}")).unwrap(); + + assert!( + valkey.eval("return {1}", &["k".to_owned()], &[]).await.is_ok(), + "sanity check: a trivial script must succeed through an unwedged proxy" + ); + + wedged.store(true, Ordering::SeqCst); + let started = std::time::Instant::now(); + let result = valkey.eval("return {1}", &["k".to_owned()], &[]).await; + let elapsed = started.elapsed(); + + assert!( + matches!(&result, Err(BackendError::Unavailable(message)) if message.starts_with("Valkey command:")), + "a wedged command must fail closed with a command-phase error, not hang or panic: {result:?}" + ); + assert!( + elapsed < VALKEY_TIMEOUT * 3, + "must fail closed at ~VALKEY_TIMEOUT ({VALKEY_TIMEOUT:?}), not wait indefinitely \ + for the wedge to clear: took {elapsed:?}" + ); + + // Unwedge and confirm the connection was invalidated, not left + // cached in its half-dead state -- the next call must + // re-establish cleanly rather than time out again. + wedged.store(false, Ordering::SeqCst); + assert!( + valkey.eval("return {1}", &["k".to_owned()], &[]).await.is_ok(), + "eval must recover on the next call after invalidating a timed-out connection" + ); + } + + /// Proves the actual mechanism the filter-level (not per-rule) + /// `backend:` config relies on to avoid opening a redundant Valkey + /// connection per rule: cloning a [`ValkeyEval`] must share the same + /// underlying connection cache, not build an independent one. + #[test] + fn cloning_valkey_eval_shares_the_same_connection_cache() { + let valkey = ValkeyEval::new("redis://127.0.0.1:1".into()).unwrap(); + let cloned = valkey.clone(); + assert!( + Arc::ptr_eq(&valkey.connection, &cloned.connection), + "a ValkeyEval clone (as handed to every Valkey-backed rule) must share one cached \ + connection, not each hold its own independent cache" + ); + } + + #[test] + fn valkey_backend_has_no_local_state_to_reconcile_or_clean_up_synchronously() { + // Business behavior under test: a networked backend must never + // silently answer a synchronous, no-I/O query -- the filter relies + // on `None` here to route reconciliation through the background + // worker (`enqueue_reconcile`) instead, and to skip gauge + // reporting rather than publish misleading zeros. No live Valkey + // is required: both methods short-circuit before any I/O. + let backend = ValkeyTokenRateLimitBackend::new(ValkeyBackendConfig { + valkey: ValkeyEval::new("redis://127.0.0.1:1".into()).unwrap(), + namespace: "ns".into(), + rule: "default".into(), + budgets: vec![Budget { + window_ms: 1_000, + capacity: 10, + }], + reservation_timeout_ms: 1_000, + max_keys: 8, + max_active_reservations: 8, + }); + + assert!( + backend.cleanup(0, 8).is_none(), + "Valkey-backed state has no local ledger to clean up in-process" + ); + assert!( + backend + .reconcile_sync(&ReconcileRequest { + key: "a".into(), + reservation_id: 1, + actual: Some(1), + estimate: 1, + now_ms: 0, + }) + .is_none(), + "Valkey-backed reconciliation must go through enqueue_reconcile, not reconcile_sync" + ); + } + + // ------------------------------------------------------------------------- + // InMemoryTokenBucketBackend (trait-contract level -- exhaustive + // business-scenario coverage for refill/refund/overage/DoS bounds + // lives in `token_bucket_ledger::tests`; these confirm the backend + // wrapper faithfully exposes that ledger through the shared trait). + // ------------------------------------------------------------------------- + + fn bucket_backend(capacity: u64, refill_rate: f64) -> InMemoryTokenBucketBackend { + let ledger = TokenBucketLedger::new(token_bucket_ledger::TokenBucketConfig { + capacity, + refill_rate, + reservation_timeout_ms: 1_000, + max_keys: 8, + max_key_length: 64, + max_active_reservations: 8, + }) + .unwrap(); + InMemoryTokenBucketBackend::new(ledger) + } + + #[tokio::test] + async fn token_bucket_backend_admits_within_capacity_and_denies_over_it() { + let backend = bucket_backend(10, 1.0); + assert!(matches!( + backend + .reserve(ReserveRequest { + key: "a".into(), + estimate: 10, + now_ms: 0 + }) + .await + .unwrap(), + BackendReserve::Admitted { .. } + )); + assert!(matches!( + backend + .reserve(ReserveRequest { + key: "a".into(), + estimate: 1, + now_ms: 0 + }) + .await + .unwrap(), + BackendReserve::Denied { .. } + )); + } + + #[tokio::test] + async fn token_bucket_backend_limit_reports_configured_capacity() { + let backend = bucket_backend(250, 5.0); + assert_eq!(backend.limit(), 250); + } + + #[tokio::test] + async fn token_bucket_backend_reconcile_sync_credits_back_unused_estimate() { + let backend = bucket_backend(100, 1.0); + let admitted = backend + .reserve(ReserveRequest { + key: "a".into(), + estimate: 50, + now_ms: 0, + }) + .await + .unwrap(); + let BackendReserve::Admitted { reservation_id, .. } = admitted else { + panic!("expected admission") + }; + let settlement = backend.reconcile_sync(&ReconcileRequest { + key: "a".into(), + reservation_id, + actual: Some(10), + estimate: 50, + now_ms: 0, + }); + assert_eq!( + settlement, + Some(BackendSettlement::Applied { + actual: 10, + refund: 40, + overage: 0 + }) + ); + } + + #[tokio::test] + async fn token_bucket_backend_cleanup_reports_active_reservations_and_keys() { + let backend = bucket_backend(100, 1.0); + backend + .reserve(ReserveRequest { + key: "a".into(), + estimate: 10, + now_ms: 0, + }) + .await + .unwrap(); + let report = backend + .cleanup(0, 8) + .expect("in-process token bucket backend must report cleanup state for gauges"); + assert_eq!(report.active_reservations, 1); + assert_eq!(report.active_keys, 1); + } + + // ------------------------------------------------------------------------- + // ValkeyTokenBucketBackend construction-time validation and the + // same no-I/O trait-contract checks as the sliding-window backend. + // ------------------------------------------------------------------------- + + #[test] + fn valkey_token_bucket_backend_rejects_zero_capacity_or_refill_rate() { + let base = || ValkeyTokenBucketConfig { + valkey: ValkeyEval::new("redis://127.0.0.1:1".into()).unwrap(), + namespace: "ns".into(), + rule: "default".into(), + capacity: 10, + refill_rate: 1.0, + reservation_timeout_ms: 1_000, + max_keys: 8, + max_active_reservations: 8, + }; + assert!(ValkeyTokenBucketBackend::new(ValkeyTokenBucketConfig { capacity: 0, ..base() }).is_err()); + assert!( + ValkeyTokenBucketBackend::new(ValkeyTokenBucketConfig { + refill_rate: 0.0, + ..base() + }) + .is_err() + ); + } + + #[test] + fn valkey_token_bucket_backend_rejects_non_finite_refill_rate() { + // Proves the shared validator (see non_finite_refill_rate_is_rejected) + // is actually wired into this backend's constructor too. + let base = || ValkeyTokenBucketConfig { + valkey: ValkeyEval::new("redis://127.0.0.1:1".into()).unwrap(), + namespace: "ns".into(), + rule: "default".into(), + capacity: 10, + refill_rate: 1.0, + reservation_timeout_ms: 1_000, + max_keys: 8, + max_active_reservations: 8, + }; + for bad_rate in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { + assert!( + ValkeyTokenBucketBackend::new(ValkeyTokenBucketConfig { + refill_rate: bad_rate, + ..base() + }) + .is_err(), + "refill_rate {bad_rate} must be rejected as non-finite/non-positive" + ); + } + } + + #[test] + fn valkey_token_bucket_backend_rejects_capacity_exceeding_f64_safe_integer() { + // See MAX_F64_SAFE_INTEGER's doc comment for why this bound exists; + // proven here for the Valkey backend's own constructor. + let base = || ValkeyTokenBucketConfig { + valkey: ValkeyEval::new("redis://127.0.0.1:1".into()).unwrap(), + namespace: "ns".into(), + rule: "default".into(), + capacity: 10, + refill_rate: 1.0, + reservation_timeout_ms: 1_000, + max_keys: 8, + max_active_reservations: 8, + }; + assert!( + ValkeyTokenBucketBackend::new(ValkeyTokenBucketConfig { + capacity: token_bucket_ledger::MAX_F64_SAFE_INTEGER + 1, + ..base() + }) + .is_err(), + "capacity above 2^53 must be rejected before precision is silently lost" + ); + } + + #[test] + fn valkey_token_bucket_backend_rejects_a_refill_rate_ratio_exceeding_the_pexpire_ttl_bound() { + // See MAX_CAPACITY_REFILL_RATE_RATIO_SECS's doc comment for why + // this bound exists; proven here for the Valkey backend's own + // constructor. + let base = || ValkeyTokenBucketConfig { + valkey: ValkeyEval::new("redis://127.0.0.1:1".into()).unwrap(), + namespace: "ns".into(), + rule: "default".into(), + capacity: 10, + refill_rate: 1.0, + reservation_timeout_ms: 1_000, + max_keys: 8, + max_active_reservations: 8, + }; + assert!( + ValkeyTokenBucketBackend::new(ValkeyTokenBucketConfig { + capacity: 10, + refill_rate: 10.0 / (token_bucket_ledger::MAX_CAPACITY_REFILL_RATE_RATIO_SECS * 2.0), + ..base() + }) + .is_err(), + "a capacity/refill_rate ratio beyond the PEXPIRE TTL bound must be rejected" + ); + } + + #[test] + fn valkey_token_bucket_backend_has_no_local_state_to_reconcile_or_clean_up_synchronously() { + let backend = ValkeyTokenBucketBackend::new(ValkeyTokenBucketConfig { + valkey: ValkeyEval::new("redis://127.0.0.1:1".into()).unwrap(), + namespace: "ns".into(), + rule: "default".into(), + capacity: 10, + refill_rate: 1.0, + reservation_timeout_ms: 1_000, + max_keys: 8, + max_active_reservations: 8, + }) + .unwrap(); + assert!(backend.cleanup(0, 8).is_none()); + assert!( + backend + .reconcile_sync(&ReconcileRequest { + key: "a".into(), + reservation_id: 1, + actual: Some(1), + estimate: 1, + now_ms: 0, + }) + .is_none() + ); + } + + /// A [`ValkeyTokenBucketBackend`] and a [`ValkeyTokenRateLimitBackend`] + /// sharing the same `namespace`/`rule` -- the plausible, even likely, + /// operator config that + /// [`valkey_token_bucket_and_sliding_window_key_parts_never_collide_even_in_the_same_namespace`] exercises. + fn same_namespace_backends() -> (ValkeyTokenBucketBackend, ValkeyTokenRateLimitBackend) { + let bucket = ValkeyTokenBucketBackend::new(ValkeyTokenBucketConfig { + valkey: ValkeyEval::new("redis://127.0.0.1:1".into()).unwrap(), + namespace: "shared".into(), + rule: "same-rule-name".into(), + capacity: 10, + refill_rate: 1.0, + reservation_timeout_ms: 1_000, + max_keys: 8, + max_active_reservations: 8, + }) + .unwrap(); + let sliding = ValkeyTokenRateLimitBackend::new(ValkeyBackendConfig { + valkey: ValkeyEval::new("redis://127.0.0.1:1".into()).unwrap(), + namespace: "shared".into(), + rule: "same-rule-name".into(), + budgets: vec![Budget { + window_ms: 1_000, + capacity: 10, + }], + reservation_timeout_ms: 1_000, + max_keys: 8, + max_active_reservations: 8, + }); + (bucket, sliding) + } + + #[test] + fn valkey_token_bucket_and_sliding_window_key_parts_never_collide_even_in_the_same_namespace() { + // Two rules sharing one `namespace:` must never let one + // algorithm's Lua script reap or mutate the other's + // physical/bookkeeping keys. + let (bucket, sliding) = same_namespace_backends(); + let bucket_keys = bucket.key_parts("same-key"); + let sliding_keys = sliding.key_parts("same-key"); + for bucket_key in &bucket_keys { + assert!( + !sliding_keys.contains(bucket_key), + "token_bucket key {bucket_key} collided with a sliding_window key" + ); + } + } + + #[test] + fn worker_enqueue_fails_once_its_receiver_is_gone() { + let worker = ReconcileWorker::detached(); + let request = ReconcileRequest { + key: "a".into(), + reservation_id: 1, + actual: Some(1), + estimate: 1, + now_ms: 0, + }; + assert!(worker.enqueue(request).is_err()); + } + + /// A backend whose `reconcile` always fails, to drive + /// [`run_reconcile_worker`]'s bounded-retry-then-abandon path. + struct AlwaysFailsReconcile { + attempts: Arc, + } + + #[async_trait] + impl TokenRateLimitStateBackend for AlwaysFailsReconcile { + async fn reserve(&self, _request: ReserveRequest) -> Result { + panic!("not exercised by this test") + } + + async fn reconcile(&self, _request: ReconcileRequest) -> Result { + self.attempts.fetch_add(1, Ordering::SeqCst); + Err(BackendError::Unavailable("simulated failure".into())) + } + + fn enqueue_reconcile(&self, _request: ReconcileRequest) -> Result<(), BackendError> { + panic!("not exercised by this test") + } + + fn limit(&self) -> u64 { + 0 + } + } + + #[tokio::test] + async fn reconcile_worker_retries_then_abandons_a_persistently_failing_reconcile() { + let attempts = Arc::new(AtomicUsize::new(0)); + let (tx, rx) = mpsc::channel(1); + tokio::spawn(run_reconcile_worker( + AlwaysFailsReconcile { + attempts: Arc::clone(&attempts), + }, + rx, + )); + tx.send(ReconcileRequest { + key: "a".into(), + reservation_id: 1, + actual: Some(1), + estimate: 1, + now_ms: 0, + }) + .await + .unwrap(); + drop(tx); + + // 1 initial attempt + 2 retries (25ms, 50ms backoff) before abandoning. + tokio::time::sleep(Duration::from_millis(300)).await; + assert_eq!( + attempts.load(Ordering::SeqCst), + 3, + "must retry exactly twice, then abandon" + ); + } +} diff --git a/filters/src/token_rate_limit/config.rs b/filters/src/token_rate_limit/config.rs new file mode 100644 index 0000000000..c2f5e7e5ba --- /dev/null +++ b/filters/src/token_rate_limit/config.rs @@ -0,0 +1,375 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Praxis Contributors + +//! Deserialized YAML configuration for the `token_rate_limit` filter. + +use std::collections::BTreeMap; + +use serde::Deserialize; + +// ----------------------------------------------------------------------------- +// TokenRateLimitConfig +// ----------------------------------------------------------------------------- + +/// Deserialized YAML config for the `token_rate_limit` filter: an ordered +/// list of `rules`, each binding an optional match condition to an +/// algorithm choice (`sliding_window` or `token_bucket`) and that rule's +/// own budget. +/// +/// Experimental: requires the `token-rate-limit-filter` cargo feature, +/// which is off by default and activates the `experimental` marker. +/// This filter delivers the agreed M1/M2/M6 milestone scope, but its +/// parent proposal is not yet `accepted` and open questions remain +/// (HA/clustered-Valkey failure modes, and the relationship to +/// Kuadrant's `TokenRateLimitPolicy` -- see `ai#127`). The +/// configuration surface may change between releases. +/// +/// Mirrors the `rules:`/`match:` shape from the `00121_token-rate-limiting` +/// proposal in `praxis-proxy/enhancements`, scoped to this milestone's +/// static header-value matchers and per-rule algorithm choice. CEL +/// matchers, soft-limit tiers, weighted per-type accounting, and +/// configurable estimation strategies are still out of scope (see the +/// module doc comment) -- upstream itself defers the first two; the +/// latter two are deferred to a separate follow-up by design, not by +/// upstream mandate. +/// +/// Assumes request identity has already been resolved upstream (this +/// filter doesn't authenticate callers) -- a catch-all rule (no +/// `match:`) reserves quota for every request that reaches it, +/// including probes and health checks. Scope rules with explicit +/// `match:` conditions, or place an identity/auth filter earlier in +/// the pipeline. Tracked as follow-on integration work in `grid#101`. +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub(super) struct TokenRateLimitConfig { + /// Evaluated in order; the first rule whose `match` is satisfied (or + /// which has no `match` at all) applies to a given request. A + /// request satisfying no rule's `match` is not rate limited by this + /// filter instance -- add a trailing rule with no `match` to enforce + /// a catch-all budget instead. + pub rules: Vec, + + /// Where every rule's admission state lives: in-process (default, + /// one budget per gateway instance) or a shared Valkey backend (one + /// budget shared across every gateway instance/replica). One + /// backend for the whole filter, not per rule -- rules already + /// share Valkey key-space isolation via `namespace`/rule-name + /// hashing, so per-rule backend selection bought no isolation + /// benefit, only a separate Valkey connection per rule pointed at + /// the same URL. Revisit if a real deployment ever needs to mix + /// in-process and Valkey rules in one filter instance. + #[serde(default)] + pub backend: BackendConfig, +} + +/// One `rules:` entry: an optional match condition, an algorithm choice +/// with that algorithm's own parameters, and this rule's own +/// estimation/keying configuration. Backend selection is shared across +/// every rule, see [`TokenRateLimitConfig::backend`]. +// `deny_unknown_fields` is deliberately omitted here: serde's flatten +// mechanism (`algorithm` below) is fundamentally incompatible with +// `deny_unknown_fields` on the containing struct -- the flattened +// enum's own fields get misreported as "unknown" because flatten +// collects remaining fields into an intermediate map before the tagged +// enum ever gets a chance to claim them. `RuleAlgorithm` itself still +// enforces `deny_unknown_fields` per-variant, so a genuinely unknown +// field (e.g. a typo, or an old flat-schema field like `window` on a +// `token_bucket` rule) is still rejected -- just attributed to the +// flattened enum's own error path instead of this struct's. +#[derive(Debug, Deserialize)] +pub(super) struct RuleConfig { + /// Human-readable rule identifier, folded into Valkey key + /// namespacing so distinct rules sharing one backend never collide. + /// + /// Renaming a live `valkey`-backed rule is therefore not a + /// no-op for operators: it changes the Valkey key hash, so the old + /// name's tracked budget is orphaned (left to expire on its own TTL) + /// and the new name starts with a fresh budget. There's no + /// migration/rename path today -- routine config hygiene (e.g. + /// renaming `"gold"` to `"gold-tier"`) silently resets that rule's + /// state. + pub name: String, + + /// Static header-value match condition. Every listed header must be + /// present on the request with an exact value match (`ANDed`) for + /// this rule to apply. Omit entirely for a catch-all rule. + #[serde(default)] + pub r#match: Option, + + /// Which admission algorithm this rule enforces, and that + /// algorithm's own parameters. + #[serde(flatten)] + pub algorithm: RuleAlgorithm, + + /// Fixed token cost reserved at admission time, before actual usage + /// is known. + /// + /// Placeholder pending M3 (configurable estimation strategies). + /// Real deployments will want this derived from request metadata + /// (e.g. `max_tokens`) rather than a single fixed constant -- that's + /// out of scope for this milestone. + pub reserved_tokens: u64, + + /// How long an admitted-but-never-reconciled reservation (lost + /// request: timeout, connection reset, upstream crash) is tracked as + /// active before that already-reserved-at-admission charge against + /// its estimate becomes irreversibly locked in (sliding-window: + /// folded into the settled total so it survives the window's normal + /// aging-out; token-bucket: the tokens were already decremented at + /// reserve time regardless, this only bounds how long the + /// reservation is tracked as pending). This does **not** defer when + /// the charge first applies -- it applies immediately at admission, + /// same as any other reservation. + /// + /// Answers the proposal's still-open "lost request handling" + /// question for this milestone. Defaults to [`DEFAULT_RESERVATION_TIMEOUT`] + /// when unset. + #[serde(default)] + pub reservation_timeout: Option, +} + +/// Static header-value match condition for a [`RuleConfig`]. +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub(super) struct MatchConfig { + /// Every header must be present on the request with this exact + /// value for the rule to match (`ANDed` across all entries). + pub headers: BTreeMap, +} + +/// Per-rule algorithm choice and its own parameters. +/// +/// Placed at the rule level (not per-budget), matching the maintainer's +/// own comparison on `ai#789`/`praxis#551` to `praxis#548`/`#856`'s +/// "per-rule" `shadow`/enforcement-action knobs. +#[derive(Debug, Deserialize)] +#[serde(tag = "algorithm", rename_all = "snake_case", deny_unknown_fields)] +pub(super) enum RuleAlgorithm { + /// Exact sliding-window admission (see [`super::ledger`]): tracks + /// usage over a continuous trailing `window`. + SlidingWindow { + /// Sliding window duration (e.g. `"1h"`, `"60s"`). + window: String, + /// Maximum tokens admitted within `window`. + capacity: u64, + }, + /// Token-bucket admission: `capacity` tokens available at once, + /// continuously refilled at `refill_rate` tokens/second. + TokenBucket { + /// Maximum tokens held at once (the bucket's ceiling). + capacity: u64, + /// Tokens refilled per second, up to `capacity`. + refill_rate: f64, + }, +} + +impl RuleAlgorithm { + /// This algorithm's configured capacity, regardless of variant. + pub(super) fn capacity(&self) -> u64 { + match self { + Self::SlidingWindow { capacity, .. } | Self::TokenBucket { capacity, .. } => *capacity, + } + } +} + +/// Default reservation timeout when `reservation_timeout` is unset. +pub(super) const DEFAULT_RESERVATION_TIMEOUT: &str = "30s"; + +/// Backend selection and connection details. +#[derive(Debug, Default, Deserialize)] +#[serde(deny_unknown_fields)] +pub(super) struct BackendConfig { + /// Which backend implementation to use. + #[serde(default)] + pub kind: BackendKind, + + /// Backend connection URL. Supports one `${ENV_VAR}` reference, so + /// credentials/hostnames don't need to be committed to config. + /// Required when `kind: valkey`, ignored otherwise. + #[serde(default)] + pub url: Option, + + /// Key namespace prefix, so multiple filter rules or deployments can + /// share one Valkey instance without colliding. Ignored for + /// `kind: memory`. Defaults to `"praxis:token_rate_limit"` when unset. + #[serde(default)] + pub namespace: Option, +} + +/// Which state backend a `token_rate_limit` rule uses. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +pub(super) enum BackendKind { + /// In-process state: fast, no extra infrastructure, but not shared + /// across gateway instances/replicas. + #[default] + Memory, + + /// Valkey-backed shared state: one budget shared across every gateway + /// instance pointed at the same `namespace`. + Valkey, +} + +#[cfg(test)] +#[expect(clippy::allow_attributes, reason = "blanket test suppressions")] +#[allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::indexing_slicing, + clippy::match_wildcard_for_single_variants, + reason = "tests intentionally fail fast on impossible fixture states" +)] +mod tests { + use super::*; + + fn parse(yaml: &str) -> Result { + serde_yaml::from_str(yaml) + } + + #[test] + fn parses_a_single_sliding_window_rule_with_no_match() { + let cfg = parse( + "rules:\n - name: default\n algorithm: sliding_window\n window: 1h\n capacity: 1000\n \ + reserved_tokens: 50\n", + ) + .unwrap(); + assert_eq!(cfg.rules.len(), 1); + let rule = &cfg.rules[0]; + assert_eq!(rule.name, "default"); + assert!(rule.r#match.is_none(), "a rule without match: is a catch-all"); + assert!(matches!( + rule.algorithm, + RuleAlgorithm::SlidingWindow { capacity: 1000, .. } + )); + assert_eq!(rule.reserved_tokens, 50); + } + + #[test] + fn parses_a_token_bucket_rule() { + let cfg = parse( + "rules:\n - name: bucket-rule\n algorithm: token_bucket\n capacity: 200\n refill_rate: 10.5\n \ + reserved_tokens: 20\n", + ) + .unwrap(); + match &cfg.rules[0].algorithm { + RuleAlgorithm::TokenBucket { capacity, refill_rate } => { + assert_eq!(*capacity, 200); + assert!((*refill_rate - 10.5).abs() < f64::EPSILON); + }, + other => panic!("expected token_bucket, got {other:?}"), + } + } + + #[test] + fn parses_multiple_rules_with_mixed_algorithms_and_header_match() { + // The customer-facing scenario this feature exists for: two apps, + // each with their own algorithm and budget, disambiguated by a + // shared header (e.g. x-app-id). + let cfg = parse( + "rules:\n\ + \x20 - name: team-alpha\n\ + \x20 match:\n\ + \x20 headers:\n\ + \x20 x-app-id: alpha\n\ + \x20 algorithm: sliding_window\n\ + \x20 window: 1h\n\ + \x20 capacity: 1000\n\ + \x20 reserved_tokens: 50\n\ + \x20 - name: team-beta\n\ + \x20 match:\n\ + \x20 headers:\n\ + \x20 x-app-id: beta\n\ + \x20 algorithm: token_bucket\n\ + \x20 capacity: 500\n\ + \x20 refill_rate: 5\n\ + \x20 reserved_tokens: 20\n", + ) + .unwrap(); + assert_eq!(cfg.rules.len(), 2); + assert_eq!(match_header(&cfg, 0, "x-app-id"), "alpha"); + assert!(matches!(cfg.rules[0].algorithm, RuleAlgorithm::SlidingWindow { .. })); + assert_eq!(match_header(&cfg, 1, "x-app-id"), "beta"); + assert!(matches!(cfg.rules[1].algorithm, RuleAlgorithm::TokenBucket { .. })); + } + + /// Fetch a header-match value off `cfg.rules[idx]` for assertions. + fn match_header<'a>(cfg: &'a TokenRateLimitConfig, idx: usize, header: &str) -> &'a str { + cfg.rules[idx].r#match.as_ref().unwrap().headers.get(header).unwrap() + } + + #[test] + fn rejects_an_empty_rules_list_shape_is_still_valid_yaml_but_filter_construction_validates_non_empty() { + // Config-level parsing accepts an empty list (YAML shape is + // valid); business-rule validation that at least one rule is + // required belongs to filter construction (`from_config`), not + // deserialization -- covered in `tests.rs`. + let cfg = parse("rules: []\n").unwrap(); + assert!(cfg.rules.is_empty()); + } + + #[test] + fn rejects_an_unknown_top_level_field() { + assert!( + parse("window: 1h\ncapacity: 100\nreserved_tokens: 5\n").is_err(), + "the old flat (pre-rules) shape must be rejected, not silently ignored" + ); + } + + #[test] + fn rejects_a_rule_missing_its_algorithm_tag() { + let err = parse("rules:\n - name: bad\n window: 1h\n capacity: 100\n reserved_tokens: 5\n") + .expect_err("should error"); + assert!(err.to_string().contains("algorithm"), "got: {err}"); + } + + #[test] + fn rejects_a_sliding_window_rule_missing_window() { + let err = parse( + "rules:\n - name: bad\n algorithm: sliding_window\n capacity: 100\n \ + reserved_tokens: 5\n", + ) + .expect_err("should error"); + assert!(err.to_string().contains("window"), "got: {err}"); + } + + #[test] + fn rejects_a_token_bucket_rule_missing_refill_rate() { + let err = + parse("rules:\n - name: bad\n algorithm: token_bucket\n capacity: 100\n reserved_tokens: 5\n") + .expect_err("should error"); + assert!(err.to_string().contains("refill_rate"), "got: {err}"); + } + + #[test] + fn rejects_mixing_sliding_window_and_token_bucket_fields_on_one_rule() { + assert!( + parse( + "rules:\n - name: bad\n algorithm: sliding_window\n window: 1h\n capacity: 100\n \ + refill_rate: 5\n reserved_tokens: 5\n" + ) + .is_err(), + "refill_rate is not a sliding_window field, deny_unknown_fields should reject it" + ); + } + + #[test] + fn algorithm_config_capacity_reads_either_variant() { + assert_eq!( + RuleAlgorithm::SlidingWindow { + window: "1h".into(), + capacity: 42 + } + .capacity(), + 42 + ); + assert_eq!( + RuleAlgorithm::TokenBucket { + capacity: 7, + refill_rate: 1.0 + } + .capacity(), + 7 + ); + } +} diff --git a/filters/src/token_rate_limit/ledger.rs b/filters/src/token_rate_limit/ledger.rs new file mode 100644 index 0000000000..8346e71edc --- /dev/null +++ b/filters/src/token_rate_limit/ledger.rs @@ -0,0 +1,767 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Praxis Contributors + +//! Exact local sliding-window reservation ledger. +//! +//! Adapted, unmodified in logic, from the `token_rate_limit::ledger` module +//! on nerdalert's `poc/distributed-token-rate-limit-demo` spike branch +//! (). +//! Replaces this filter's token-bucket state with a true sliding window per +//! the proposal's design doc ("Windows are sliding: a `window: 1h` budget +//! tracks usage in the most recent 60 minutes from the current instant"), +//! and answers that same proposal's still-open "lost request handling" +//! question via `reservation_timeout_ms` + conservative charge-at-estimate +//! on expiry. + +#![allow( + missing_docs, + clippy::missing_docs_in_private_items, + clippy::too_many_lines, + reason = "private ledger implementation is covered by its public filter contract and focused tests" +)] + +use std::{ + collections::{HashMap, VecDeque}, + sync::{ + Arc, Mutex, + atomic::{AtomicU64, AtomicUsize, Ordering}, + }, +}; + +use dashmap::DashMap; + +/// A positive rolling-window budget. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) struct Budget { + /// Window length in milliseconds. + pub(super) window_ms: u64, + /// Maximum settled plus active tokens in the window. + pub(super) capacity: u64, +} + +/// Bounds and timing for a ledger. +#[derive(Clone, Debug)] +pub(super) struct LedgerConfig { + /// All budgets in one atomic reservation rule. + pub(super) budgets: Vec, + /// Time after which an ambiguous reservation is charged at its estimate. + pub(super) reservation_timeout_ms: u64, + /// Maximum logical keys retained by the ledger. + pub(super) max_keys: usize, + /// Maximum key length retained by the ledger. + pub(super) max_key_length: usize, + /// Maximum active reservations retained by the ledger. + pub(super) max_active_reservations: usize, +} + +impl LedgerConfig { + /// Validate configuration before constructing a ledger. + pub(super) fn validate(&self) -> Result<(), String> { + if self.budgets.is_empty() { + return Err("at least one budget is required".into()); + } + if self.budgets.iter().any(|b| b.window_ms == 0 || b.capacity == 0) { + return Err("budget window and capacity must be positive".into()); + } + if self.budgets.windows(2).any(|w| { + w.first() + .zip(w.get(1)) + .is_some_and(|(left, right)| left.window_ms == right.window_ms) + }) { + return Err("budget windows must be unique".into()); + } + if self.reservation_timeout_ms == 0 { + return Err("reservation timeout must be positive".into()); + } + if self.max_keys == 0 || self.max_key_length == 0 || self.max_active_reservations == 0 { + return Err("ledger bounds must be positive".into()); + } + Ok(()) + } +} + +/// A reservation admitted atomically across every configured budget. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) struct Reservation { + /// Opaque identifier used for idempotent reconciliation. + pub(super) id: u64, + /// Estimated token cost reserved at admission. + pub(super) estimate: u64, + /// Monotonic timestamp at admission, in milliseconds. + pub(super) created_at_ms: u64, +} + +/// Result of attempting admission. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) enum Decision { + /// Request may proceed with this reservation. + Admitted(Reservation), + /// Request must be rejected before routing. + Denied { + /// Conservative delay before another admission attempt. + retry_after_ms: u64, + /// Bounded reason used for operational counters. + reason: DenialReason, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum DenialReason { + InvalidKey, + KeyCapacity, + WindowCapacity, + ReservationCapacity, +} + +/// Result of reconciling a reservation. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) enum Settlement { + /// Actual usage was applied exactly once. + Applied { + /// Actual tokens charged to the rolling ledger. + actual: u64, + /// Estimate returned to the ledger. + refund: u64, + /// Usage above the estimate. + overage: u64, + }, + /// The reservation was already reconciled or conservatively expired. + Noop, +} + +#[derive(Debug)] +struct Usage { + at_ms: u64, + tokens: u64, +} + +#[derive(Debug)] +struct ActiveReservation { + estimate: u64, + created_at_ms: u64, +} + +#[derive(Debug, Default)] +struct KeyState { + settled: VecDeque, + active: HashMap, +} + +impl KeyState { + fn reap(&mut self, now_ms: u64, config: &LedgerConfig) -> Vec { + let expired: Vec = self + .active + .iter() + .filter_map(|(id, reservation)| { + (now_ms.saturating_sub(reservation.created_at_ms) >= config.reservation_timeout_ms).then_some(*id) + }) + .collect(); + + for id in &expired { + if let Some(reservation) = self.active.remove(id) { + // An ambiguous request is never free traffic. Charge the + // estimate at admission time so the normal window expiry + // rules still apply. + self.settled.push_back(Usage { + at_ms: reservation.created_at_ms, + tokens: reservation.estimate, + }); + } + } + + let max_window = config.budgets.iter().map(|b| b.window_ms).max().unwrap_or(0); + while self + .settled + .front() + .is_some_and(|entry| now_ms.saturating_sub(entry.at_ms) >= max_window) + { + self.settled.pop_front(); + } + expired + } + + fn usage_in_window(&self, now_ms: u64, window_ms: u64) -> u64 { + let settled = self + .settled + .iter() + .filter(|entry| now_ms.saturating_sub(entry.at_ms) < window_ms) + .fold(0_u64, |sum, entry| sum.saturating_add(entry.tokens)); + let active = self + .active + .values() + .fold(0_u64, |sum, reservation| sum.saturating_add(reservation.estimate)); + settled.saturating_add(active) + } + + fn retry_after_ms(&self, now_ms: u64, config: &LedgerConfig) -> u64 { + config + .budgets + .iter() + .flat_map(|budget| { + self.settled.iter().filter_map(move |entry| { + let expiry = entry.at_ms.saturating_add(budget.window_ms); + (expiry > now_ms).then_some(expiry - now_ms) + }) + }) + .max() + .unwrap_or(config.reservation_timeout_ms) + .max(config.reservation_timeout_ms) + } + + fn is_empty(&self) -> bool { + self.active.is_empty() && self.settled.is_empty() + } +} + +/// Thread-safe exact local ledger with independent locks per key. +pub(super) struct Ledger { + config: LedgerConfig, + keys: DashMap>>, + reservations: DashMap, + next_id: AtomicU64, + key_count: AtomicUsize, + active_reservations: AtomicUsize, +} + +impl Ledger { + /// Construct a validated empty ledger. + pub(super) fn new(config: LedgerConfig) -> Result { + config.validate()?; + Ok(Self { + config, + keys: DashMap::new(), + reservations: DashMap::new(), + next_id: AtomicU64::new(1), + key_count: AtomicUsize::new(0), + active_reservations: AtomicUsize::new(0), + }) + } + + /// Return the smallest configured capacity for bounded quota headers. + pub(super) fn limit(&self) -> u64 { + self.config + .budgets + .iter() + .map(|budget| budget.capacity) + .min() + .unwrap_or(0) + } + + /// Current number of active reservations. + pub(super) fn active_count(&self) -> usize { + self.active_reservations.load(Ordering::Relaxed) + } + + /// Current number of retained logical keys. + pub(super) fn key_count(&self) -> usize { + self.key_count.load(Ordering::Relaxed) + } + + /// Reserve an estimate atomically across all configured windows. + pub(super) fn reserve(&self, key: &str, estimate: u64, now_ms: u64) -> Decision { + if key.is_empty() || key.len() > self.config.max_key_length || estimate == 0 { + return Decision::Denied { + retry_after_ms: 0, + reason: DenialReason::InvalidKey, + }; + } + + let state = match self.keys.entry(key.to_owned()) { + dashmap::mapref::entry::Entry::Occupied(entry) => Arc::clone(entry.get()), + dashmap::mapref::entry::Entry::Vacant(entry) => { + if self + .key_count + .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |count| { + (count < self.config.max_keys).then_some(count + 1) + }) + .is_err() + { + return Decision::Denied { + retry_after_ms: 0, + reason: DenialReason::KeyCapacity, + }; + } + let state = Arc::new(Mutex::new(KeyState::default())); + entry.insert(Arc::clone(&state)); + state + }, + }; + let mut state = match state.lock() { + Ok(state) => state, + Err(poisoned) => poisoned.into_inner(), + }; + let expired = state.reap(now_ms, &self.config); + for id in &expired { + self.reservations.remove(id); + } + self.active_reservations.fetch_sub(expired.len(), Ordering::Relaxed); + + if self + .config + .budgets + .iter() + .any(|budget| state.usage_in_window(now_ms, budget.window_ms).saturating_add(estimate) > budget.capacity) + { + return Decision::Denied { + retry_after_ms: state.retry_after_ms(now_ms, &self.config), + reason: DenialReason::WindowCapacity, + }; + } + if self + .active_reservations + .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |active| { + (active < self.config.max_active_reservations).then_some(active + 1) + }) + .is_err() + { + return Decision::Denied { + retry_after_ms: self.config.reservation_timeout_ms, + reason: DenialReason::ReservationCapacity, + }; + } + + let id = self.next_id.fetch_add(1, Ordering::Relaxed); + state.active.insert( + id, + ActiveReservation { + estimate, + created_at_ms: now_ms, + }, + ); + self.reservations.insert(id, key.to_owned()); + drop(state); + Decision::Admitted(Reservation { + id, + estimate, + created_at_ms: now_ms, + }) + } + + /// Reconcile actual usage. Repeated calls for one ID are no-ops. + pub(super) fn reconcile(&self, id: u64, actual: Option, now_ms: u64) -> Settlement { + let Some((_, key)) = self.reservations.remove(&id) else { + return Settlement::Noop; + }; + let Some(state) = self.keys.get(&key).map(|entry| Arc::clone(entry.value())) else { + return Settlement::Noop; + }; + let mut state = match state.lock() { + Ok(state) => state, + Err(poisoned) => poisoned.into_inner(), + }; + let Some(reservation) = state.active.remove(&id) else { + return Settlement::Noop; + }; + self.active_reservations.fetch_sub(1, Ordering::Relaxed); + let actual = actual.unwrap_or(reservation.estimate); + state.settled.push_back(Usage { + at_ms: reservation.created_at_ms, + tokens: actual, + }); + let refund = reservation.estimate.saturating_sub(actual); + let overage = actual.saturating_sub(reservation.estimate); + let expired = state.reap(now_ms, &self.config); + for expired_id in expired { + self.reservations.remove(&expired_id); + self.active_reservations.fetch_sub(1, Ordering::Relaxed); + } + drop(state); + Settlement::Applied { + actual, + refund, + overage, + } + } + + /// Conservatively expire a bounded number of keys and reclaim idle state. + pub(super) fn cleanup(&self, now_ms: u64, max_keys_to_scan: usize) -> usize { + let mut orphaned = 0; + let keys: Vec = self + .keys + .iter() + .take(max_keys_to_scan) + .map(|entry| entry.key().clone()) + .collect(); + for key in keys { + let Some(entry) = self.keys.get_mut(&key) else { + continue; + }; + let state_arc = Arc::clone(entry.value()); + let mut state = match state_arc.lock() { + Ok(state) => state, + Err(poisoned) => poisoned.into_inner(), + }; + let expired = state.reap(now_ms, &self.config); + orphaned += expired.len(); + for id in &expired { + self.reservations.remove(id); + } + self.active_reservations.fetch_sub(expired.len(), Ordering::Relaxed); + let empty = state.is_empty(); + drop(state); + drop(entry); + if empty + && self + .keys + .remove_if(&key, |_, candidate| { + Arc::ptr_eq(candidate, &state_arc) + && match candidate.lock() { + Ok(state) => state.is_empty(), + Err(poisoned) => poisoned.into_inner().is_empty(), + } + }) + .is_some() + { + self.key_count.fetch_sub(1, Ordering::Relaxed); + } + } + orphaned + } +} + +#[cfg(test)] +#[expect(clippy::allow_attributes, reason = "blanket test suppressions")] +#[allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::indexing_slicing, + clippy::panic, + clippy::manual_let_else, + clippy::match_wildcard_for_single_variants, + reason = "ledger tests intentionally fail fast on impossible fixture states" +)] +mod tests { + use std::sync::{Arc, Barrier}; + + use super::*; + + fn ledger(budgets: &[(u64, u64)]) -> Ledger { + Ledger::new(LedgerConfig { + budgets: budgets + .iter() + .map(|&(window_ms, capacity)| Budget { window_ms, capacity }) + .collect(), + reservation_timeout_ms: 100, + max_keys: 8, + max_key_length: 256, + max_active_reservations: 32, + }) + .unwrap() + } + + #[test] + fn admits_and_denies_one_window() { + let l = ledger(&[(60_000, 10)]); + assert!(matches!(l.reserve("a", 10, 0), Decision::Admitted(_))); + assert!(matches!(l.reserve("a", 1, 0), Decision::Denied { .. })); + } + + #[test] + fn concurrent_same_key_admission_cannot_oversubscribe() { + let ledger = Arc::new(ledger(&[(1_000, 100)])); + let barrier = Arc::new(Barrier::new(16)); + let handles = (0..16) + .map(|_| { + let ledger = Arc::clone(&ledger); + let barrier = Arc::clone(&barrier); + std::thread::spawn(move || { + barrier.wait(); + matches!(ledger.reserve("same", 10, 0), Decision::Admitted(_)) + }) + }) + .collect::>(); + let admitted = handles + .into_iter() + .filter_map(|handle| handle.join().ok()) + .filter(|ok| *ok) + .count(); + assert_eq!(admitted, 10, "exactly the capacity should be admitted"); + } + + #[test] + fn concurrent_different_keys_respect_global_reservation_bound() { + let ledger = Arc::new( + Ledger::new(LedgerConfig { + budgets: vec![Budget { + window_ms: 1_000, + capacity: 1_000, + }], + reservation_timeout_ms: 100, + max_keys: 32, + max_key_length: 256, + max_active_reservations: 4, + }) + .unwrap(), + ); + let barrier = Arc::new(Barrier::new(16)); + let handles = (0..16) + .map(|index| { + let ledger = Arc::clone(&ledger); + let barrier = Arc::clone(&barrier); + std::thread::spawn(move || { + barrier.wait(); + matches!(ledger.reserve(&format!("key-{index}"), 1, 0), Decision::Admitted(_)) + }) + }) + .collect::>(); + let admitted = handles + .into_iter() + .filter_map(|handle| handle.join().ok()) + .filter(|ok| *ok) + .count(); + assert_eq!(admitted, 4, "the global reservation bound must be atomic"); + assert_eq!(ledger.active_count(), 4); + } + + #[test] + fn exact_boundary_expires_usage() { + let l = ledger(&[(100, 10)]); + let r = match l.reserve("a", 10, 0) { + Decision::Admitted(r) => r, + _ => panic!(), + }; + assert!(matches!(l.reconcile(r.id, Some(10), 0), Settlement::Applied { .. })); + assert!(matches!(l.reserve("a", 10, 100), Decision::Admitted(_))); + } + + #[test] + fn multiple_windows_are_atomic() { + let l = ledger(&[(100, 10), (1_000, 15)]); + let r = match l.reserve("a", 10, 0) { + Decision::Admitted(r) => r, + _ => panic!(), + }; + l.reconcile(r.id, Some(10), 0); + assert!(matches!(l.reserve("a", 6, 0), Decision::Denied { .. })); + assert!(matches!(l.reserve("b", 6, 0), Decision::Admitted(_))); + } + + #[test] + fn active_reservations_count_and_orphans_are_charged() { + let l = ledger(&[(1_000, 10)]); + assert!(matches!(l.reserve("a", 10, 0), Decision::Admitted(_))); + assert_eq!(l.active_count(), 1); + assert!(matches!(l.reserve("a", 1, 0), Decision::Denied { .. })); + l.cleanup(100, 8); + assert_eq!(l.active_count(), 0); + assert!(matches!(l.reserve("a", 1, 100), Decision::Denied { .. })); + } + + #[test] + fn idle_settled_keys_are_reclaimed_atomically() { + let l = ledger(&[(100, 10)]); + let r = match l.reserve("a", 10, 0) { + Decision::Admitted(r) => r, + _ => panic!(), + }; + assert!(matches!(l.reconcile(r.id, Some(10), 0), Settlement::Applied { .. })); + assert_eq!(l.key_count(), 1); + l.cleanup(100, 8); + assert_eq!(l.key_count(), 0); + } + + #[test] + fn refund_exact_and_overage_are_recorded() { + let l = ledger(&[(1_000, 100)]); + let r = match l.reserve("a", 50, 0) { + Decision::Admitted(r) => r, + _ => panic!(), + }; + assert_eq!( + l.reconcile(r.id, Some(20), 0), + Settlement::Applied { + actual: 20, + refund: 30, + overage: 0 + } + ); + let r = match l.reserve("a", 50, 0) { + Decision::Admitted(r) => r, + _ => panic!(), + }; + assert_eq!( + l.reconcile(r.id, Some(70), 0), + Settlement::Applied { + actual: 70, + refund: 0, + overage: 20 + } + ); + } + + #[test] + fn duplicate_reconciliation_is_noop() { + let l = ledger(&[(1_000, 100)]); + let r = match l.reserve("a", 5, 0) { + Decision::Admitted(r) => r, + _ => panic!(), + }; + assert!(matches!( + l.reconcile(r.id, None, 0), + Settlement::Applied { actual: 5, .. } + )); + assert_eq!(l.reconcile(r.id, Some(99), 0), Settlement::Noop); + } + + #[test] + fn keys_are_independent_and_bounded() { + let l = ledger(&[(1_000, 10)]); + assert!(matches!(l.reserve("a", 10, 0), Decision::Admitted(_))); + assert!(matches!(l.reserve("b", 10, 0), Decision::Admitted(_))); + assert!(matches!(l.reserve("", 1, 0), Decision::Denied { .. })); + } + + #[test] + fn invalid_config_is_rejected() { + assert!( + Ledger::new(LedgerConfig { + budgets: vec![], + reservation_timeout_ms: 1, + max_keys: 1, + max_key_length: 1, + max_active_reservations: 1 + }) + .is_err() + ); + assert!( + Ledger::new(LedgerConfig { + budgets: vec![Budget { + window_ms: 0, + capacity: 1 + }], + reservation_timeout_ms: 1, + max_keys: 1, + max_key_length: 1, + max_active_reservations: 1 + }) + .is_err() + ); + } + + #[test] + fn duplicate_budget_windows_are_rejected() { + assert!( + Ledger::new(LedgerConfig { + budgets: vec![ + Budget { + window_ms: 1_000, + capacity: 10 + }, + Budget { + window_ms: 1_000, + capacity: 20 + }, + ], + reservation_timeout_ms: 1, + max_keys: 1, + max_key_length: 1, + max_active_reservations: 1 + }) + .is_err(), + "two budgets sharing the same window are ambiguous" + ); + } + + #[test] + fn zero_reservation_timeout_is_rejected() { + assert!( + Ledger::new(LedgerConfig { + budgets: vec![Budget { + window_ms: 1_000, + capacity: 10 + }], + reservation_timeout_ms: 0, + max_keys: 1, + max_key_length: 1, + max_active_reservations: 1 + }) + .is_err() + ); + } + + #[test] + fn zero_ledger_bounds_are_rejected() { + for (max_keys, max_key_length, max_active_reservations) in [(0, 1, 1), (1, 0, 1), (1, 1, 0)] { + assert!( + Ledger::new(LedgerConfig { + budgets: vec![Budget { + window_ms: 1_000, + capacity: 10 + }], + reservation_timeout_ms: 1, + max_keys, + max_key_length, + max_active_reservations + }) + .is_err(), + "bounds ({max_keys}, {max_key_length}, {max_active_reservations}) must be rejected" + ); + } + } + + #[test] + fn key_capacity_denies_new_keys_beyond_the_configured_limit() { + let l = Ledger::new(LedgerConfig { + budgets: vec![Budget { + window_ms: 1_000, + capacity: 100, + }], + reservation_timeout_ms: 100, + max_keys: 1, + max_key_length: 256, + max_active_reservations: 32, + }) + .unwrap(); + assert!(matches!(l.reserve("a", 1, 0), Decision::Admitted(_))); + assert!(matches!( + l.reserve("b", 1, 0), + Decision::Denied { + reason: DenialReason::KeyCapacity, + .. + } + )); + } + + #[test] + fn expired_active_reservation_is_reaped_on_next_reserve_for_the_same_key() { + let l = ledger(&[(100, 10)]); + let r = match l.reserve("a", 10, 0) { + Decision::Admitted(r) => r, + other => panic!("expected admission, got {other:?}"), + }; + assert_eq!(l.active_count(), 1); + // Never reconciled, and now well past both the window and the + // reservation timeout -- the next reserve for "a" must reap the + // stale reservation before evaluating capacity, not deny it as + // still active. + assert!(matches!(l.reserve("a", 10, 200), Decision::Admitted(_))); + assert_eq!(l.active_count(), 1, "old reservation reaped, new one admitted"); + assert_eq!( + l.reconcile(r.id, Some(1), 200), + Settlement::Noop, + "the reaped id was actually dropped, not just shadowed" + ); + } + + #[test] + fn expired_sibling_reservation_is_reaped_during_reconcile() { + let l = ledger(&[(1_000, 10)]); + let stale = match l.reserve("a", 1, 0) { + Decision::Admitted(r) => r, + other => panic!("expected admission, got {other:?}"), + }; + let fresh = match l.reserve("a", 1, 50) { + Decision::Admitted(r) => r, + other => panic!("expected admission, got {other:?}"), + }; + assert_eq!(l.active_count(), 2); + // Reconciling `fresh` at t=150 (>= reservation_timeout_ms=100 past + // `stale`'s creation) must also reap `stale` as a side effect. + assert!(matches!( + l.reconcile(fresh.id, Some(1), 150), + Settlement::Applied { .. } + )); + assert_eq!(l.active_count(), 0, "the stale sibling must be reaped too"); + assert_eq!(l.reconcile(stale.id, Some(1), 150), Settlement::Noop); + } +} diff --git a/filters/src/token_rate_limit/mod.rs b/filters/src/token_rate_limit/mod.rs new file mode 100644 index 0000000000..c8c37be15b --- /dev/null +++ b/filters/src/token_rate_limit/mod.rs @@ -0,0 +1,952 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Praxis Contributors + +//! Token-denominated rate limiting filter. +//! +//! **Experimental.** Requires the `token-rate-limit-filter` cargo +//! feature, which is off by default and activates the `experimental` +//! marker. This filter delivers the epic's agreed M1/M2/M6 scope (see +//! below), but its parent proposal (`00121_token-rate-limiting.md`) is +//! not yet `accepted`, and open questions remain: HA/clustered-Valkey +//! failure modes, and this filter's relationship to Kuadrant's +//! `TokenRateLimitPolicy` (a separate, already-shipped mechanism for +//! the same problem -- see `ai#127`). The configuration surface may +//! change between releases. Anything beyond the agreed M1/M2/M6 scope +//! belongs in `praxis-proxy/experimental` first, not here -- see +//! `grid#101`. +//! +//! Implements the agreed M1/M2/M6 core of the token rate limiting +//! proposal (`00121_token-rate-limiting.md` in `praxis-proxy/enhancements`, +//! tracked by epic `ai#121`): an ordered list of `rules`, each an optional +//! static header-value match condition ("static matchers") +//! bound to an independently-chosen admission algorithm and its own +//! token budget, reservation-based admission reconciled against actual +//! provider-reported usage, and standard 429 responses with +//! token-denominated rate limit headers. +//! +//! Two algorithms are supported per rule, chosen via `algorithm:`: +//! +//! - `sliding_window` (see [`ledger`]): exact sliding-window admission, adapted from nerdalert's +//! `poc/distributed-token-rate-limit-demo` spike branch. Tracks usage over a continuous trailing `window`. +//! - `token_bucket` (see [`token_bucket_ledger`]): continuous refill up to `capacity` at `refill_rate` tokens/second, +//! reusing the refill formula from Praxis's own lock-free `traffic_management::token_bucket`, extended with the +//! reserve/reconcile split this filter needs. +//! +//! Both sit behind the same pluggable [`backend`] trait, so either +//! algorithm runs in-process (default) or against a shared Valkey +//! backend (`backend: {kind: valkey}`) for state shared across gateway +//! instances/replicas -- see `praxis-proxy/grid#83` for the fuller +//! Valkey-backend spec this milestone is a narrower slice of. +//! +//! Per-rule algorithm choice, rather than one fixed algorithm for the +//! whole filter, mirrors the `GuardrailsFilter`'s own `rules: Vec` +//! architecture and answers the maintainer's own framing on +//! `ai#789`/`praxis#551` ("this looks like a per-rule choice, similar to +//! `shadow`/enforcement-action knobs elsewhere"). +//! +//! Deliberately deferred, pending the proposal's own open design questions: +//! +//! - **CEL-expression matchers**: only static, exact header-value equality matching is implemented (overlapping +//! `praxis#189`/`#232`). +//! - **Composite/multi-dimension keys, per-model keys**: flagged as TBD under the proposal's own M5 goal (see +//! `ai#123`/`ai#232`); `ai#129`'s single-header-value keying (one budget applied uniformly per key, fallback to +//! global) is implemented per rule, and intentionally does not resolve identity to a key itself -- it keys off +//! whatever header value an upstream component has already put there. +//! - **Configurable estimation (M3)**: `reserved_tokens` is a fixed constant per rule, not derived from request +//! metadata (e.g. `max_tokens`). +//! - **Token-type-aware accounting (M4)**: reconciles against `token.total` only; per-type (input/output/cached) +//! weighting is not modeled yet. +//! - **Multiple budgets per rule, soft-limit tiers**: the proposal allows several `token_budgets` (e.g. hourly + daily) +//! and graduated tiers per rule; this milestone admits exactly one budget per rule with a hard deny at capacity. +//! - **Observability (M7/M8) and metering (S3)**: out of scope here -- both are recommended to split into their own +//! follow-on proposals. +//! - **Trust boundary, non-inference traffic scoping**: this filter assumes request identity has already been resolved +//! upstream (the proposal's own Non-Goals) and does not itself authenticate callers or exempt probes/health +//! checks/malformed requests from a catch-all rule's reservation. Scope rules with explicit `match:` conditions, or +//! place an identity/auth filter earlier in the pipeline. Tracked as follow-on integration work in `grid#101`. +//! +//! `X-RateLimit-*` headers are emitted on 429 rejection only, matching +//! the validated pattern on the source spike branch: computing +//! "remaining" for every successful response would need an extra read +//! on the Valkey path (doubling backend round-trips per request), so +//! both backends behave identically here rather than diverging by +//! backend. +//! +//! Depends on `token_count` running earlier in the response phase to +//! populate `token.total` in `filter_metadata`; if that metadata is +//! absent when the response completes, the reservation is left as +//! final rather than guessed at. + +#[cfg(test)] +#[expect(clippy::allow_attributes, reason = "blanket test suppressions")] +#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic, reason = "tests")] +mod tests; + +mod backend; +mod config; +mod ledger; +mod token_bucket_ledger; + +use std::{collections::HashSet, sync::Arc, time::Instant}; + +use async_trait::async_trait; +use bytes::Bytes; +use http::header::HeaderName; +use metrics::{counter, gauge}; +use praxis_filter::{ + BodyAccess, BodyMode, FilterAction, FilterError, HttpFilter, HttpFilterContext, Rejection, parse_filter_config, +}; + +use self::{ + backend::{ + BackendError, BackendReserve, BackendSettlement, CleanupReport, InMemoryTokenBucketBackend, + InMemoryTokenRateLimitBackend, ReconcileRequest, ReserveRequest, TokenRateLimitStateBackend, + ValkeyBackendConfig, ValkeyEval, ValkeyTokenBucketBackend, ValkeyTokenBucketConfig, + ValkeyTokenRateLimitBackend, + }, + config::{ + BackendConfig, BackendKind, DEFAULT_RESERVATION_TIMEOUT, MatchConfig, RuleAlgorithm, RuleConfig, + TokenRateLimitConfig, + }, + ledger::{Budget, Ledger, LedgerConfig}, + token_bucket_ledger::{TokenBucketConfig, TokenBucketLedger}, +}; +use crate::token_usage::META_TOKEN_TOTAL; + +// ----------------------------------------------------------------------------- +// Constants +// ----------------------------------------------------------------------------- + +/// Metadata key stashing this request's reservation ID, read back during +/// response-phase reconciliation. +const META_RESERVATION_ID: &str = "token_rate_limit.reservation_id"; + +/// Metadata key stashing the resolved bucket key, read back in +/// reconciliation so it operates on the same budget `on_request` reserved +/// from. +const META_BUCKET_KEY: &str = "token_rate_limit.bucket_key"; + +/// Metadata key stashing the index of the [`CompiledRule`] that admitted +/// this request, so reconciliation settles against the same rule's +/// backend/estimate even when other rules exist. +const META_RULE_INDEX: &str = "token_rate_limit.rule_index"; + +/// The single budget key every request resolves to in this milestone: all +/// requests matching a rule share one budget. Kept as a named sentinel +/// (rather than threading `Option`/`&str` through the backend +/// APIs) so a future per-request keying mechanism (M5, deliberately out +/// of scope here -- see the proposal's open question and `ai#790`'s +/// quota-key design) can slot in without changing the backend trait. +const FALLBACK_KEY: &str = "__fallback__"; + +/// Bound on distinct budget keys retained at once, per rule. +/// +/// Always `1` in this milestone (just [`FALLBACK_KEY`]); sized for future +/// per-request keying, mirroring the soft cap `rate_limit` uses for +/// per-IP entries. +const MAX_KEYS: usize = 100_000; + +/// Bound on a single budget key's length. +const MAX_KEY_LENGTH: usize = 256; + +/// Bound on reservations awaiting reconciliation across all keys, per rule. +const MAX_ACTIVE_RESERVATIONS: usize = 200_000; + +/// Rate limit header: configured token budget. +/// +/// Uses the `-Tokens` suffix per `ai#124`'s spec, distinct from the +/// existing `rate_limit` filter's unsuffixed `X-RateLimit-Limit`, to +/// avoid a header collision when both filters run in the same +/// pipeline. +const HEADER_RATELIMIT_LIMIT_TOKENS: &str = "X-RateLimit-Limit-Tokens"; + +/// Rate limit header: remaining tokens (always `0` -- only sent on 429). +const HEADER_RATELIMIT_REMAINING_TOKENS: &str = "X-RateLimit-Remaining-Tokens"; + +/// Rate limit header: seconds until another admission attempt may succeed. +const HEADER_RATELIMIT_RESET: &str = "X-RateLimit-Reset-Tokens"; + +/// Resolved, ready-to-use form of the filter-level `backend:` config: +/// either every rule uses in-process state, or every rule shares one +/// already-open Valkey connection (see [`ValkeyEval`]'s doc comment for +/// why this is built once and `Clone`d, not once per rule). +enum BackendResource { + /// Every rule gets its own in-process ledger (the default). + Memory, + /// Every rule shares this one Valkey connection, differentiated by + /// `namespace`/rule-name key hashing. Boxed: `ValkeyEval` embeds a + /// `redis::Client`/`ConnectionInfo`, large enough that an unboxed + /// field here would size the whole enum (including the zero-data + /// `Memory` variant) up to match it. + Valkey { + /// Filter-shared connection, `Clone`d into each Valkey-backed + /// rule's own backend. + valkey: Box, + /// Key namespace prefix, see [`BackendConfig::namespace`]. + namespace: String, + }, +} + +/// Resolve the filter's `backend:` block into a [`BackendResource`], +/// opening the Valkey connection once up front if configured. +/// +/// # Errors +/// +/// Returns [`FilterError`] if `backend.kind: valkey` is set without a +/// `url`, or the URL fails to parse/expand. +fn build_backend_resource(backend: &BackendConfig) -> Result { + match backend.kind { + BackendKind::Memory => Ok(BackendResource::Memory), + BackendKind::Valkey => { + let url = backend + .url + .as_deref() + .ok_or("token_rate_limit: backend.url is required for backend.kind: valkey")?; + let url = expand_backend_url(url)?; + let namespace = backend + .namespace + .clone() + .unwrap_or_else(|| "praxis:token_rate_limit".to_owned()); + let valkey = Box::new(ValkeyEval::new(url)?); + Ok(BackendResource::Valkey { valkey, namespace }) + }, + } +} + +/// Build the configured state backend for a `sliding_window` rule +/// (in-process ledger, or a handle onto the filter's shared Valkey +/// connection). +/// +/// # Errors +/// +/// Returns [`FilterError`] if the ledger config is invalid. +fn build_sliding_window_backend( + backend: &BackendResource, + rule_name: &str, + budgets: Vec, + reservation_timeout_ms: u64, +) -> Result, FilterError> { + match backend { + BackendResource::Memory => { + let ledger = Ledger::new(LedgerConfig { + budgets, + reservation_timeout_ms, + max_keys: MAX_KEYS, + max_key_length: MAX_KEY_LENGTH, + max_active_reservations: MAX_ACTIVE_RESERVATIONS, + }) + .map_err(|error| format!("token_rate_limit: rule '{rule_name}': {error}"))?; + Ok(Arc::new(InMemoryTokenRateLimitBackend::new(ledger))) + }, + BackendResource::Valkey { valkey, namespace } => { + Ok(Arc::new(ValkeyTokenRateLimitBackend::new(ValkeyBackendConfig { + valkey: (**valkey).clone(), + namespace: namespace.clone(), + rule: rule_name.to_owned(), + budgets, + reservation_timeout_ms, + max_keys: MAX_KEYS, + max_active_reservations: MAX_ACTIVE_RESERVATIONS, + }))) + }, + } +} + +/// Build the configured state backend for a `token_bucket` rule +/// (in-process ledger, or a handle onto the filter's shared Valkey +/// connection). +/// +/// # Errors +/// +/// Returns [`FilterError`] if the ledger config is invalid. +fn build_token_bucket_backend( + backend: &BackendResource, + rule_name: &str, + capacity: u64, + refill_rate: f64, + reservation_timeout_ms: u64, +) -> Result, FilterError> { + match backend { + BackendResource::Memory => { + let ledger = TokenBucketLedger::new(TokenBucketConfig { + capacity, + refill_rate, + reservation_timeout_ms, + max_keys: MAX_KEYS, + max_key_length: MAX_KEY_LENGTH, + max_active_reservations: MAX_ACTIVE_RESERVATIONS, + }) + .map_err(|error| format!("token_rate_limit: rule '{rule_name}': {error}"))?; + Ok(Arc::new(InMemoryTokenBucketBackend::new(ledger))) + }, + BackendResource::Valkey { valkey, namespace } => { + Ok(Arc::new(ValkeyTokenBucketBackend::new(ValkeyTokenBucketConfig { + valkey: (**valkey).clone(), + namespace: namespace.clone(), + rule: rule_name.to_owned(), + capacity, + refill_rate, + reservation_timeout_ms, + max_keys: MAX_KEYS, + max_active_reservations: MAX_ACTIVE_RESERVATIONS, + })?)) + }, + } +} + +// ----------------------------------------------------------------------------- +// CompiledRule +// ----------------------------------------------------------------------------- + +/// One `rules:` entry, fully resolved: its match condition (if any), +/// backend, and estimation. +struct CompiledRule { + /// Human-readable identifier, used in metrics labels and error + /// messages, and folded into Valkey key namespacing. + name: String, + + /// Static header-value match condition. `None` matches every + /// request unconditionally (a catch-all rule). + matcher: Option, + + /// This rule's own admission state: in-process, or shared via + /// Valkey; sliding-window or token-bucket. + backend: Arc, + + /// Fixed token cost reserved at admission (M3 placeholder). + reserved_tokens: u64, +} + +impl CompiledRule { + /// Whether `headers` satisfies this rule's match condition (or the + /// rule is a catch-all with no condition at all). + fn matches(&self, headers: &http::HeaderMap) -> bool { + self.matcher.as_deref().is_none_or(|conditions| { + conditions + .iter() + .all(|(name, value)| headers.get(name).and_then(|v| v.to_str().ok()) == Some(value.as_str())) + }) + } +} + +/// Validate a rule's `capacity`/`reserved_tokens` bounds and resolve its +/// `reservation_timeout` string to milliseconds. +/// +/// # Errors +/// +/// Returns [`FilterError`] if `capacity` is zero, exceeds the Lua +/// `f64` safe-integer bound, `reserved_tokens` is zero or exceeds +/// `capacity`, or `reservation_timeout` isn't a valid duration. +fn validate_rule_bounds(rule: &RuleConfig, capacity: u64) -> Result { + validate_capacity_safe_integer_bound(&rule.name, capacity)?; + if rule.reserved_tokens == 0 { + return Err(format!( + "token_rate_limit: rule '{}': reserved_tokens must be greater than 0", + rule.name + ) + .into()); + } + if rule.reserved_tokens > capacity { + return Err(format!( + "token_rate_limit: rule '{}': reserved_tokens must not exceed capacity", + rule.name + ) + .into()); + } + parse_duration_ms( + rule.reservation_timeout + .as_deref() + .unwrap_or(DEFAULT_RESERVATION_TIMEOUT), + ) +} + +/// Reject a zero `capacity`, or one beyond the Lua `f64` safe-integer +/// bound. +/// +/// `token_bucket_ledger` re-checks this same bound on its own +/// construction path (see `MAX_F64_SAFE_INTEGER`'s doc comment), but +/// `ledger` (`sliding_window`) has no such gate downstream -- checking +/// it here, before either algorithm's backend is built, closes that +/// gap for both. +fn validate_capacity_safe_integer_bound(rule_name: &str, capacity: u64) -> Result<(), FilterError> { + if capacity == 0 { + return Err(format!("token_rate_limit: rule '{rule_name}': capacity must be greater than 0").into()); + } + if capacity > token_bucket_ledger::MAX_F64_SAFE_INTEGER { + return Err(format!( + "token_rate_limit: rule '{rule_name}': capacity must not exceed {} (2^53)", + token_bucket_ledger::MAX_F64_SAFE_INTEGER + ) + .into()); + } + Ok(()) +} + +/// A compiled `match: {headers: ...}` condition: an ordered list of +/// header-name/expected-value pairs that must all match (see +/// [`CompiledRule::matches`]). +type HeaderMatchers = Vec<(HeaderName, String)>; + +/// Parse a rule's optional `match: {headers: ...}` block into concrete +/// [`HeaderName`]s, validating each header name along the way. +/// +/// # Errors +/// +/// Returns [`FilterError`] if any header name is invalid. +fn compile_matcher(rule_name: &str, r#match: Option) -> Result, FilterError> { + r#match + .map(|m| { + m.headers + .into_iter() + .map(|(name, value)| { + HeaderName::try_from(name.as_str()) + .map(|header_name| (header_name, value)) + .map_err(|error| { + FilterError::from(format!( + "token_rate_limit: rule '{rule_name}': invalid match header '{name}': {error}" + )) + }) + }) + .collect::, FilterError>>() + }) + .transpose() +} + +/// Build this rule's state backend per its chosen `algorithm:` variant. +/// +/// # Errors +/// +/// See [`build_sliding_window_backend`]/[`build_token_bucket_backend`]. +fn build_rule_backend( + algorithm: &RuleAlgorithm, + backend: &BackendResource, + rule_name: &str, + reservation_timeout_ms: u64, +) -> Result, FilterError> { + match algorithm { + RuleAlgorithm::SlidingWindow { window, capacity } => { + let window_ms = parse_duration_ms(window)?; + let budgets = vec![Budget { + window_ms, + capacity: *capacity, + }]; + build_sliding_window_backend(backend, rule_name, budgets, reservation_timeout_ms) + }, + RuleAlgorithm::TokenBucket { capacity, refill_rate } => { + build_token_bucket_backend(backend, rule_name, *capacity, *refill_rate, reservation_timeout_ms) + }, + } +} + +/// Compile one YAML `rules:` entry into a [`CompiledRule`], validating +/// and constructing its backend. +/// +/// # Errors +/// +/// Returns [`FilterError`] if `capacity` is zero, `reserved_tokens` is +/// zero or exceeds `capacity`, `window`/`reservation_timeout` aren't +/// valid durations, a `match` header name is invalid, or the rule's +/// backend fails to construct (see [`build_rule_backend`]). +fn compile_rule(rule: RuleConfig, backend: &BackendResource) -> Result { + let capacity = rule.algorithm.capacity(); + let reservation_timeout_ms = validate_rule_bounds(&rule, capacity)?; + let backend = build_rule_backend(&rule.algorithm, backend, &rule.name, reservation_timeout_ms)?; + let matcher = compile_matcher(&rule.name, rule.r#match)?; + + Ok(CompiledRule { + name: rule.name, + matcher, + backend, + reserved_tokens: rule.reserved_tokens, + }) +} + +// ----------------------------------------------------------------------------- +// TokenRateLimitFilter +// ----------------------------------------------------------------------------- + +/// Token-denominated rate limiter: reserves an estimated cost at +/// admission, reconciles against actual usage after the response +/// completes. Evaluates an ordered list of rules, each with its own +/// optional match condition, algorithm choice, and budget. +/// +/// # YAML configuration +/// +/// ```yaml +/// filter: token_rate_limit +/// backend: # optional: defaults to in-process state, shared by every rule +/// kind: valkey # memory (default) | valkey +/// url: "${TOKEN_RATE_LIMIT_VALKEY_URL}" +/// namespace: praxis:token_rate_limit +/// rules: +/// - name: team-alpha # human-readable, unique per filter instance +/// match: # optional: omit for a catch-all rule +/// headers: +/// x-app-id: alpha +/// algorithm: sliding_window # sliding_window | token_bucket +/// window: 1h # sliding_window only: window duration +/// capacity: 100000 # max tokens admitted (sliding_window) or held (token_bucket) +/// reserved_tokens: 500 # fixed cost reserved per request at admission +/// - name: team-beta +/// match: +/// headers: +/// x-app-id: beta +/// algorithm: token_bucket +/// capacity: 50000 +/// refill_rate: 50 # token_bucket only: tokens refilled per second +/// reserved_tokens: 200 +/// ``` +/// +/// Rules are evaluated in order; the first whose `match` is satisfied +/// (or which has no `match` at all) applies. A request satisfying no +/// rule's `match` is **not** rate limited by this filter instance -- +/// add a trailing rule with no `match` for a catch-all budget instead. +pub struct TokenRateLimitFilter { + /// Compiled rules, evaluated in configured order. + rules: Vec, + + /// Monotonic clock reference; all timestamps are offsets from this. + epoch: Instant, +} + +impl TokenRateLimitFilter { + /// Create a filter from parsed YAML config. + /// + /// # Errors + /// + /// Returns [`FilterError`] if the YAML config is invalid, `rules` is + /// empty, two rules share a `name`, or any individual rule fails to + /// compile (see `compile_rule`). + pub fn from_config(config: &serde_yaml::Value) -> Result, FilterError> { + let cfg: TokenRateLimitConfig = parse_filter_config("token_rate_limit", config)?; + if cfg.rules.is_empty() { + return Err("token_rate_limit: at least one rule is required".into()); + } + let mut seen_names = HashSet::with_capacity(cfg.rules.len()); + for rule in &cfg.rules { + if !seen_names.insert(rule.name.clone()) { + return Err(format!("token_rate_limit: duplicate rule name '{}'", rule.name).into()); + } + } + + let backend = build_backend_resource(&cfg.backend)?; + let rules = cfg + .rules + .into_iter() + .map(|rule| compile_rule(rule, &backend)) + .collect::, _>>()?; + + Ok(Box::new(Self { + rules, + epoch: Instant::now(), + })) + } + + /// Milliseconds elapsed since this filter's epoch. + #[expect( + clippy::cast_possible_truncation, + reason = "millis fit u64 for any realistic process uptime" + )] + fn now_ms(&self) -> u64 { + self.epoch.elapsed().as_millis().min(u128::from(u64::MAX)) as u64 + } + + /// The first rule (in configured order) whose `match` is satisfied + /// by `headers`, alongside its index for reconciliation. + fn matching_rule(&self, headers: &http::HeaderMap) -> Option<(usize, &CompiledRule)> { + self.rules.iter().enumerate().find(|(_, rule)| rule.matches(headers)) + } + + /// Reclaim idle/orphaned in-process state for one rule and publish + /// its gauges. + /// + /// No-ops for a Valkey backend (`cleanup()` returns `None`): expiry + /// there is handled by the Lua reserve script itself. + fn cleanup_and_record_state(rule: &CompiledRule, now_ms: u64) { + let Some(report) = rule.backend.cleanup(now_ms, 1) else { + return; + }; + record_cleanup_metrics(&rule.name, report); + } + + /// Record metrics/metadata for an admitted reservation. + fn record_admission( + ctx: &mut HttpFilterContext<'_>, + rule_index: usize, + rule: &CompiledRule, + admitted: AdmittedReservation, + ) { + counter!("praxis_ai_token_rate_limit_requests_total", "decision" => "admitted", "rule" => rule.name.clone()) + .increment(1); + counter!("praxis_ai_token_rate_limit_tokens_total", "kind" => "estimated", "rule" => rule.name.clone()) + .increment(admitted.estimate); + ctx.set_metadata(META_RESERVATION_ID, admitted.reservation_id.to_string()); + ctx.set_metadata(META_BUCKET_KEY, admitted.key); + ctx.set_metadata(META_RULE_INDEX, rule_index.to_string()); + } + + /// Build the 429 rejection for a denied reservation, including the + /// token-denominated rate limit headers. + fn denied_action(rule: &CompiledRule, retry_after_ms: u64) -> FilterAction { + counter!("praxis_ai_token_rate_limit_requests_total", "decision" => "denied", "rule" => rule.name.clone()) + .increment(1); + let retry_secs = retry_after_ms.saturating_add(999) / 1000; + let retry_secs = retry_secs.max(1); + FilterAction::Reject( + Rejection::status(429) + .with_header("Retry-After", retry_secs.to_string()) + .with_header(HEADER_RATELIMIT_LIMIT_TOKENS, rule.backend.limit().to_string()) + .with_header(HEADER_RATELIMIT_REMAINING_TOKENS, "0") + .with_header(HEADER_RATELIMIT_RESET, retry_secs.to_string()), + ) + } + + /// Turn a completed `reserve()` call into the `on_request` result: + /// record admission metadata/metrics, build the 429 rejection, or + /// fail closed (503) on a backend error. + fn handle_reserve_outcome( + ctx: &mut HttpFilterContext<'_>, + rule_index: usize, + rule: &CompiledRule, + key: String, + outcome: Result, + ) -> FilterAction { + match outcome { + Ok(BackendReserve::Admitted { + reservation_id, + estimate, + }) => { + let admitted = AdmittedReservation { + key, + reservation_id, + estimate, + }; + Self::record_admission(ctx, rule_index, rule, admitted); + FilterAction::Continue + }, + Ok(BackendReserve::Denied { retry_after_ms }) => { + tracing::info!( + estimate = rule.reserved_tokens, + key, + rule = rule.name, + "token_rate_limit: rejecting request (429)" + ); + Self::denied_action(rule, retry_after_ms) + }, + Err(error) => { + counter!("praxis_ai_token_rate_limit_backend_errors_total", "operation" => "reserve", "rule" => rule.name.clone()) + .increment(1); + tracing::error!(%error, rule = rule.name, "token_rate_limit: admission backend failed, failing closed"); + FilterAction::Reject(Rejection::status(503)) + }, + } + } + + /// Look up the reservation/key/rule metadata `on_request` stashed for + /// this exchange, if all three are present and the rule index still + /// resolves -- the shared precondition for [`Self::reconcile`]. + fn reconciliation_context(&self, ctx: &HttpFilterContext<'_>) -> Option<(ReconcileRequest, &CompiledRule)> { + let reservation_id = ctx + .get_metadata(META_RESERVATION_ID) + .and_then(|v| v.parse::().ok())?; + let key = ctx.get_metadata(META_BUCKET_KEY).map(str::to_owned)?; + let rule = ctx + .get_metadata(META_RULE_INDEX) + .and_then(|v| v.parse::().ok()) + .and_then(|index| self.rules.get(index))?; + let actual = ctx.get_metadata(META_TOKEN_TOTAL).and_then(|v| v.parse::().ok()); + if actual.is_none() { + tracing::trace!("token_rate_limit: no token.total metadata at end of stream, charging at estimate"); + } + let request = ReconcileRequest { + key, + reservation_id, + actual, + estimate: rule.reserved_tokens, + now_ms: self.now_ms(), + }; + Some((request, rule)) + } + + /// Reconcile a prior reservation against actual usage, if known. + /// + /// No-ops if the reservation, bucket key, or originating rule index + /// metadata is absent -- reservation stands as final rather than + /// guessing. + /// + /// In-process state reconciles synchronously and immediately (cheap, + /// no I/O). A Valkey backend instead enqueues the reconciliation onto + /// a background worker (see `backend::ValkeyTokenRateLimitBackend`/ + /// `backend::ValkeyTokenBucketBackend`) so the response is never held + /// up on a network round-trip that has no bearing on whether *this* + /// request was admitted. + fn reconcile(&self, ctx: &HttpFilterContext<'_>) { + let Some((request, rule)) = self.reconciliation_context(ctx) else { + return; + }; + + if let Some(settlement) = rule.backend.reconcile_sync(&request) { + record_settlement_metrics(&rule.name, &settlement); + tracing::debug!( + ?settlement, + rule = rule.name, + "token_rate_limit: reconciled reservation against actual usage" + ); + return; + } + if let Err(error) = rule.backend.enqueue_reconcile(request) { + tracing::error!(%error, rule = rule.name, "token_rate_limit: failed to enqueue reconciliation"); + } + } +} + +/// The bucket key, reservation ID, and estimate an admitted +/// [`BackendReserve::Admitted`] carries, bundled so +/// [`TokenRateLimitFilter::record_admission`] stays within clippy's +/// argument-count budget. +struct AdmittedReservation { + /// The budget key this reservation was admitted under (see + /// [`FALLBACK_KEY`] -- always that sentinel in this milestone). + key: String, + /// The backend-issued reservation ID, stashed for later reconciliation. + reservation_id: u64, + /// Tokens reserved at admission (the rule's `reserved_tokens`, echoed + /// back by the backend). + estimate: u64, +} + +/// Emit gauges/counters for one rule's cleanup pass. +fn record_cleanup_metrics(rule_name: &str, report: CleanupReport) { + if report.orphaned > 0 { + counter!("praxis_ai_token_rate_limit_reservations_total", "result" => "orphaned", "rule" => rule_name.to_owned()) + .increment(report.orphaned as u64); + } + #[expect( + clippy::cast_precision_loss, + reason = "metrics gauges use f64, bounded by config caps in practice" + )] + { + gauge!("praxis_ai_token_rate_limit_active_reservations", "rule" => rule_name.to_owned()) + .set(report.active_reservations as f64); + gauge!("praxis_ai_token_rate_limit_active_keys", "rule" => rule_name.to_owned()).set(report.active_keys as f64); + } +} + +/// Emit counters for a completed synchronous (in-process) reconciliation. +fn record_settlement_metrics(rule_name: &str, settlement: &BackendSettlement) { + if let BackendSettlement::Applied { + actual, + refund, + overage, + } = *settlement + { + counter!("praxis_ai_token_rate_limit_reservations_total", "result" => "reconciled", "rule" => rule_name.to_owned()) + .increment(1); + counter!("praxis_ai_token_rate_limit_tokens_total", "kind" => "actual", "rule" => rule_name.to_owned()) + .increment(actual); + counter!("praxis_ai_token_rate_limit_tokens_total", "kind" => "refunded", "rule" => rule_name.to_owned()) + .increment(refund); + counter!("praxis_ai_token_rate_limit_tokens_total", "kind" => "overage", "rule" => rule_name.to_owned()) + .increment(overage); + } +} + +#[async_trait] +impl HttpFilter for TokenRateLimitFilter { + fn name(&self) -> &'static str { + "token_rate_limit" + } + + fn response_body_access(&self) -> BodyAccess { + BodyAccess::ReadOnly + } + + fn response_body_mode(&self) -> BodyMode { + BodyMode::Stream + } + + async fn on_request(&self, ctx: &mut HttpFilterContext<'_>) -> Result { + let now_ms = self.now_ms(); + let Some((rule_index, rule)) = self.matching_rule(&ctx.request.headers) else { + // No configured rule applies to this request -- not rate + // limited by this filter instance. Operators wanting a + // catch-all budget add a trailing rule with no `match`. + return Ok(FilterAction::Continue); + }; + Self::cleanup_and_record_state(rule, now_ms); + let key = FALLBACK_KEY.to_owned(); + let outcome = rule + .backend + .reserve(ReserveRequest { + key: key.clone(), + estimate: rule.reserved_tokens, + now_ms, + }) + .await; + Ok(Self::handle_reserve_outcome(ctx, rule_index, rule, key, outcome)) + } + + fn on_response_body( + &self, + ctx: &mut HttpFilterContext<'_>, + _body: &mut Option, + end_of_stream: bool, + ) -> Result { + if end_of_stream { + self.reconcile(ctx); + ctx.filter_metadata.remove(META_RESERVATION_ID); + ctx.filter_metadata.remove(META_BUCKET_KEY); + ctx.filter_metadata.remove(META_RULE_INDEX); + } + Ok(FilterAction::Continue) + } +} + +/// Expand one `${ENV_VAR}` reference in a backend URL, if present. +/// +/// Takes an explicit lookup function (rather than calling +/// [`std::env::var`] directly) so tests can exercise both branches +/// without mutating real process-global environment state. +fn expand_backend_url_with( + url: &str, + lookup: impl Fn(&str) -> Result, +) -> Result { + let Some(start) = url.find("${") else { + return Ok(url.to_owned()); + }; + let Some(name) = url.strip_prefix("${").and_then(|value| value.strip_suffix('}')) else { + return Err("token_rate_limit: backend.url supports one complete ${ENV_VAR} reference".into()); + }; + if start != 0 || name.contains("${") { + return Err("token_rate_limit: backend.url supports one complete ${ENV_VAR} reference".into()); + } + if name.is_empty() + || !name + .bytes() + .enumerate() + .all(|(index, byte)| byte == b'_' || byte.is_ascii_uppercase() || (index > 0 && byte.is_ascii_digit())) + { + return Err("token_rate_limit: backend.url contains an invalid environment variable reference".into()); + } + lookup(name).map_err(|_error| "token_rate_limit: backend.url environment variable is not set".into()) +} + +/// Expand one `${ENV_VAR}` reference in a backend URL against the real +/// process environment. +fn expand_backend_url(url: &str) -> Result { + expand_backend_url_with(url, |name| std::env::var(name)) +} + +/// Parse a simple `` duration (`ms`, `s`, `m`, `h`) into +/// milliseconds, as used by `window` and `reservation_timeout`. +fn parse_duration_ms(value: &str) -> Result { + let value = value.trim(); + let (number, multiplier) = if let Some(value) = value.strip_suffix("ms") { + (value, 1_u64) + } else if let Some(value) = value.strip_suffix('s') { + (value, 1_000_u64) + } else if let Some(value) = value.strip_suffix('m') { + (value, 60_000_u64) + } else if let Some(value) = value.strip_suffix('h') { + (value, 3_600_000_u64) + } else { + return Err(format!("token_rate_limit: invalid duration '{value}'").into()); + }; + let amount = number + .parse::() + .map_err(|error| format!("token_rate_limit: invalid duration '{value}': {error}"))?; + let millis = amount + .checked_mul(multiplier) + .ok_or("token_rate_limit: duration overflow")?; + if millis == 0 { + return Err("token_rate_limit: duration must be positive".into()); + } + Ok(millis) +} + +impl std::fmt::Debug for TokenRateLimitFilter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TokenRateLimitFilter") + .field("rules", &self.rules.iter().map(|rule| &rule.name).collect::>()) + .finish_non_exhaustive() + } +} + +/// White-box tests that construct a [`CompiledRule`] directly with a +/// purpose-built [`backend::TokenRateLimitStateBackend`], for behavior +/// unreachable through [`TokenRateLimitFilter::from_config`] with any +/// real backend. Kept separate from `tests.rs`, which deliberately only +/// drives the filter through its public `HttpFilter` surface. +#[cfg(test)] +#[expect(clippy::allow_attributes, reason = "blanket test suppressions")] +#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic, reason = "tests")] +mod backend_injection_tests { + use praxis_filter::HttpFilter as _; + + use super::{ + CompiledRule, TokenRateLimitFilter, + backend::{ + BackendError, BackendReserve, BackendSettlement, ReconcileRequest, ReserveRequest, + TokenRateLimitStateBackend, + }, + }; + + /// A backend that admits every reservation but always fails to + /// enqueue its reconciliation -- the one way + /// [`TokenRateLimitFilter::reconcile`]'s enqueue-failure log line is + /// reachable in production (a real Valkey-backed rule's + /// background-worker channel saturated or its receiver gone), and + /// impractical to drive there for real without either exhausting a + /// live worker's 1024-deep channel or tearing down its receiver + /// mid-test. + struct EnqueueAlwaysFailsBackend; + + #[async_trait::async_trait] + impl TokenRateLimitStateBackend for EnqueueAlwaysFailsBackend { + async fn reserve(&self, _request: ReserveRequest) -> Result { + Ok(BackendReserve::Admitted { + reservation_id: 1, + estimate: 1, + }) + } + + async fn reconcile(&self, _request: ReconcileRequest) -> Result { + panic!("not exercised by this test") + } + + fn enqueue_reconcile(&self, _request: ReconcileRequest) -> Result<(), BackendError> { + Err(BackendError::Unavailable("enqueue failed (test)".into())) + } + + fn limit(&self) -> u64 { + 1 + } + } + + /// [`TokenRateLimitFilter::reconcile`] falls back to + /// `enqueue_reconcile` when `reconcile_sync` returns `None` (the + /// default, unimplemented by [`EnqueueAlwaysFailsBackend`]). When + /// that enqueue itself fails, the error must be logged and + /// swallowed, not propagated -- a reconciliation failure is never + /// the inbound request's fault, so it must not affect the response + /// already on its way out. + #[tokio::test] + async fn reconcile_logs_rather_than_propagates_an_enqueue_reconcile_failure() { + let filter = TokenRateLimitFilter { + rules: vec![CompiledRule { + name: "default".to_owned(), + matcher: None, + backend: std::sync::Arc::new(EnqueueAlwaysFailsBackend), + reserved_tokens: 1, + }], + epoch: std::time::Instant::now(), + }; + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut ctx = crate::test_utils::make_filter_context(&req); + drop(filter.on_request(&mut ctx).await.unwrap()); + + let mut body = None; + drop(filter.on_response_body(&mut ctx, &mut body, true).unwrap()); + } +} diff --git a/filters/src/token_rate_limit/tests.rs b/filters/src/token_rate_limit/tests.rs new file mode 100644 index 0000000000..5d35d1225a --- /dev/null +++ b/filters/src/token_rate_limit/tests.rs @@ -0,0 +1,1157 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Praxis Contributors + +//! Tests for the `token_rate_limit` filter. + +use praxis_filter::{FilterAction, HttpFilter}; + +use super::TokenRateLimitFilter; +use crate::token_usage::META_TOKEN_TOTAL; + +/// Wrap one rule body (already-valid YAML lines, unindented) into a +/// full one-rule `rules:` config, named `"default"`. Most scenarios +/// pre-date per-rule algorithm choice and only care about one rule's +/// behavior in isolation -- multi-rule dispatch itself is covered +/// separately below. +fn single_rule(body: &str) -> String { + let indented = body + .lines() + .map(|line| format!(" {line}")) + .collect::>() + .join("\n"); + format!("rules:\n - name: default\n{indented}\n") +} + +/// [`single_rule`], parsed straight into a [`serde_yaml::Value`]. +fn single_rule_yaml(body: &str) -> serde_yaml::Value { + serde_yaml::from_str(&single_rule(body)).unwrap() +} + +/// [`single_rule`], with a filter-level `top_level` block (e.g. +/// `backend: {...}`) prepended as a sibling of `rules:`. +fn single_rule_yaml_with(top_level: &str, body: &str) -> serde_yaml::Value { + serde_yaml::from_str(&format!("{top_level}\n{}", single_rule(body))).unwrap() +} + +/// Build a request carrying a single extra header, for `match` tests. +fn make_request_with_header(name: &str, value: &str) -> praxis_filter::Request { + let mut req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + req.headers.insert( + http::header::HeaderName::from_bytes(name.as_bytes()).unwrap(), + http::HeaderValue::from_str(value).unwrap(), + ); + req +} + +// ----------------------------------------------------------------------------- +// Config Validation +// ----------------------------------------------------------------------------- + +#[test] +fn from_config_parses_valid_config() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 100000\nreserved_tokens: 500"); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + assert_eq!(filter.name(), "token_rate_limit"); +} + +#[test] +fn from_config_rejects_an_empty_rules_list() { + let yaml: serde_yaml::Value = serde_yaml::from_str("rules: []\n").unwrap(); + let err = TokenRateLimitFilter::from_config(&yaml).err().expect("should error"); + assert!(err.to_string().contains("at least one rule"), "got: {err}"); +} + +#[test] +fn from_config_rejects_duplicate_rule_names() { + let yaml: serde_yaml::Value = serde_yaml::from_str( + "rules:\n\ + \x20 - name: dup\n\ + \x20 algorithm: sliding_window\n\ + \x20 window: 1h\n\ + \x20 capacity: 100\n\ + \x20 reserved_tokens: 10\n\ + \x20 - name: dup\n\ + \x20 algorithm: token_bucket\n\ + \x20 capacity: 100\n\ + \x20 refill_rate: 1\n\ + \x20 reserved_tokens: 10\n", + ) + .unwrap(); + let err = TokenRateLimitFilter::from_config(&yaml).err().expect("should error"); + assert!(err.to_string().contains("duplicate rule name"), "got: {err}"); +} + +#[test] +fn from_config_rejects_zero_capacity() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 0\nreserved_tokens: 10"); + let err = TokenRateLimitFilter::from_config(&yaml).err().expect("should error"); + assert!(err.to_string().contains("capacity must be"), "got: {err}"); +} + +/// `validate_rule_bounds` is the single gate both algorithms pass +/// through in `compile_rule`, before either builds its own backend. +/// `token_bucket_ledger` happens to re-check this same bound on its +/// own construction path, but `ledger` (`sliding_window`) does not -- +/// so this rejection has to come from the shared gate, not either +/// algorithm's downstream validation, to protect both. +#[test] +fn from_config_rejects_capacity_above_the_lua_safe_integer_bound() { + let over_bound = super::token_bucket_ledger::MAX_F64_SAFE_INTEGER + 1; + + let sliding_window = single_rule_yaml(&format!( + "algorithm: sliding_window\nwindow: 1h\ncapacity: {over_bound}\nreserved_tokens: 10" + )); + let err = TokenRateLimitFilter::from_config(&sliding_window) + .err() + .expect("sliding_window should reject a capacity beyond the f64 safe-integer bound"); + assert!(err.to_string().contains("must not exceed"), "got: {err}"); + + let token_bucket = single_rule_yaml(&format!( + "algorithm: token_bucket\ncapacity: {over_bound}\nrefill_rate: 1\nreserved_tokens: 10" + )); + let err = TokenRateLimitFilter::from_config(&token_bucket) + .err() + .expect("token_bucket should reject a capacity beyond the f64 safe-integer bound"); + assert!(err.to_string().contains("must not exceed"), "got: {err}"); +} + +#[test] +fn from_config_rejects_zero_estimate() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 0"); + let err = TokenRateLimitFilter::from_config(&yaml).err().expect("should error"); + assert!(err.to_string().contains("reserved_tokens"), "got: {err}"); +} + +#[test] +fn from_config_rejects_estimate_exceeding_capacity() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 500"); + let err = TokenRateLimitFilter::from_config(&yaml).err().expect("should error"); + assert!(err.to_string().contains("must not exceed capacity"), "got: {err}"); +} + +#[test] +fn from_config_rejects_invalid_window() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: not-a-duration\ncapacity: 100\nreserved_tokens: 5"); + let err = TokenRateLimitFilter::from_config(&yaml).err().expect("should error"); + assert!(err.to_string().contains("invalid duration"), "got: {err}"); +} + +#[test] +fn from_config_rejects_unknown_field() { + let yaml = single_rule_yaml( + "algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 5\nbucket_key: header", + ); + assert!( + TokenRateLimitFilter::from_config(&yaml).is_err(), + "composite/CEL bucket keys are still deliberately unsupported, config should reject the unknown field" + ); +} + +#[test] +fn from_config_rejects_the_old_flat_pre_rules_shape() { + let yaml: serde_yaml::Value = serde_yaml::from_str("window: 1h\ncapacity: 100000\nreserved_tokens: 500").unwrap(); + let err = TokenRateLimitFilter::from_config(&yaml).err().expect("should error"); + assert!(err.to_string().contains("rules"), "got: {err}"); +} + +// ----------------------------------------------------------------------------- +// Backend config (Valkey opt-in for shared/distributed state) +// ----------------------------------------------------------------------------- + +#[test] +fn from_config_defaults_to_memory_backend_when_backend_block_absent() { + // No `backend:` block at all must keep working exactly like before this + // field was added -- a config written before the Valkey backend existed + // should never start silently expecting shared state it never asked for. + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 5"); + assert!(TokenRateLimitFilter::from_config(&yaml).is_ok()); +} + +#[test] +fn from_config_accepts_explicit_memory_backend() { + let yaml = single_rule_yaml_with( + "backend:\n kind: memory", + "algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 5", + ); + assert!(TokenRateLimitFilter::from_config(&yaml).is_ok()); +} + +#[test] +fn from_config_rejects_valkey_backend_without_url() { + // A distributed deployment that forgets `backend.url` must fail loudly + // at startup, not silently fall back to per-instance state -- silent + // fallback would defeat the whole point of asking for a shared backend. + let yaml = single_rule_yaml_with( + "backend:\n kind: valkey", + "algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 5", + ); + let err = TokenRateLimitFilter::from_config(&yaml).err().expect("should error"); + assert!(err.to_string().contains("backend.url is required"), "got: {err}"); +} + +#[test] +fn from_config_accepts_valkey_backend_with_url() { + let yaml = single_rule_yaml_with( + "backend:\n kind: valkey\n url: redis://127.0.0.1:6399", + "algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 5", + ); + assert!(TokenRateLimitFilter::from_config(&yaml).is_ok()); +} + +#[test] +fn from_config_accepts_two_rules_with_different_algorithms_sharing_one_filter_level_valkey_backend() { + // The whole point of moving `backend:` from per-rule to per-filter: + // one `backend:` block, two rules with two different algorithms, one + // shared Valkey connection underneath -- not one connection per rule. + let yaml: serde_yaml::Value = serde_yaml::from_str( + "backend:\n\ + \x20 kind: valkey\n\ + \x20 url: redis://127.0.0.1:6399\n\ + \x20 namespace: shared\n\ + rules:\n\ + \x20 - name: team-alpha\n\ + \x20 algorithm: sliding_window\n\ + \x20 window: 1h\n\ + \x20 capacity: 100\n\ + \x20 reserved_tokens: 5\n\ + \x20 - name: team-beta\n\ + \x20 algorithm: token_bucket\n\ + \x20 capacity: 100\n\ + \x20 refill_rate: 1\n\ + \x20 reserved_tokens: 5\n", + ) + .unwrap(); + assert!(TokenRateLimitFilter::from_config(&yaml).is_ok()); +} + +// `${ENV_VAR}` expansion is tested directly against `expand_backend_url_with` +// below (dependency-injected lookup) rather than through real process +// environment mutation, which is unsafe in this edition and would be +// racy across parallel test threads regardless. + +#[test] +fn expand_backend_url_with_substitutes_a_resolved_env_var() { + let expanded = super::expand_backend_url_with("${TOKEN_RATE_LIMIT_TEST_URL}", |name| { + assert_eq!(name, "TOKEN_RATE_LIMIT_TEST_URL"); + Ok("redis://127.0.0.1:6399".to_owned()) + }) + .unwrap(); + assert_eq!(expanded, "redis://127.0.0.1:6399"); +} + +#[test] +fn expand_backend_url_with_passes_through_a_url_without_any_reference() { + let expanded = super::expand_backend_url_with("redis://127.0.0.1:6399", |_name| { + panic!("lookup should not be called when there is no ${{ENV_VAR}} reference") + }) + .unwrap(); + assert_eq!(expanded, "redis://127.0.0.1:6399"); +} + +#[test] +fn expand_backend_url_with_rejects_an_unset_env_var() { + let err = super::expand_backend_url_with("${UNSET_VAR}", |_name| Err(std::env::VarError::NotPresent)) + .expect_err("should error"); + assert!( + err.to_string().contains("environment variable is not set"), + "got: {err}" + ); +} + +#[test] +fn expand_backend_url_with_rejects_an_embedded_reference_not_spanning_the_whole_url() { + // A misconfigured distributed deployment silently connecting to the + // wrong host (e.g. a typo'd literal instead of the intended env var) + // would be a config-integrity failure that's easy to miss in review + // -- only whole-value substitution is supported, so any other shape + // fails config load loudly rather than doing partial/ambiguous + // substitution. + let err = super::expand_backend_url_with("redis://${REDIS_HOST}:6379", |_name| { + panic!("lookup must not run for an unsupported reference shape") + }) + .expect_err("should error"); + assert!(err.to_string().contains("one complete"), "got: {err}"); +} + +#[test] +fn expand_backend_url_with_rejects_multiple_references() { + let err = super::expand_backend_url_with("${A}${B}", |_name| { + panic!("lookup must not run for an unsupported reference shape") + }) + .expect_err("should error"); + assert!(err.to_string().contains("one complete"), "got: {err}"); +} + +#[test] +fn expand_backend_url_with_rejects_an_invalid_variable_name() { + // A misconfigured `${...}` shape (rather than a clean uppercase + // env-var name) must fail loudly with a clear reason at startup, + // not be silently passed through as a literal, non-functioning URL. + let err = super::expand_backend_url_with("${lower_case}", |_name| { + panic!("lookup must not run for an invalid variable name") + }) + .expect_err("should error"); + assert!( + err.to_string().contains("invalid environment variable reference"), + "got: {err}" + ); +} + +// ----------------------------------------------------------------------------- +// Admission (on_request) +// ----------------------------------------------------------------------------- + +#[tokio::test] +async fn admits_request_within_budget() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 1000\nreserved_tokens: 200"); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + let action = filter.on_request(&mut ctx).await.unwrap(); + assert!(matches!(action, FilterAction::Continue), "should admit within budget"); + assert!( + ctx.get_metadata("token_rate_limit.reservation_id").is_some(), + "reservation id should be stashed for reconciliation" + ); +} + +#[tokio::test] +async fn rejects_with_429_when_budget_exhausted() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 60"); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut first_ctx = crate::test_utils::make_filter_context(&req); + let mut second_ctx = crate::test_utils::make_filter_context(&req); + + // First request consumes 60 of 100; second needs another 60, only 40 left. + let first = filter.on_request(&mut first_ctx).await.unwrap(); + assert!( + matches!(first, FilterAction::Continue), + "first request should be admitted" + ); + + let second = filter.on_request(&mut second_ctx).await.unwrap(); + match second { + FilterAction::Reject(rejection) => { + assert_eq!(rejection.status, 429); + let has_header = |name: &str| rejection.headers.iter().any(|(n, _)| n == name); + assert!(has_header("Retry-After"), "429 should carry Retry-After"); + assert!( + has_header("X-RateLimit-Limit-Tokens"), + "429 should carry token-suffixed limit header" + ); + assert!( + has_header("X-RateLimit-Remaining-Tokens"), + "429 should carry token-suffixed remaining header" + ); + assert!(has_header("X-RateLimit-Reset-Tokens"), "429 should carry reset header"); + }, + other => panic!("second request should be rejected, insufficient tokens remain, got {other:?}"), + } +} + +#[tokio::test] +async fn rejection_does_not_consume_budget() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 10\nreserved_tokens: 10"); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + + // First request consumes the entire capacity; every request after + // that should be rejected, and rejecting must not partially drain + // (or otherwise corrupt) the already-exhausted budget. + let mut first_ctx = crate::test_utils::make_filter_context(&req); + let first = filter.on_request(&mut first_ctx).await.unwrap(); + assert!( + matches!(first, FilterAction::Continue), + "first request should exactly exhaust the capacity" + ); + + for _ in 0..3 { + let mut ctx = crate::test_utils::make_filter_context(&req); + let action = filter.on_request(&mut ctx).await.unwrap(); + assert!( + matches!(action, FilterAction::Reject(_)), + "exhausted budget should always reject" + ); + } +} + +#[tokio::test] +async fn a_request_matching_no_rule_is_not_rate_limited() { + // Business behavior: a rule scoped to one app must not silently + // become a global rate limiter for traffic it was never configured + // to cover -- operators who want a catch-all budget add a trailing + // rule with no `match`. + let yaml: serde_yaml::Value = serde_yaml::from_str( + "rules:\n\ + \x20 - name: alpha-only\n\ + \x20 match:\n\ + \x20 headers:\n\ + \x20 x-app-id: alpha\n\ + \x20 algorithm: sliding_window\n\ + \x20 window: 1h\n\ + \x20 capacity: 1\n\ + \x20 reserved_tokens: 1\n", + ) + .unwrap(); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let unmatched_req = make_request_with_header("x-app-id", "beta"); + for _ in 0..5 { + let mut ctx = crate::test_utils::make_filter_context(&unmatched_req); + assert!( + matches!(filter.on_request(&mut ctx).await.unwrap(), FilterAction::Continue), + "traffic matching no configured rule must pass through, even repeatedly" + ); + } +} + +// ----------------------------------------------------------------------------- +// Reconciliation (on_response_body) +// ----------------------------------------------------------------------------- + +#[tokio::test] +async fn reconcile_releases_unused_tokens_on_overestimate() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 50"); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + drop(filter.on_request(&mut ctx).await.unwrap()); // reserves 50, 50 left + ctx.set_metadata(META_TOKEN_TOTAL, "30"); // actual usage was only 30 + + let mut body = None; + drop(filter.on_response_body(&mut ctx, &mut body, true).unwrap()); + + // Reserved 50, actual 30 -> release 20 back -> 70 should now remain. + // A next 50-token request should succeed (70 >= 50, leaving 20)... + let mut second_ctx = crate::test_utils::make_filter_context(&req); + let second = filter.on_request(&mut second_ctx).await.unwrap(); + assert!( + matches!(second, FilterAction::Continue), + "70 remaining should admit a 50-token request" + ); + + // ...but a third 50-token request should now fail (only 20 left). + let mut third_ctx = crate::test_utils::make_filter_context(&req); + let third = filter.on_request(&mut third_ctx).await.unwrap(); + assert!( + matches!(third, FilterAction::Reject(_)), + "only 20 remaining should reject a 50-token request" + ); +} + +#[tokio::test] +async fn reconcile_draws_more_tokens_on_underestimate_and_can_starve_next_request() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 50"); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + drop(filter.on_request(&mut ctx).await.unwrap()); // reserves 50, 50 left + ctx.set_metadata(META_TOKEN_TOTAL, "90"); // actual usage exceeded the estimate + + let mut body = None; + drop(filter.on_response_body(&mut ctx, &mut body, true).unwrap()); + + // Reserved 50, actual 90 -> the window now holds 90 of 100 -> only 10 left. + let mut next_ctx = crate::test_utils::make_filter_context(&req); + let next_action = filter.on_request(&mut next_ctx).await.unwrap(); + assert!( + matches!(next_action, FilterAction::Reject(_)), + "underestimate should have drawn the window down enough to starve the next 50-token request" + ); +} + +#[tokio::test] +async fn reconcile_charges_the_estimate_without_token_total_metadata() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 50"); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + drop(filter.on_request(&mut ctx).await.unwrap()); // reserves 50, 50 left + // No token.total metadata set (e.g. token_count filter not configured upstream). + + let mut body = None; + drop(filter.on_response_body(&mut ctx, &mut body, true).unwrap()); + + // Settled at the estimate (50 of 100 used): a second 50-token request + // should still fit exactly... + let mut next_ctx = crate::test_utils::make_filter_context(&req); + let next_action = filter.on_request(&mut next_ctx).await.unwrap(); + assert!( + matches!(next_action, FilterAction::Continue), + "50 of 100 already settled leaves exactly 50 for the next request" + ); + + // ...but a third would exceed the window's capacity. + let mut third_ctx = crate::test_utils::make_filter_context(&req); + let third_action = filter.on_request(&mut third_ctx).await.unwrap(); + assert!( + matches!(third_action, FilterAction::Reject(_)), + "window is now fully settled at 100/100" + ); +} + +#[tokio::test] +async fn does_not_reconcile_before_end_of_stream() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 50"); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + drop(filter.on_request(&mut ctx).await.unwrap()); + ctx.set_metadata(META_TOKEN_TOTAL, "5"); + + let mut body = None; + drop(filter.on_response_body(&mut ctx, &mut body, false).unwrap()); + + // Reconciliation must not have run yet: 50 tokens are still an active + // reservation (not settled down to actual=5), so a fresh 50-token + // request only has the remaining 50 of capacity=100 to draw from, and + // a second one on top of that should fail. + let mut second_ctx = crate::test_utils::make_filter_context(&req); + let second = filter.on_request(&mut second_ctx).await.unwrap(); + assert!( + matches!(second, FilterAction::Continue), + "50 remaining should admit one more 50-token request" + ); + + let mut third_ctx = crate::test_utils::make_filter_context(&req); + let third = filter.on_request(&mut third_ctx).await.unwrap(); + assert!( + matches!(third, FilterAction::Reject(_)), + "window should be fully committed now (no premature release happened pre-end_of_stream)" + ); +} + +/// End-of-stream on an exchange that was never admitted by this filter +/// (e.g. it matched no rule, or a prior filter already short-circuited +/// the request) must not panic or reconcile phantom state -- `reconcile` +/// no-ops when `on_request` never stashed reservation/key/rule metadata. +#[tokio::test] +async fn reconcile_is_a_noop_without_prior_admission_metadata() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 50"); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut ctx = crate::test_utils::make_filter_context(&req); + // Deliberately skip `on_request` -- no reservation/key/rule metadata + // is present on `ctx`. + let mut body = None; + let action = filter.on_response_body(&mut ctx, &mut body, true).unwrap(); + assert!(matches!(action, FilterAction::Continue)); + + // The full 100-token budget must still be there for a real request -- + // the no-op above must not have reserved or settled anything against it. + let mut fresh_ctx = crate::test_utils::make_filter_context(&req); + assert!(matches!( + filter.on_request(&mut fresh_ctx).await.unwrap(), + FilterAction::Continue + )); +} + +// ----------------------------------------------------------------------------- +// Lost-request handling (the proposal's still-open question, answered here +// via reservation_timeout) +// ----------------------------------------------------------------------------- + +#[tokio::test] +async fn lost_request_is_charged_at_its_estimate_and_cannot_bypass_the_budget() { + // A client that aborts a request before the response completes + // (connection reset, client timeout, upstream crash) must not be + // able to dodge the budget entirely by ensuring on_response_body/ + // reconciliation never runs -- that would make token rate limiting + // trivially bypassable by just not waiting for the response. + // reservation_timeout bounds how long such a reservation is trusted + // before being conservatively + // charged at its estimate, matching the "lost request handling" + // question the proposal's own design doc leaves open. + let yaml = single_rule_yaml( + "algorithm: sliding_window\nwindow: 300ms\ncapacity: 50\nreserved_tokens: 50\nreservation_timeout: 50ms", + ); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut ctx = crate::test_utils::make_filter_context(&req); + assert!( + matches!(filter.on_request(&mut ctx).await.unwrap(), FilterAction::Continue), + "first request should be admitted, reserving the entire 50-token capacity" + ); + // Simulate an aborted request: on_response_body is deliberately never + // called, so this reservation is never explicitly reconciled. + drop(ctx); + + tokio::time::sleep(std::time::Duration::from_millis(80)).await; // past reservation_timeout, still inside window + + let mut second_ctx = crate::test_utils::make_filter_context(&req); + let second = filter.on_request(&mut second_ctx).await.unwrap(); + assert!( + matches!(second, FilterAction::Reject(_)), + "the aborted request's reservation must still be charged against the window once it times out -- it \ + must not grant free/unmetered capacity just because the response was never observed" + ); + + tokio::time::sleep(std::time::Duration::from_millis(250)).await; // past the window's own expiry too + + let mut third_ctx = crate::test_utils::make_filter_context(&req); + let third = filter.on_request(&mut third_ctx).await.unwrap(); + assert!( + matches!(third, FilterAction::Continue), + "once the window rolls over, a one-time lost request must not permanently lock the key out" + ); +} + +#[test] +fn from_config_rejects_an_invalid_match_header_name() { + let yaml = single_rule_yaml( + "algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 5\nmatch:\n headers:\n \"x \ + app\": bad\n", + ); + let err = TokenRateLimitFilter::from_config(&yaml).err().expect("should error"); + assert!(err.to_string().contains("invalid match header"), "got: {err}"); +} + +#[test] +fn from_config_accepts_minute_suffix_durations() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 5m\ncapacity: 100\nreserved_tokens: 5"); + assert!(TokenRateLimitFilter::from_config(&yaml).is_ok()); +} + +#[test] +fn from_config_rejects_a_zero_duration_window() { + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 0s\ncapacity: 100\nreserved_tokens: 5"); + let err = TokenRateLimitFilter::from_config(&yaml).err().expect("should error"); + assert!(err.to_string().contains("must be positive"), "got: {err}"); +} + +#[test] +fn debug_format_lists_configured_rule_names() { + // `from_config` returns `Box`, which has no `Debug` + // impl -- build the concrete type directly to exercise its own + // `Debug` impl instead. + let yaml = single_rule_yaml("algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 5"); + let cfg: super::config::TokenRateLimitConfig = + praxis_filter::parse_filter_config("token_rate_limit", &yaml).unwrap(); + let backend = super::build_backend_resource(&cfg.backend).unwrap(); + let rules = cfg + .rules + .into_iter() + .map(|rule| super::compile_rule(rule, &backend)) + .collect::, _>>() + .unwrap(); + let filter = TokenRateLimitFilter { + rules, + epoch: std::time::Instant::now(), + }; + let debug = format!("{filter:?}"); + assert!(debug.contains("default"), "got: {debug}"); +} + +#[test] +fn from_config_accepts_custom_reservation_timeout() { + let yaml = single_rule_yaml( + "algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 5\nreservation_timeout: 10s", + ); + assert!(TokenRateLimitFilter::from_config(&yaml).is_ok()); +} + +#[test] +fn from_config_rejects_invalid_reservation_timeout() { + let yaml = single_rule_yaml( + "algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 5\nreservation_timeout: \ + not-a-duration", + ); + let err = TokenRateLimitFilter::from_config(&yaml).err().expect("should error"); + assert!(err.to_string().contains("invalid duration"), "got: {err}"); +} + +// ----------------------------------------------------------------------------- +// Per-rule algorithm choice (ai#789/praxis#551): mixed sliding_window and +// token_bucket rules, disambiguated by a header match -- the customer +// scenario this feature exists for (each app/team picks its own +// algorithm and budget). +// ----------------------------------------------------------------------------- + +/// Two rules, one per algorithm, matched by `x-app-id`: `alpha` gets a +/// tiny sliding-window budget, `beta` gets a tiny token-bucket budget. +fn two_algorithm_rules_yaml() -> serde_yaml::Value { + serde_yaml::from_str( + "rules:\n\ + \x20 - name: team-alpha\n\ + \x20 match:\n\ + \x20 headers:\n\ + \x20 x-app-id: alpha\n\ + \x20 algorithm: sliding_window\n\ + \x20 window: 1h\n\ + \x20 capacity: 100\n\ + \x20 reserved_tokens: 100\n\ + \x20 - name: team-beta\n\ + \x20 match:\n\ + \x20 headers:\n\ + \x20 x-app-id: beta\n\ + \x20 algorithm: token_bucket\n\ + \x20 capacity: 100\n\ + \x20 refill_rate: 1\n\ + \x20 reserved_tokens: 100\n", + ) + .unwrap() +} + +#[tokio::test] +async fn dispatches_to_the_first_matching_rule_by_algorithm_and_enforces_its_own_budget() { + let filter = TokenRateLimitFilter::from_config(&two_algorithm_rules_yaml()).unwrap(); + + let alpha_req = make_request_with_header("x-app-id", "alpha"); + let mut ctx = crate::test_utils::make_filter_context(&alpha_req); + assert!( + matches!(filter.on_request(&mut ctx).await.unwrap(), FilterAction::Continue), + "alpha's sliding-window rule should admit its first 100-token request" + ); + let mut ctx = crate::test_utils::make_filter_context(&alpha_req); + assert!( + matches!(filter.on_request(&mut ctx).await.unwrap(), FilterAction::Reject(_)), + "alpha's sliding-window budget is now exhausted" + ); + + // beta's independent token-bucket rule/budget is completely untouched + // by alpha's exhaustion, proving the two rules (and algorithms) are + // fully isolated from one another. + let beta_req = make_request_with_header("x-app-id", "beta"); + let mut ctx = crate::test_utils::make_filter_context(&beta_req); + assert!( + matches!(filter.on_request(&mut ctx).await.unwrap(), FilterAction::Continue), + "beta's token-bucket rule must be unaffected by alpha's exhausted sliding-window rule" + ); + let mut ctx = crate::test_utils::make_filter_context(&beta_req); + assert!( + matches!(filter.on_request(&mut ctx).await.unwrap(), FilterAction::Reject(_)), + "beta's token-bucket budget is now exhausted too" + ); +} + +/// A request matching neither rule's `match:` condition (e.g. a readiness +/// probe with no `x-app-id`) isn't rate limited by this filter instance at +/// all -- it's admitted without reserving against *any* rule's budget. +/// This is the documented mitigation for unrelated/non-inference traffic +/// under a scoped (non-catch-all) rule set (see `on_request`'s doc comment). +#[tokio::test] +async fn on_request_with_no_matching_rule_admits_without_reserving_any_budget() { + let filter = TokenRateLimitFilter::from_config(&two_algorithm_rules_yaml()).unwrap(); + + let unmatched = crate::test_utils::make_request(http::Method::GET, "/healthz"); + for _ in 0..5 { + let mut ctx = crate::test_utils::make_filter_context(&unmatched); + assert!( + matches!(filter.on_request(&mut ctx).await.unwrap(), FilterAction::Continue), + "a request matching no rule's `match:` condition must never be rejected by this filter" + ); + } + + // Prove the five unmatched requests above didn't silently draw down + // alpha's budget: it must still have its full 100-token capacity. + let alpha_req = make_request_with_header("x-app-id", "alpha"); + let mut ctx = crate::test_utils::make_filter_context(&alpha_req); + assert!( + matches!(filter.on_request(&mut ctx).await.unwrap(), FilterAction::Continue), + "alpha's full budget must be untouched by requests that matched no rule" + ); +} + +/// Two-rule config for [`reconciliation_settles_against_the_same_rule_that_admitted_the_request`]: +/// alpha (sliding window, capacity 5) and beta (token bucket, capacity +/// 100, `reserved_tokens` 40 -- smaller than capacity so a correct vs. +/// wrong/no-op credit-back is observably distinguishable). +fn team_alpha_sliding_and_team_beta_bucket_config() -> serde_yaml::Value { + serde_yaml::from_str( + "rules:\n\ + \x20 - name: team-alpha\n\ + \x20 match:\n\ + \x20 headers:\n\ + \x20 x-app-id: alpha\n\ + \x20 algorithm: sliding_window\n\ + \x20 window: 1h\n\ + \x20 capacity: 5\n\ + \x20 reserved_tokens: 5\n\ + \x20 - name: team-beta\n\ + \x20 match:\n\ + \x20 headers:\n\ + \x20 x-app-id: beta\n\ + \x20 algorithm: token_bucket\n\ + \x20 capacity: 100\n\ + \x20 refill_rate: 0.0001\n\ + \x20 reserved_tokens: 40\n", + ) + .unwrap() +} + +#[tokio::test] +async fn reconciliation_settles_against_the_same_rule_that_admitted_the_request() { + // Regression guard for the rule-index bookkeeping: reconciling a + // token-bucket-admitted request must credit back into *that same* + // rule's own bucket, not silently no-op or corrupt a different rule's + // (e.g. the sliding-window one's) state. + let filter = TokenRateLimitFilter::from_config(&team_alpha_sliding_and_team_beta_bucket_config()).unwrap(); + + let beta_req = make_request_with_header("x-app-id", "beta"); + // 100 - 40 - 40 = 20 remaining, then denied on a third 40-token ask. + let mut first_ctx = crate::test_utils::make_filter_context(&beta_req); + assert!(matches!( + filter.on_request(&mut first_ctx).await.unwrap(), + FilterAction::Continue + )); + assert!(matches!( + request_action(&*filter, &beta_req).await, + FilterAction::Continue + )); + assert!(matches!( + request_action(&*filter, &beta_req).await, + FilterAction::Reject(_) + )); + + // Reconcile the *first* reservation down to actual usage of 10 + // (refunding 30): if this credited the wrong rule, or no-op'd, beta's + // bucket would still be stuck at 20 and stay denied below. + first_ctx.set_metadata(META_TOKEN_TOTAL, "10"); + let mut body = None; + drop(filter.on_response_body(&mut first_ctx, &mut body, true).unwrap()); + + // 20 + 30 refund = 50 available, enough for one more 40-token request. + assert!( + matches!(request_action(&*filter, &beta_req).await, FilterAction::Continue), + "the refund from reconciling beta's own reservation must land in beta's own bucket" + ); + + // alpha's untouched sliding-window rule must still have its full + // capacity -- proving the refund didn't leak into the wrong rule. + let alpha_req = make_request_with_header("x-app-id", "alpha"); + assert!( + matches!(request_action(&*filter, &alpha_req).await, FilterAction::Continue), + "alpha's rule must be completely unaffected by beta's reconciliation" + ); +} + +#[test] +fn from_config_accepts_a_token_bucket_rule() { + let yaml = single_rule_yaml("algorithm: token_bucket\ncapacity: 100\nrefill_rate: 10\nreserved_tokens: 5"); + assert!(TokenRateLimitFilter::from_config(&yaml).is_ok()); +} + +/// Regression test for the actual reported vulnerability: `.nan`/`.inf` +/// parse cleanly from YAML via `serde_yaml` into an `f64` field with no +/// deserialization error, so this must be caught by `from_config`'s +/// validation, not just by unit tests that construct the Rust config +/// struct directly (which bypass the YAML layer entirely and can't catch +/// a future regression in the YAML-to-ledger wiring). +#[test] +fn from_config_rejects_non_finite_refill_rate_from_yaml() { + for literal in [".nan", ".inf", "-.inf"] { + let yaml = single_rule_yaml(&format!( + "algorithm: token_bucket\ncapacity: 100\nrefill_rate: {literal}\nreserved_tokens: 5" + )); + let err = TokenRateLimitFilter::from_config(&yaml) + .err() + .unwrap_or_else(|| panic!("refill_rate: {literal} must be rejected")); + assert!( + err.to_string().contains("refill_rate"), + "got: {err} for refill_rate: {literal}" + ); + } +} + +#[test] +fn from_config_rejects_a_refill_rate_that_would_overflow_the_valkey_reserve_scripts_pexpire_ttl() { + // capacity / refill_rate = 1e11 seconds -- a config typo away from + // plausible (e.g. an extra zero on refill_rate against a large + // capacity meant for a generous burst rule), not a contrived extreme. + // See MAX_CAPACITY_REFILL_RATE_RATIO_SECS's doc comment for why an + // unbounded ratio here is a real, silently budget-draining bug on the + // Valkey backend, not just a cosmetic validation gap. + let yaml = single_rule_yaml("algorithm: token_bucket\ncapacity: 1000000000\nrefill_rate: 0.01\nreserved_tokens: 5"); + let err = TokenRateLimitFilter::from_config(&yaml).err().expect("should error"); + assert!(err.to_string().contains("capacity / refill_rate"), "got: {err}"); +} + +#[tokio::test] +async fn token_bucket_rule_admits_within_capacity_and_denies_over_it() { + let yaml = single_rule_yaml("algorithm: token_bucket\ncapacity: 100\nrefill_rate: 1\nreserved_tokens: 100"); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut ctx = crate::test_utils::make_filter_context(&req); + assert!(matches!( + filter.on_request(&mut ctx).await.unwrap(), + FilterAction::Continue + )); + + let mut ctx = crate::test_utils::make_filter_context(&req); + assert!(matches!( + filter.on_request(&mut ctx).await.unwrap(), + FilterAction::Reject(_) + )); +} + +// ----------------------------------------------------------------------------- +// Valkey-backed shared state (final scenario: state shared across gateway +// instances/replicas) -- gated on a live Valkey/Redis instance via +// TOKEN_RATE_LIMIT_VALKEY_URL, skipped (not failed) when unset so +// contributors without a local Valkey aren't blocked. Set up locally with: +// brew install valkey && valkey-server --port 6400 --daemonize yes --save "" +// TOKEN_RATE_LIMIT_VALKEY_URL=redis://127.0.0.1:6400 cargo test -p praxis-ai-filters token_rate_limit +// ----------------------------------------------------------------------------- + +/// [`single_rule`], with a filter-level `backend: {kind: valkey}` block +/// prepended -- shared by the cross-instance/worker-reconciliation +/// scenarios below, which only vary the algorithm-specific rule body. +fn single_rule_valkey_yaml(algorithm_body: &str, url: &str, namespace: &str) -> serde_yaml::Value { + single_rule_yaml_with( + &format!("backend:\n kind: valkey\n url: {url}\n namespace: {namespace}"), + algorithm_body, + ) +} + +/// `filter.on_request` against a fresh context for one test request. +async fn request_action(filter: &dyn HttpFilter, req: &praxis_filter::Request) -> FilterAction { + let mut ctx = crate::test_utils::make_filter_context(req); + filter.on_request(&mut ctx).await.unwrap() +} + +/// Assert `req` is admitted by `filter`, with a business-behavior message +/// explaining why (for the many cross-instance/cross-algorithm Valkey +/// scenarios below). +async fn assert_admitted(filter: &dyn HttpFilter, req: &praxis_filter::Request, why: &str) { + assert!( + matches!(request_action(filter, req).await, FilterAction::Continue), + "{why}" + ); +} + +/// Assert `req` is denied (429) by `filter`, with a business-behavior +/// message explaining why. +async fn assert_denied(filter: &dyn HttpFilter, req: &praxis_filter::Request, why: &str) { + assert!( + matches!(request_action(filter, req).await, FilterAction::Reject(_)), + "{why}" + ); +} + +/// Poll `filter.on_request` for `req` up to `attempts` times, sleeping +/// briefly between each, until it's admitted. Used to await an +/// asynchronous (background-worker) Valkey reconciliation without a +/// fixed, flaky sleep. +async fn poll_until_admitted(filter: &dyn HttpFilter, req: &praxis_filter::Request, attempts: u32) -> bool { + for _ in 0..attempts { + if matches!(request_action(filter, req).await, FilterAction::Continue) { + return true; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + false +} + +#[tokio::test] +async fn valkey_budget_exhausted_on_one_instance_is_denied_on_another() { + let Ok(url) = std::env::var("TOKEN_RATE_LIMIT_VALKEY_URL") else { + tracing::warn!("skipping: TOKEN_RATE_LIMIT_VALKEY_URL not set"); + return; + }; + let namespace = format!("praxis-test-cross-instance-{}", std::process::id()); + let yaml = single_rule_valkey_yaml( + "algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 100", + &url, + &namespace, + ); + + // Two independent filter instances, exactly as two gateway replicas + // would each build their own filter from the same config. + let instance_one = TokenRateLimitFilter::from_config(&yaml).unwrap(); + let instance_two = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + assert_admitted(instance_one.as_ref(), &req, "admitted on instance one").await; + + // The *second* gateway instance must see the budget as already + // exhausted -- this is the property that makes Valkey worth the + // added complexity over in-process state (final scenario). + assert_denied( + instance_two.as_ref(), + &req, + "exhausted budget visible via shared Valkey state", + ) + .await; +} + +#[tokio::test] +async fn valkey_worker_reconciles_usage_off_the_response_path() { + let Ok(url) = std::env::var("TOKEN_RATE_LIMIT_VALKEY_URL") else { + tracing::warn!("skipping: TOKEN_RATE_LIMIT_VALKEY_URL not set"); + return; + }; + let namespace = format!("praxis-test-valkey-worker-{}", std::process::id()); + let yaml = single_rule_valkey_yaml( + "algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 50", + &url, + &namespace, + ); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut ctx = crate::test_utils::make_filter_context(&req); + let action = filter.on_request(&mut ctx).await.unwrap(); + assert!(matches!(action, FilterAction::Continue)); + ctx.set_metadata(META_TOKEN_TOTAL, "10"); // actual usage far below the 50-token estimate + + // Reconciliation for a Valkey backend is enqueued onto a background + // worker rather than awaited inline (the response must not be held + // up on a network round-trip that has no bearing on this request's + // own admission) -- so the freed budget becomes visible asynchronously. + let mut body = None; + let action = filter.on_response_body(&mut ctx, &mut body, true).unwrap(); + assert!(matches!(action, FilterAction::Continue)); + + // 10 (settled) + 85 should just fit under capacity=100 only once the + // worker has actually released the 40 unused reserved tokens (50 + // estimate - 10 actual); before that, 50 (still-active reservation) + // + 85 would exceed capacity and be denied. + let yaml_probe = single_rule_valkey_yaml( + "algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 85", + &url, + &namespace, + ); + let probe_filter = TokenRateLimitFilter::from_config(&yaml_probe).unwrap(); + + let settled = poll_until_admitted(probe_filter.as_ref(), &req, 40).await; + assert!( + settled, + "worker-based reconciliation should eventually release the unused reservation into the shared Valkey budget" + ); +} + +#[tokio::test] +async fn valkey_failure_fails_closed() { + // An unreachable backend (no server on this port) must reject, not + // admit -- a rate limiter that silently lets every request through + // when its state store is unavailable defeats the point of rate + // limiting it at all, right when a backend outage makes runaway + // spend/load most likely. + let yaml = single_rule_yaml_with( + "backend:\n kind: valkey\n url: redis://127.0.0.1:1", + "algorithm: sliding_window\nwindow: 1h\ncapacity: 100\nreserved_tokens: 10", + ); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut ctx = crate::test_utils::make_filter_context(&req); + match filter.on_request(&mut ctx).await.unwrap() { + FilterAction::Reject(rejection) => { + assert_eq!(rejection.status, 503, "unreachable backend should fail closed with 503"); + }, + other => panic!("unreachable Valkey backend must not admit the request, got {other:?}"), + } +} + +#[tokio::test] +async fn valkey_token_bucket_budget_exhausted_on_one_instance_is_denied_on_another() { + // The token-bucket analog of `valkey_budget_exhausted_on_one_instance_is_denied_on_another`: + // proves the *second* algorithm also gets the distributed-state + // property that's the whole point of the Valkey backend, not just + // the sliding-window one. + let Ok(url) = std::env::var("TOKEN_RATE_LIMIT_VALKEY_URL") else { + tracing::warn!("skipping: TOKEN_RATE_LIMIT_VALKEY_URL not set"); + return; + }; + let namespace = format!("praxis-test-tb-cross-instance-{}", std::process::id()); + let yaml = single_rule_valkey_yaml( + "algorithm: token_bucket\ncapacity: 100\nrefill_rate: 0.001\nreserved_tokens: 100", + &url, + &namespace, + ); + + let instance_one = TokenRateLimitFilter::from_config(&yaml).unwrap(); + let instance_two = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + assert_admitted(instance_one.as_ref(), &req, "admitted on instance one").await; + assert_denied( + instance_two.as_ref(), + &req, + "exhausted bucket visible via shared Valkey state", + ) + .await; +} + +#[tokio::test] +async fn valkey_token_bucket_failure_fails_closed() { + // The token-bucket analog of `valkey_failure_fails_closed`: an + // unreachable Valkey backend must not silently admit token-bucket + // requests either -- fail-closed has to hold for both algorithms, + // not just the sliding-window one it was first proven on. + let yaml = single_rule_yaml_with( + "backend:\n kind: valkey\n url: redis://127.0.0.1:1", + "algorithm: token_bucket\ncapacity: 100\nrefill_rate: 1\nreserved_tokens: 10", + ); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut ctx = crate::test_utils::make_filter_context(&req); + match filter.on_request(&mut ctx).await.unwrap() { + FilterAction::Reject(rejection) => { + assert_eq!(rejection.status, 503, "unreachable backend should fail closed with 503"); + }, + other => panic!("unreachable Valkey backend must not admit the token-bucket request, got {other:?}"), + } +} + +#[tokio::test] +async fn valkey_token_bucket_worker_reconciles_usage_off_the_response_path() { + // Token-bucket analog of `valkey_worker_reconciles_usage_off_the_response_path`: + // reconciliation is enqueued onto the background worker, not awaited + // inline, and its credit becomes visible once the worker runs. + let Ok(url) = std::env::var("TOKEN_RATE_LIMIT_VALKEY_URL") else { + tracing::warn!("skipping: TOKEN_RATE_LIMIT_VALKEY_URL not set"); + return; + }; + let namespace = format!("praxis-test-tb-worker-{}", std::process::id()); + let yaml = single_rule_valkey_yaml( + "algorithm: token_bucket\ncapacity: 100\nrefill_rate: 0.0001\nreserved_tokens: 50", + &url, + &namespace, + ); + let filter = TokenRateLimitFilter::from_config(&yaml).unwrap(); + + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat"); + let mut ctx = crate::test_utils::make_filter_context(&req); + let action = filter.on_request(&mut ctx).await.unwrap(); + assert!(matches!(action, FilterAction::Continue)); + ctx.set_metadata(META_TOKEN_TOTAL, "10"); // actual usage far below the 50-token estimate + + let mut body = None; + let action = filter.on_response_body(&mut ctx, &mut body, true).unwrap(); + assert!(matches!(action, FilterAction::Continue)); + + // 10 (settled) + 85 should just fit under capacity=100 only once the + // worker has actually credited back the 40 unused reserved tokens + // (50 estimate - 10 actual); before that, 50 (still-reserved) + 85 + // would exceed capacity and be denied. + let yaml_probe = single_rule_valkey_yaml( + "algorithm: token_bucket\ncapacity: 100\nrefill_rate: 0.0001\nreserved_tokens: 85", + &url, + &namespace, + ); + let probe_filter = TokenRateLimitFilter::from_config(&yaml_probe).unwrap(); + + let settled = poll_until_admitted(probe_filter.as_ref(), &req, 40).await; + assert!( + settled, + "worker-based reconciliation should eventually credit into the shared bucket" + ); +} diff --git a/filters/src/token_rate_limit/token_bucket_ledger.rs b/filters/src/token_rate_limit/token_bucket_ledger.rs new file mode 100644 index 0000000000..1491f4cf6e --- /dev/null +++ b/filters/src/token_rate_limit/token_bucket_ledger.rs @@ -0,0 +1,779 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Praxis Contributors + +//! In-process token-bucket reservation ledger. +//! +//! Reuses, unmodified, the refill formula from Praxis's own lock-free +//! `praxis_filter::builtins::http::traffic_management::token_bucket` +//! (`tokens = (tokens + elapsed_secs * rate).min(capacity)`), extended +//! with the reserve/reconcile split this filter needs: that lock-free +//! bucket is a single atomic decrement with no way to credit tokens +//! back, so a reservation that overestimates its actual cost would +//! strand unused capacity until the next natural refill. Reconcile here +//! explicitly credits back `estimate - actual` (or debits the +//! shortfall) once actual provider-reported usage is known, mirroring +//! [`super::ledger::Ledger::reconcile`]'s refund/overage semantics for +//! the sliding-window algorithm. +//! +//! Structurally mirrors [`super::ledger::Ledger`] (per-key locking, +//! active-reservation bookkeeping, bounded keys/reservations) so both +//! algorithms get the same operational guarantees; only the "how much +//! is available right now" computation differs (refill-and-cap here vs. +//! sum-in-window there). + +#![allow( + missing_docs, + clippy::missing_docs_in_private_items, + clippy::too_many_lines, + reason = "private ledger implementation is covered by its public filter contract and focused tests, matching \ + the sliding-window ledger's own module-level allow" +)] + +use std::{ + collections::HashMap, + sync::{ + Arc, Mutex, + atomic::{AtomicU64, AtomicUsize, Ordering}, + }, +}; + +use dashmap::DashMap; + +/// Upper bound, in seconds, on `capacity / refill_rate` -- the time to +/// fill an empty bucket from scratch. +/// +/// [`super::backend::TOKEN_BUCKET_RESERVE_SCRIPT`]'s Valkey/Lua path +/// folds this ratio into a millisecond `PEXPIRE` TTL. Lua 5.1's `%.14g` +/// number formatting switches to scientific notation past ~1e14, which +/// `PEXPIRE`'s strict-integer parser rejects -- and since Redis doesn't +/// roll back a script's earlier `redis.call()`s on a later error, that +/// failure would permanently drain the bucket instead of just denying +/// one request. Enforced here, shared by both backends, so a config +/// rejected on one is rejected on both. 1e9 stays ~1e5x below the +/// threshold, with margin for `reservation_timeout_ms` on top. +pub(super) const MAX_CAPACITY_REFILL_RATE_RATIO_SECS: f64 = 1e9; + +/// Upper bound on `capacity`/`reserved_tokens`, matching f64's 2^53 +/// mantissa: every `as f64` cast site in this module (refill math, +/// deficit/overage arithmetic) assumes its input stays below this via +/// `#[expect(cast_precision_loss)]`. A `capacity` above it would +/// silently lose precision instead of erroring. +pub(super) const MAX_F64_SAFE_INTEGER: u64 = 1 << 53; + +/// Bounds and parameters for a [`TokenBucketLedger`]. +#[derive(Clone, Debug)] +pub(super) struct TokenBucketConfig { + /// Maximum tokens held at once (the bucket's ceiling), per key. + pub(super) capacity: u64, + /// Tokens refilled per second, up to `capacity`. + pub(super) refill_rate: f64, + /// Time after which an ambiguous reservation is left charged at its + /// estimate (tokens were already decremented at reserve time; this + /// only bounds how long it's tracked as "active"). + pub(super) reservation_timeout_ms: u64, + /// Maximum logical keys retained by the ledger. + pub(super) max_keys: usize, + /// Maximum key length retained by the ledger. + pub(super) max_key_length: usize, + /// Maximum active reservations retained by the ledger. + pub(super) max_active_reservations: usize, +} + +/// Validate `capacity`/`refill_rate` bounds shared by both the in-memory +/// and Valkey token-bucket backends, so a config that's rejected on one +/// is rejected identically on the other (see [`MAX_F64_SAFE_INTEGER`] +/// and [`MAX_CAPACITY_REFILL_RATE_RATIO_SECS`] for why each bound +/// exists). +pub(super) fn validate_capacity_and_refill_rate(capacity: u64, refill_rate: f64) -> Result<(), String> { + if capacity == 0 { + return Err("capacity must be positive".into()); + } + if capacity > MAX_F64_SAFE_INTEGER { + return Err(format!("capacity must not exceed {MAX_F64_SAFE_INTEGER} (2^53)")); + } + if !refill_rate.is_finite() || refill_rate <= 0.0 { + return Err("refill_rate must be a positive, finite number".into()); + } + #[expect( + clippy::cast_precision_loss, + reason = "capacity is already checked above to not exceed f64's 2^53 mantissa" + )] + let capacity_f64 = capacity as f64; + if capacity_f64 / refill_rate > MAX_CAPACITY_REFILL_RATE_RATIO_SECS { + return Err(format!( + "capacity / refill_rate must not exceed {MAX_CAPACITY_REFILL_RATE_RATIO_SECS} seconds (time to fill \ + an empty bucket)" + )); + } + Ok(()) +} + +impl TokenBucketConfig { + /// Validate configuration before constructing a ledger. + pub(super) fn validate(&self) -> Result<(), String> { + validate_capacity_and_refill_rate(self.capacity, self.refill_rate)?; + if self.reservation_timeout_ms == 0 { + return Err("reservation timeout must be positive".into()); + } + if self.max_keys == 0 || self.max_key_length == 0 || self.max_active_reservations == 0 { + return Err("ledger bounds must be positive".into()); + } + Ok(()) + } +} + +/// A reservation admitted against one key's bucket. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) struct Reservation { + /// Opaque identifier used for idempotent reconciliation. + pub(super) id: u64, + /// Estimated token cost reserved (and already decremented) at + /// admission. + pub(super) estimate: u64, + /// Monotonic timestamp at admission, in milliseconds. + pub(super) created_at_ms: u64, +} + +/// Result of attempting admission. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) enum Decision { + /// Request may proceed with this reservation. + Admitted(Reservation), + /// Request must be rejected before routing. + Denied { + /// Conservative delay before the bucket refills enough to admit + /// the same estimate. + retry_after_ms: u64, + }, +} + +/// Result of reconciling a reservation. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) enum Settlement { + /// Actual usage was applied exactly once. + Applied { + /// Actual tokens charged to the bucket. + actual: u64, + /// Estimate credited back to the bucket. + refund: u64, + /// Usage above the estimate, debited from the bucket. + overage: u64, + }, + /// The reservation was already reconciled or conservatively expired. + Noop, +} + +#[derive(Debug)] +struct ActiveReservation { + estimate: u64, + created_at_ms: u64, +} + +#[derive(Debug)] +struct BucketState { + /// Current tokens, refilled lazily on access (same lazy-refill + /// pattern as the core lock-free `TokenBucket`). + tokens: f64, + last_refill_ms: u64, + active: HashMap, +} + +impl BucketState { + fn new(capacity: u64) -> Self { + Self { + #[expect( + clippy::cast_precision_loss, + reason = "token capacities are far below f64's 2^53 mantissa" + )] + tokens: capacity as f64, + last_refill_ms: 0, + active: HashMap::new(), + } + } + + /// Refill tokens for elapsed time, capped at `capacity`. Mirrors + /// `praxis_filter::builtins::http::traffic_management::token_bucket::TokenBucket::try_acquire`'s + /// refill formula exactly. + fn refill(&mut self, now_ms: u64, config: &TokenBucketConfig) { + let elapsed_ms = now_ms.saturating_sub(self.last_refill_ms); + if elapsed_ms > 0 { + #[expect( + clippy::cast_precision_loss, + reason = "elapsed_ms/capacity are far below f64's 2^53 mantissa for any realistic reservation window" + )] + { + let elapsed_secs = elapsed_ms as f64 / 1000.0; + self.tokens = (self.tokens + elapsed_secs * config.refill_rate).min(config.capacity as f64); + } + self.last_refill_ms = now_ms; + } + } + + /// Drop active reservations that exceeded `reservation_timeout_ms` + /// without being reconciled. + /// + /// Unlike the sliding-window ledger, no further charge happens here + /// -- the estimate was already decremented from the bucket at + /// reserve time (immediate-decrement design), so an abandoned + /// reservation is already charged. This only stops it from + /// permanently occupying the active-reservation bookkeeping. + fn reap(&mut self, now_ms: u64, config: &TokenBucketConfig) -> Vec { + let expired: Vec = self + .active + .iter() + .filter_map(|(id, reservation)| { + (now_ms.saturating_sub(reservation.created_at_ms) >= config.reservation_timeout_ms).then_some(*id) + }) + .collect(); + for id in &expired { + self.active.remove(id); + } + expired + } + + /// Whether this key's state can be safely forgotten: no reservations + /// pending reconciliation, *and* fully refilled back to capacity. + /// + /// Unlike the sliding-window ledger (where a settled entry simply + /// ages out of the window and becomes irrelevant), a token bucket's + /// balance never "expires" -- it only recovers via refill over time. + /// Evicting a key while it's still short of `capacity` would + /// silently grant it a full, unrecovered bucket the next time it's + /// touched (a fresh key always starts full) -- effectively free + /// bonus tokens. Only a fully-recovered, reservation-free key is + /// truly equivalent to "not tracked at all". + fn is_empty(&self, capacity: u64) -> bool { + #[expect( + clippy::cast_precision_loss, + reason = "token capacities are far below f64's 2^53 mantissa" + )] + let capacity = capacity as f64; + self.active.is_empty() && self.tokens >= capacity + } +} + +/// Thread-safe exact local token-bucket ledger with independent locks +/// per key. +pub(super) struct TokenBucketLedger { + config: TokenBucketConfig, + keys: DashMap>>, + reservations: DashMap, + next_id: AtomicU64, + key_count: AtomicUsize, + active_reservations: AtomicUsize, +} + +impl TokenBucketLedger { + /// Construct a validated empty ledger. + pub(super) fn new(config: TokenBucketConfig) -> Result { + config.validate()?; + Ok(Self { + config, + keys: DashMap::new(), + reservations: DashMap::new(), + next_id: AtomicU64::new(1), + key_count: AtomicUsize::new(0), + active_reservations: AtomicUsize::new(0), + }) + } + + /// The configured capacity, for bounded quota headers. + pub(super) fn limit(&self) -> u64 { + self.config.capacity + } + + /// Current number of active reservations. + pub(super) fn active_count(&self) -> usize { + self.active_reservations.load(Ordering::Relaxed) + } + + /// Current number of retained logical keys. + pub(super) fn key_count(&self) -> usize { + self.key_count.load(Ordering::Relaxed) + } + + /// Reserve an estimate against one key's bucket: refill to `now_ms`, + /// admit and immediately decrement if enough tokens are available, + /// deny otherwise. + pub(super) fn reserve(&self, key: &str, estimate: u64, now_ms: u64) -> Decision { + if key.is_empty() || key.len() > self.config.max_key_length || estimate == 0 { + return Decision::Denied { retry_after_ms: 0 }; + } + + let state = match self.keys.entry(key.to_owned()) { + dashmap::mapref::entry::Entry::Occupied(entry) => Arc::clone(entry.get()), + dashmap::mapref::entry::Entry::Vacant(entry) => { + if self + .key_count + .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |count| { + (count < self.config.max_keys).then_some(count + 1) + }) + .is_err() + { + return Decision::Denied { retry_after_ms: 0 }; + } + let state = Arc::new(Mutex::new(BucketState::new(self.config.capacity))); + entry.insert(Arc::clone(&state)); + state + }, + }; + let mut state = match state.lock() { + Ok(state) => state, + Err(poisoned) => poisoned.into_inner(), + }; + state.refill(now_ms, &self.config); + let expired = state.reap(now_ms, &self.config); + for id in &expired { + self.reservations.remove(id); + } + self.active_reservations.fetch_sub(expired.len(), Ordering::Relaxed); + + #[expect( + clippy::cast_precision_loss, + reason = "token estimates are far below f64's 2^53 mantissa" + )] + let estimate_f64 = estimate as f64; + if estimate_f64 > state.tokens { + let deficit = estimate_f64 - state.tokens; + let retry_after_ms = (deficit / self.config.refill_rate * 1000.0).ceil(); + #[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "retry_after_ms is a small positive duration bounded by realistic refill rates" + )] + let retry_after_ms = retry_after_ms.max(1.0) as u64; + return Decision::Denied { retry_after_ms }; + } + if self + .active_reservations + .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |active| { + (active < self.config.max_active_reservations).then_some(active + 1) + }) + .is_err() + { + return Decision::Denied { + retry_after_ms: self.config.reservation_timeout_ms, + }; + } + + state.tokens -= estimate_f64; + let id = self.next_id.fetch_add(1, Ordering::Relaxed); + state.active.insert( + id, + ActiveReservation { + estimate, + created_at_ms: now_ms, + }, + ); + self.reservations.insert(id, key.to_owned()); + drop(state); + Decision::Admitted(Reservation { + id, + estimate, + created_at_ms: now_ms, + }) + } + + /// Reconcile actual usage: refill to `now_ms`, then credit back + /// `estimate - actual` (capped at capacity) or debit the shortfall + /// (floored at 0). Repeated calls for one ID are no-ops. + pub(super) fn reconcile(&self, id: u64, actual: Option, now_ms: u64) -> Settlement { + let Some((_, key)) = self.reservations.remove(&id) else { + return Settlement::Noop; + }; + let Some(state) = self.keys.get(&key).map(|entry| Arc::clone(entry.value())) else { + return Settlement::Noop; + }; + let mut state = match state.lock() { + Ok(state) => state, + Err(poisoned) => poisoned.into_inner(), + }; + let Some(reservation) = state.active.remove(&id) else { + return Settlement::Noop; + }; + self.active_reservations.fetch_sub(1, Ordering::Relaxed); + state.refill(now_ms, &self.config); + + let actual = actual.unwrap_or(reservation.estimate); + let refund = reservation.estimate.saturating_sub(actual); + let overage = actual.saturating_sub(reservation.estimate); + #[expect( + clippy::cast_precision_loss, + reason = "token deltas are far below f64's 2^53 mantissa" + )] + { + if refund > 0 { + state.tokens = (state.tokens + refund as f64).min(self.config.capacity as f64); + } else if overage > 0 { + state.tokens = (state.tokens - overage as f64).max(0.0); + } + } + drop(state); + Settlement::Applied { + actual, + refund, + overage, + } + } + + /// Conservatively expire a bounded number of keys and reclaim idle + /// state. + pub(super) fn cleanup(&self, now_ms: u64, max_keys_to_scan: usize) -> usize { + let mut orphaned = 0; + let keys: Vec = self + .keys + .iter() + .take(max_keys_to_scan) + .map(|entry| entry.key().clone()) + .collect(); + for key in keys { + let Some(entry) = self.keys.get_mut(&key) else { + continue; + }; + let state_arc = Arc::clone(entry.value()); + let mut state = match state_arc.lock() { + Ok(state) => state, + Err(poisoned) => poisoned.into_inner(), + }; + state.refill(now_ms, &self.config); + let expired = state.reap(now_ms, &self.config); + orphaned += expired.len(); + for id in &expired { + self.reservations.remove(id); + } + self.active_reservations.fetch_sub(expired.len(), Ordering::Relaxed); + let empty = state.is_empty(self.config.capacity); + drop(state); + drop(entry); + if empty + && self + .keys + .remove_if(&key, |_, candidate| { + Arc::ptr_eq(candidate, &state_arc) + && match candidate.lock() { + Ok(state) => state.is_empty(self.config.capacity), + Err(poisoned) => poisoned.into_inner().is_empty(self.config.capacity), + } + }) + .is_some() + { + self.key_count.fetch_sub(1, Ordering::Relaxed); + } + } + orphaned + } +} + +#[cfg(test)] +#[expect(clippy::allow_attributes, reason = "blanket test suppressions")] +#[allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::indexing_slicing, + clippy::panic, + clippy::manual_let_else, + clippy::match_wildcard_for_single_variants, + reason = "ledger tests intentionally fail fast on impossible fixture states" +)] +mod tests { + use std::sync::{Arc, Barrier}; + + use super::*; + + fn ledger(capacity: u64, refill_rate: f64) -> TokenBucketLedger { + TokenBucketLedger::new(TokenBucketConfig { + capacity, + refill_rate, + reservation_timeout_ms: 100, + max_keys: 8, + max_key_length: 256, + max_active_reservations: 32, + }) + .unwrap() + } + + #[test] + fn admits_within_capacity_and_denies_over_capacity() { + let l = ledger(10, 1.0); + assert!(matches!(l.reserve("a", 10, 0), Decision::Admitted(_))); + assert!(matches!(l.reserve("a", 1, 0), Decision::Denied { .. })); + } + + #[test] + fn refills_over_time_at_the_configured_rate() { + let l = ledger(10, 10.0); // 10 tokens/sec + assert!(matches!(l.reserve("a", 10, 0), Decision::Admitted(_))); + assert!(matches!(l.reserve("a", 1, 0), Decision::Denied { .. })); + // 200ms at 10 tokens/sec refills 2 tokens. + assert!(matches!(l.reserve("a", 2, 200), Decision::Admitted(_))); + } + + #[test] + fn refill_never_exceeds_capacity() { + let l = ledger(5, 1000.0); + assert!(matches!(l.reserve("a", 1, 0), Decision::Admitted(_))); + // A huge amount of elapsed time should cap at capacity (5), not overflow. + assert!(matches!(l.reserve("a", 5, 1_000_000), Decision::Admitted(_))); + assert!(matches!(l.reserve("a", 1, 1_000_000), Decision::Denied { .. })); + } + + #[test] + fn reconcile_credits_back_unused_estimate_on_overestimate() { + let l = ledger(100, 1.0); + let r = match l.reserve("a", 50, 0) { + Decision::Admitted(r) => r, + other => panic!("expected admission, got {other:?}"), + }; + assert_eq!( + l.reconcile(r.id, Some(20), 0), + Settlement::Applied { + actual: 20, + refund: 30, + overage: 0 + } + ); + // 100 - 50 (reserved) + 30 (refund) = 80 available. + assert!(matches!(l.reserve("a", 80, 0), Decision::Admitted(_))); + } + + #[test] + fn reconcile_debits_the_shortfall_on_underestimate() { + let l = ledger(100, 1.0); + let r = match l.reserve("a", 50, 0) { + Decision::Admitted(r) => r, + other => panic!("expected admission, got {other:?}"), + }; + assert_eq!( + l.reconcile(r.id, Some(70), 0), + Settlement::Applied { + actual: 70, + refund: 0, + overage: 20 + } + ); + // 100 - 50 (reserved) - 20 (extra overage debit) = 30 available. + assert!(matches!(l.reserve("a", 30, 0), Decision::Admitted(_))); + assert!( + matches!(l.reserve("b", 1, 0), Decision::Admitted(_)), + "other keys unaffected" + ); + } + + #[test] + fn overage_debit_floors_at_zero_rather_than_going_negative() { + let l = ledger(10, 1.0); + let r = match l.reserve("a", 5, 0) { + Decision::Admitted(r) => r, + other => panic!("expected admission, got {other:?}"), + }; + // actual (100) wildly exceeds both the estimate and the bucket's + // total capacity -- the bucket must floor at 0, not panic or + // wrap on the unsigned subtraction. + assert_eq!( + l.reconcile(r.id, Some(100), 0), + Settlement::Applied { + actual: 100, + refund: 0, + overage: 95 + } + ); + assert!(matches!(l.reserve("a", 1, 0), Decision::Denied { .. })); + } + + #[test] + fn duplicate_reconciliation_is_noop() { + let l = ledger(100, 1.0); + let r = match l.reserve("a", 5, 0) { + Decision::Admitted(r) => r, + other => panic!("expected admission, got {other:?}"), + }; + assert!(matches!( + l.reconcile(r.id, None, 0), + Settlement::Applied { actual: 5, .. } + )); + assert_eq!(l.reconcile(r.id, Some(99), 0), Settlement::Noop); + } + + #[test] + fn keys_are_independent_and_bounded() { + let l = ledger(10, 1.0); + assert!(matches!(l.reserve("a", 10, 0), Decision::Admitted(_))); + assert!(matches!(l.reserve("b", 10, 0), Decision::Admitted(_))); + assert!(matches!(l.reserve("", 1, 0), Decision::Denied { .. })); + } + + #[test] + fn abandoned_reservation_stays_charged_after_timeout() { + // A lost request (never reconciled) must not free its reserved + // tokens back into the bucket just because it timed out -- the + // tokens were already spent at reserve time under the + // immediate-decrement design, so "abandoned" must mean "stays + // charged", matching the sliding-window ledger's equivalent + // guarantee. + let l = ledger(10, 1.0); + assert!(matches!(l.reserve("a", 10, 0), Decision::Admitted(_))); + assert_eq!(l.active_count(), 1); + l.cleanup(200, 8); // past reservation_timeout_ms=100 + assert_eq!(l.active_count(), 0, "orphan should be reaped from active tracking"); + // Capacity is still fully charged (no refill credited for the + // abandoned reservation) -- only the 1 token/sec natural refill + // over 200ms (0.2 tokens) is available, nowhere near 10. + assert!(matches!(l.reserve("a", 1, 200), Decision::Denied { .. })); + } + + #[test] + fn idle_settled_keys_are_reclaimed_once_fully_refilled() { + let l = ledger(10, 1000.0); // fast refill so full recovery is reachable in the test + let r = match l.reserve("a", 10, 0) { + Decision::Admitted(r) => r, + other => panic!("expected admission, got {other:?}"), + }; + l.reconcile(r.id, Some(10), 0); + assert_eq!(l.key_count(), 1); + // Not yet reclaimed: still short of full capacity at the same instant. + l.cleanup(0, 8); + assert_eq!( + l.key_count(), + 1, + "a not-yet-refilled key must not be evicted (would grant free bonus tokens)" + ); + + // 100ms at 1000 tokens/sec fully refills the 10-token bucket. + l.cleanup(100, 8); + assert_eq!( + l.key_count(), + 0, + "a fully-refilled, reservation-free key can be safely forgotten" + ); + } + + #[test] + fn concurrent_same_key_admission_cannot_oversubscribe() { + let ledger = Arc::new(ledger(100, 0.000_001)); // negligible refill during the test window + let barrier = Arc::new(Barrier::new(16)); + let handles = (0..16) + .map(|_| { + let ledger = Arc::clone(&ledger); + let barrier = Arc::clone(&barrier); + std::thread::spawn(move || { + barrier.wait(); + matches!(ledger.reserve("same", 10, 0), Decision::Admitted(_)) + }) + }) + .collect::>(); + let admitted = handles + .into_iter() + .filter_map(|handle| handle.join().ok()) + .filter(|ok| *ok) + .count(); + assert_eq!(admitted, 10, "exactly the capacity should be admitted"); + } + + #[test] + fn invalid_config_is_rejected() { + assert!( + TokenBucketLedger::new(TokenBucketConfig { + capacity: 0, + refill_rate: 1.0, + reservation_timeout_ms: 1, + max_keys: 1, + max_key_length: 1, + max_active_reservations: 1, + }) + .is_err() + ); + assert!( + TokenBucketLedger::new(TokenBucketConfig { + capacity: 1, + refill_rate: 0.0, + reservation_timeout_ms: 1, + max_keys: 1, + max_key_length: 1, + max_active_reservations: 1, + }) + .is_err() + ); + } + + #[test] + fn non_finite_refill_rate_is_rejected() { + // `refill_rate <= 0.0` alone is always false for NaN/+-Infinity. + for bad_rate in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { + assert!( + TokenBucketLedger::new(TokenBucketConfig { + capacity: 1, + refill_rate: bad_rate, + reservation_timeout_ms: 1, + max_keys: 1, + max_key_length: 1, + max_active_reservations: 1, + }) + .is_err(), + "refill_rate {bad_rate} must be rejected as non-finite/non-positive" + ); + } + } + + #[test] + fn zero_reservation_timeout_is_rejected() { + assert!( + TokenBucketLedger::new(TokenBucketConfig { + capacity: 1, + refill_rate: 1.0, + reservation_timeout_ms: 0, + max_keys: 1, + max_key_length: 1, + max_active_reservations: 1, + }) + .is_err() + ); + } + + #[test] + fn zero_ledger_bounds_are_rejected() { + for (max_keys, max_key_length, max_active_reservations) in [(0, 1, 1), (1, 0, 1), (1, 1, 0)] { + assert!( + TokenBucketLedger::new(TokenBucketConfig { + capacity: 1, + refill_rate: 1.0, + reservation_timeout_ms: 1, + max_keys, + max_key_length, + max_active_reservations, + }) + .is_err(), + "bounds ({max_keys}, {max_key_length}, {max_active_reservations}) must be rejected" + ); + } + } + + #[test] + fn key_capacity_denies_new_keys_beyond_the_configured_limit() { + let l = TokenBucketLedger::new(TokenBucketConfig { + capacity: 100, + refill_rate: 1.0, + reservation_timeout_ms: 100, + max_keys: 1, + max_key_length: 256, + max_active_reservations: 32, + }) + .unwrap(); + assert!(matches!(l.reserve("a", 1, 0), Decision::Admitted(_))); + assert!(matches!(l.reserve("b", 1, 0), Decision::Denied { .. })); + } + + #[test] + fn reservation_capacity_denies_beyond_the_configured_limit() { + let l = TokenBucketLedger::new(TokenBucketConfig { + capacity: 1_000, + refill_rate: 1.0, + reservation_timeout_ms: 100, + max_keys: 32, + max_key_length: 256, + max_active_reservations: 1, + }) + .unwrap(); + assert!(matches!(l.reserve("a", 1, 0), Decision::Admitted(_))); + assert!(matches!(l.reserve("b", 1, 0), Decision::Denied { .. })); + } +} diff --git a/filters/src/token_usage/mod.rs b/filters/src/token_usage/mod.rs index e468f4c73e..b437a14cd6 100644 --- a/filters/src/token_usage/mod.rs +++ b/filters/src/token_usage/mod.rs @@ -23,6 +23,18 @@ const META_TOKEN_INPUT: &str = "token.input"; const META_TOKEN_OUTPUT: &str = "token.output"; /// Metadata key for the total token count. +/// +/// `pub(crate)` only under `token-rate-limit-filter`, so that experimental +/// filter's reconciliation path can reference this constant directly +/// instead of duplicating the string literal — see the duplication risk +/// this avoids: [`ai#351`](https://github.com/praxis-proxy/ai/issues/351) +/// (cached-token double-counting caused by a second, independent parsing +/// path drifting from this one). Private otherwise, so this stable filter's +/// public surface is unaffected when the experimental feature is disabled. +#[cfg(feature = "token-rate-limit-filter")] +pub(crate) const META_TOKEN_TOTAL: &str = "token.total"; +/// Metadata key for the total token count. +#[cfg(not(feature = "token-rate-limit-filter"))] const META_TOKEN_TOTAL: &str = "token.total"; /// Metadata key signaling that usage could not be captured because the diff --git a/server/Cargo.toml b/server/Cargo.toml index db219d108a..59d9755288 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -29,6 +29,8 @@ http-callout-filter = ["praxis-ai-filters/http-callout-filter", "experimental"] azure-ad-filter = ["praxis-ai-filters/azure-ad-filter", "experimental"] # Experimental: the gcp_adc (GCP ADC) auth filter (see praxis-ai-filters). gcp-adc-filter = ["praxis-ai-filters/gcp-adc-filter", "experimental"] +# Experimental: the token_rate_limit filter (see praxis-ai-filters). +token-rate-limit-filter = ["praxis-ai-filters/token-rate-limit-filter", "experimental"] # Marker feature activated transitively by any experimental feature. experimental = ["praxis-ai-filters/experimental"] opentelemetry = ["praxis-ai-filters/opentelemetry"] diff --git a/tests/integration/Cargo.toml b/tests/integration/Cargo.toml index e31c459590..f4162f74f1 100644 --- a/tests/integration/Cargo.toml +++ b/tests/integration/Cargo.toml @@ -18,6 +18,9 @@ azure-ad-filter = ["praxis-ai-proxy/azure-ad-filter"] # Experimental: gates the gcp_adc example test, which needs the gcp_adc # filter registered in the in-process proxy registry. gcp-adc-filter = ["praxis-ai-proxy/gcp-adc-filter"] +# Experimental: gates the token_rate_limit example tests, which need the +# token_rate_limit filter registered in the in-process proxy registry. +token-rate-limit-filter = ["praxis-ai-proxy/token-rate-limit-filter"] llmd-ext-proc = ["praxis-test-utils/llmd-ext-proc"] no-mac-cert-rotation-tests = [] # We use this to disable testing cert rotation on macOS praxis-main = ["praxis-test-utils/praxis-main"] From d1628c1b361b0609eeaf4db348dcc8d90aae2f7a Mon Sep 17 00:00:00 2001 From: Jordi Gil Date: Fri, 28 Aug 2026 21:45:58 -0400 Subject: [PATCH 2/5] test(token_rate_limit): add integration tests and example configs Covers single-rule and mixed-algorithm (sliding_window + token_bucket) setups end to end, including the Valkey-backed path that isolates budgets across gateway replicas. Also adds the runnable example configs referenced from the filter's docs. Signed-off-by: Jordi Gil --- examples/README.md | 2 + .../token-rate-limit-mixed-algorithms.yaml | 68 +++ examples/configs/token-rate-limit.yaml | 92 ++++ tests/integration/tests/suite/examples/mod.rs | 2 + .../tests/suite/examples/token_rate_limit.rs | 396 ++++++++++++++++++ 5 files changed, 560 insertions(+) create mode 100644 examples/configs/token-rate-limit-mixed-algorithms.yaml create mode 100644 examples/configs/token-rate-limit.yaml create mode 100644 tests/integration/tests/suite/examples/token_rate_limit.rs diff --git a/examples/README.md b/examples/README.md index 6faacef86f..6849626ef0 100644 --- a/examples/README.md +++ b/examples/README.md @@ -42,6 +42,8 @@ before sending requests. | [provider-route.yaml](configs/provider-route.yaml) | This listener requires downstream mTLS. `peer_identity_trust` authenticates and authorizes the edge gateway before AI-owned x-ai-routing-* fields can influence provider-local routing | | [time-to-first-token.yaml](configs/time-to-first-token.yaml) | Measures the elapsed time from request receipt to the first non-empty SSE body chunk and records a praxis_ai_ttft_seconds Prometheus histogram labeled by model | | [token-counting.yaml](configs/token-counting.yaml) | Extracts token usage from AI inference responses (streaming and non-streaming) and makes counts available to downstream filters via filter metadata as token.input, token.output, and token.total | +| [token-rate-limit-mixed-algorithms.yaml](configs/token-rate-limit-mixed-algorithms.yaml) | Extends token-rate-limit.yaml with per-rule algorithm choice (ai#789 / praxis#551): each rule in `rules:` independently picks sliding_window or token_bucket, matched by a static header value. team-alpha gets an exact trailing-window budget; team-beta gets a continuously-refilling bucket | +| [token-rate-limit.yaml](configs/token-rate-limit.yaml) | Reserves an estimated token cost at admission time and reconciles that reservation against actual provider-reported usage once the response completes | | [token-usage-headers.yaml](configs/token-usage-headers.yaml) | Inject Praxis-Token-Input, Praxis-Token-Output, and Praxis-Token-Total headers into downstream responses when token counts are available in filter metadata | ### Anthropic diff --git a/examples/configs/token-rate-limit-mixed-algorithms.yaml b/examples/configs/token-rate-limit-mixed-algorithms.yaml new file mode 100644 index 0000000000..bd36d369ab --- /dev/null +++ b/examples/configs/token-rate-limit-mixed-algorithms.yaml @@ -0,0 +1,68 @@ +# Token Rate Limiting -- Mixed Algorithms Per Rule +# +# Extends token-rate-limit.yaml with per-rule algorithm choice (ai#789 / +# praxis#551): each rule in `rules:` independently picks sliding_window +# or token_bucket, matched by a static header value. team-alpha gets an +# exact trailing-window budget; team-beta gets a continuously-refilling +# bucket. Requests matching neither rule are not rate limited by this +# filter instance. +# +# Usage: +# cargo run -p praxis-ai-proxy -- -c examples/configs/token-rate-limit-mixed-algorithms.yaml +# curl -i http://localhost:8080/v1/chat/completions -H 'x-app-id: alpha' -d '{"model":"gpt-4","messages":[{"role":"user","content":"hi"}]}' +# curl -i http://localhost:8080/v1/chat/completions -H 'x-app-id: beta' -d '{"model":"gpt-4","messages":[{"role":"user","content":"hi"}]}' +# +# See token-rate-limit.yaml for the full rationale on reservation-based +# admission (M2), reconciliation, the `rules:` schema, and what's still +# deliberately out of scope (composite/CEL keys, configurable +# estimation, token-type-aware accounting, observability/metering). +# +# Supported token_count provider values: +# openai | anthropic | google | bedrock | azure + +listeners: + - name: default + address: "127.0.0.1:8080" + filter_chains: + - main + +filter_chains: + - name: main + filters: + - filter: router + routes: + - path: "/v1/chat/completions" + cluster: backend + + - filter: token_rate_limit + rules: + - name: team-alpha + match: + headers: + x-app-id: alpha + algorithm: sliding_window + window: 1h # exact trailing-window budget + capacity: 100000 # max tokens admitted within `window` + reserved_tokens: 500 # fixed cost reserved per request at admission + - name: team-beta + match: + headers: + x-app-id: beta + algorithm: token_bucket + capacity: 50000 # max tokens held at once + refill_rate: 50 # tokens refilled per second, up to `capacity` + reserved_tokens: 200 + + - filter: token_count + provider: openai # openai | anthropic | google | bedrock | azure + + - filter: access_log + + - filter: load_balancer + clusters: + - name: backend + endpoints: + - "127.0.0.1:3000" + +insecure_options: + allow_private_endpoints: true # example proxies to a local backend diff --git a/examples/configs/token-rate-limit.yaml b/examples/configs/token-rate-limit.yaml new file mode 100644 index 0000000000..22b7ac66fe --- /dev/null +++ b/examples/configs/token-rate-limit.yaml @@ -0,0 +1,92 @@ +# Token Rate Limiting +# +# Reserves an estimated token cost at admission time and reconciles +# that reservation against actual provider-reported usage once the +# response completes. Rejects with 429 when the bucket can't cover +# the estimate. +# +# Usage: +# cargo run -p praxis-ai-proxy -- -c examples/configs/token-rate-limit.yaml +# curl -i http://localhost:8080/v1/chat/completions -d '{"model":"gpt-4","messages":[{"role":"user","content":"hi"}]}' +# +# This is the agreed M1/M2/M6 core of the token rate limiting proposal +# (00121_token-rate-limiting.md in praxis-proxy/enhancements, tracked by +# epic ai#121): +# - M1: a single sliding-window budget (one catch-all rule) +# - M2: reservation-based admission, reconciled against actual usage +# - M6: 429 responses with token-denominated rate limit headers +# +# `rules:` (per `ai#789`/`praxis#551`) lets each rule pick its own +# admission algorithm and match condition; a rule with no `match:` is a +# catch-all applying to every request that reaches it. This example +# uses one such catch-all rule. `window`/`capacity` match the proposal's +# own field names: "Windows are sliding: a `window: 1h` budget tracks +# usage in the most recent 60 minutes from the current instant." State +# defaults to in-process; set a filter-level `backend: {kind: valkey, +# url: ..., namespace: ...}` (a sibling of `rules:`, not per-rule) to +# share every rule's budget across every gateway instance/replica over +# one shared connection. +# +# Deliberately not yet built (see filters/src/token_rate_limit/mod.rs +# for the full rationale): CEL/composite bucket keys, configurable +# estimation strategies (M3, `reserved_tokens` below is a fixed +# placeholder), token-type-aware accounting (M4), and +# observability/metering (M7/M8/S3). +# +# token_rate_limit is declared *before* token_count: response hooks +# run in reverse declared order, so token_count's on_response_body +# (which writes token.total to filter_metadata) runs before +# token_rate_limit's on_response_body reads it back to reconcile the +# reservation. This mirrors the token_usage_headers/token_count +# ordering in examples/configs/token-counting.yaml. +# +# Assumes request identity is already resolved upstream (this filter +# doesn't authenticate callers) -- a catch-all rule like the one below +# reserves quota for every request that reaches it, including probes +# and health checks. Scope with an explicit `match:` per rule, or place +# an identity/auth filter earlier in the pipeline. Tracked in grid#101. +# +# The X-RateLimit-*-Tokens response headers always reflect the +# reservation-time snapshot, never this response's own reconciliation +# — headers are committed before the body (and thus actual usage) is +# known. Reconciliation still affects every *subsequent* request's +# admission decision. +# +# Supported token_count provider values: +# openai | anthropic | google | bedrock | azure + +listeners: + - name: default + address: "127.0.0.1:8080" + filter_chains: + - main + +filter_chains: + - name: main + filters: + - filter: router + routes: + - path: "/v1/chat/completions" + cluster: backend + + - filter: token_rate_limit + rules: + - name: default + algorithm: sliding_window + window: 1h # sliding window duration + capacity: 100000 # max tokens admitted within `window` + reserved_tokens: 500 # fixed cost reserved per request at admission + + - filter: token_count + provider: openai # openai | anthropic | google | bedrock | azure + + - filter: access_log + + - filter: load_balancer + clusters: + - name: backend + endpoints: + - "127.0.0.1:3000" + +insecure_options: + allow_private_endpoints: true # example proxies to a local backend diff --git a/tests/integration/tests/suite/examples/mod.rs b/tests/integration/tests/suite/examples/mod.rs index e72b5e5522..1c3427b716 100644 --- a/tests/integration/tests/suite/examples/mod.rs +++ b/tests/integration/tests/suite/examples/mod.rs @@ -53,6 +53,8 @@ mod session_replay; mod time_to_first_token; mod token_count; mod token_counting; +#[cfg(feature = "token-rate-limit-filter")] +mod token_rate_limit; mod token_usage_headers; mod vector_stores_routing; mod vllm_agentic_api; diff --git a/tests/integration/tests/suite/examples/token_rate_limit.rs b/tests/integration/tests/suite/examples/token_rate_limit.rs new file mode 100644 index 0000000000..580bdd93df --- /dev/null +++ b/tests/integration/tests/suite/examples/token_rate_limit.rs @@ -0,0 +1,396 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Praxis Contributors + +//! Integration tests for the `token_rate_limit` filter's example config. +//! +//! Covers the agreed M1/M2/M6 core of the token rate limiting proposal +//! (`00121_token-rate-limiting.md` in `praxis-proxy/enhancements`, epic +//! `ai#121`): reservation-based admission, 429 rejection with +//! token-denominated headers, and reconciliation against actual +//! provider-reported usage (`token_count`'s `token.total`) once the +//! response completes. +//! +//! `mixed_algorithm_rules_valkey_backend_isolates_budgets_across_gateway_replicas` +//! additionally covers `ai#789`/`praxis#551`'s per-rule algorithm choice +//! (`rules:`/`match:`/`algorithm:`) end-to-end through the real +//! [`praxis_filter::HttpFilter`] pipeline, gated on a live Valkey/Redis +//! instance the same way `filters/src/token_rate_limit/tests.rs`'s +//! unit-level Valkey tests are. + +use std::collections::HashMap; + +use praxis_test_utils::{ + Backend, example_config_path, free_port, http_send, json_post, load_example_config, parse_body, parse_header, + parse_status, patch_yaml, start_proxy, +}; + +/// Build a `POST` request carrying extra headers beyond the standard +/// JSON content-type/length, for `match`-based rule dispatch scenarios +/// that need to tag requests with an app identity. +fn json_post_with_headers(path: &str, body: &str, headers: &[(&str, &str)]) -> String { + let mut extra = String::new(); + for (name, value) in headers { + extra.push_str(&format!("{name}: {value}\r\n")); + } + format!( + "POST {path} HTTP/1.1\r\n\ + Host: localhost\r\n\ + Content-Type: application/json\r\n\ + Content-Length: {}\r\n\ + {extra}\ + Connection: close\r\n\r\n\ + {body}", + body.len() + ) +} + +// ----------------------------------------------------------------------------- +// Mock response bodies +// ----------------------------------------------------------------------------- + +/// OpenAI-shaped response reporting 10 total tokens used. +const OPENAI_LOW_USAGE_JSON: &str = + r#"{"choices":[],"usage":{"prompt_tokens":5,"completion_tokens":5,"total_tokens":10}}"#; + +/// Plain-text response with no token usage: `token_count` extracts +/// nothing, so `token_rate_limit`'s reservation is never reconciled. +const PLAIN_TEXT_BODY: &str = "ok"; + +// ----------------------------------------------------------------------------- +// Test Utilities +// ----------------------------------------------------------------------------- + +/// Build a YAML config for the token rate limiting pipeline using the +/// example file, substituting `capacity`/`reserved_tokens` with the given +/// values so tests can exercise small, deterministic budgets. +fn token_rate_limit_config( + proxy_port: u16, + backend_port: u16, + capacity: u64, + reserved_tokens: u64, +) -> praxis_core::config::Config { + let path = example_config_path("token-rate-limit.yaml"); + let yaml = std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}")); + let yaml = yaml + .replace("capacity: 100000", &format!("capacity: {capacity}")) + .replace("reserved_tokens: 500", &format!("reserved_tokens: {reserved_tokens}")); + let patched = patch_yaml(&yaml, proxy_port, &HashMap::from([("127.0.0.1:3000", backend_port)])); + praxis_core::config::Config::from_yaml(&patched).expect("config should parse") +} + +// ----------------------------------------------------------------------------- +// Admission and headers +// ----------------------------------------------------------------------------- + +#[test] +fn admits_request_within_budget() { + let backend = Backend::fixed(PLAIN_TEXT_BODY) + .header("content-type", "text/plain") + .start_with_shutdown(); + let proxy_port = free_port(); + let config = token_rate_limit_config(proxy_port, backend.port(), 100_000, 500); + let proxy = start_proxy(&config); + + let raw = http_send(proxy.addr(), &json_post("/v1/chat/completions", "{}")); + assert_eq!(parse_status(&raw), 200, "request within budget should be admitted"); + assert_eq!(parse_body(&raw), PLAIN_TEXT_BODY, "body should pass through unchanged"); +} + +#[test] +fn rejects_with_429_and_retry_after_when_estimate_budget_exhausted() { + let backend = Backend::fixed(PLAIN_TEXT_BODY) + .header("content-type", "text/plain") + .start_with_shutdown(); + let proxy_port = free_port(); + // capacity=50, estimate=40, no reconciliation (this backend returns no + // usage info): first request -40 -> 10 left (200), second request + // needs 40 more and must be rejected. + let config = token_rate_limit_config(proxy_port, backend.port(), 50, 40); + let proxy = start_proxy(&config); + + let first = http_send(proxy.addr(), &json_post("/v1/chat/completions", "{}")); + assert_eq!(parse_status(&first), 200, "first request should be admitted"); + + let second = http_send(proxy.addr(), &json_post("/v1/chat/completions", "{}")); + assert_eq!( + parse_status(&second), + 429, + "second request should be rejected, only 10 of 50 tokens remain after the first request spent 40" + ); + assert!( + parse_header(&second, "retry-after").is_some(), + "429 should carry a Retry-After header" + ); + assert!( + parse_header(&second, "x-ratelimit-limit-tokens").is_some(), + "429 should carry the token-suffixed limit header" + ); + assert!( + parse_header(&second, "x-ratelimit-remaining-tokens").is_some(), + "429 should carry the token-suffixed remaining header" + ); +} + +// ----------------------------------------------------------------------------- +// Reconciliation against actual usage +// ----------------------------------------------------------------------------- + +#[test] +fn reconciliation_frees_budget_for_next_request_after_low_actual_usage() { + let backend = Backend::fixed(OPENAI_LOW_USAGE_JSON) + .header("content-type", "application/json") + .start_with_shutdown(); + let proxy_port = free_port(); + // capacity=50, estimate=40. Every admitted request reserves 40 and + // then, since this backend's actual usage (10) is far below the + // estimate, gets 30 released back on reconciliation — so the + // window settles into a steady drain of only 10 net tokens per + // request instead of monotonically losing the full 40. Starting + // from an empty window (0/50 used): + // first: reserve 40 (50->10), reconcile +30 -> 40 [200] + // second: reserve 40 (40->0), reconcile +30 -> 30 [200] + // third: needs 40, only 30 remain [429] + // A naive (non-reconciling) reservation scheme would already have + // rejected "second" (it would see only 10 remaining after "first"). + let config = token_rate_limit_config(proxy_port, backend.port(), 50, 40); + let proxy = start_proxy(&config); + + let first = http_send(proxy.addr(), &json_post("/v1/chat/completions", "{}")); + assert_eq!(parse_status(&first), 200, "first request should be admitted"); + assert_eq!( + parse_body(&first), + OPENAI_LOW_USAGE_JSON, + "body should pass through unchanged" + ); + + let second = http_send(proxy.addr(), &json_post("/v1/chat/completions", "{}")); + assert_eq!( + parse_status(&second), + 200, + "reconciliation should have released enough budget (50 remaining) to admit a second 40-token request" + ); + + let third = http_send(proxy.addr(), &json_post("/v1/chat/completions", "{}")); + assert_eq!( + parse_status(&third), + 429, + "after two admissions the bucket should be down to 10 remaining, rejecting a third 40-token request" + ); +} + +// ----------------------------------------------------------------------------- +// Example config smoke test +// ----------------------------------------------------------------------------- + +#[test] +fn example_config_token_rate_limit() { + let backend = Backend::fixed(PLAIN_TEXT_BODY) + .header("content-type", "text/plain") + .start_with_shutdown(); + let proxy_port = free_port(); + + let config = load_example_config( + "token-rate-limit.yaml", + proxy_port, + HashMap::from([("127.0.0.1:3000", backend.port())]), + ); + let proxy = start_proxy(&config); + + let raw = http_send(proxy.addr(), &json_post("/v1/chat/completions", "{}")); + assert_eq!(parse_status(&raw), 200, "example config smoke test should return 200"); + assert_eq!(parse_body(&raw), PLAIN_TEXT_BODY, "body should pass through unchanged"); +} + +/// Smoke-tests the real `token-rate-limit-mixed-algorithms.yaml` example +/// file itself (ai#789/praxis#551), distinct from +/// `mixed_algorithm_rules_valkey_backend_isolates_budgets_across_gateway_replicas` +/// below, which builds its own hand-rolled YAML with a Valkey backend +/// clause to prove cross-replica isolation. This confirms the example +/// file that ships in the repo actually wires up both the `team-alpha` +/// (sliding_window) and `team-beta` (token_bucket) rules correctly. +#[test] +fn example_config_token_rate_limit_mixed_algorithms() { + let backend = Backend::fixed(PLAIN_TEXT_BODY) + .header("content-type", "text/plain") + .start_with_shutdown(); + let proxy_port = free_port(); + + let config = load_example_config( + "token-rate-limit-mixed-algorithms.yaml", + proxy_port, + HashMap::from([("127.0.0.1:3000", backend.port())]), + ); + let proxy = start_proxy(&config); + + let alpha = http_send( + proxy.addr(), + &json_post_with_headers("/v1/chat/completions", "{}", &[("x-app-id", "alpha")]), + ); + assert_eq!( + parse_status(&alpha), + 200, + "team-alpha's sliding_window rule should admit a request within its budget" + ); + + let beta = http_send( + proxy.addr(), + &json_post_with_headers("/v1/chat/completions", "{}", &[("x-app-id", "beta")]), + ); + assert_eq!( + parse_status(&beta), + 200, + "team-beta's token_bucket rule should admit a request within its budget" + ); +} + +// ----------------------------------------------------------------------------- +// Mixed algorithms, per rule (ai#789/praxis#551) -- Valkey-backed, driven +// through the real gateway pipeline across two independent proxy +// instances (simulated replicas), gated on a live Valkey/Redis instance +// via TOKEN_RATE_LIMIT_VALKEY_URL (see filters/src/token_rate_limit/ +// tests.rs for local setup instructions). +// ----------------------------------------------------------------------------- + +/// Two `token_rate_limit` rules sharing one Valkey namespace: `team-alpha` +/// (matched on `x-app-id: alpha`) enforces a sliding-window budget, +/// `team-beta` (matched on `x-app-id: beta`) enforces a token-bucket +/// budget. A request matching neither rule (no `x-app-id` header) is a +/// catch-all-free config here, so it passes through unrated -- this +/// config is deliberately about proving per-algorithm isolation, not +/// fallback-bucket behavior (already covered above). +fn mixed_algorithm_rules_config(proxy_port: u16, backend_port: u16, valkey_url: &str, namespace: &str) -> String { + let yaml = format!( + "listeners:\n\ + \x20 - name: default\n\ + \x20 address: \"0.0.0.0:8080\"\n\ + \x20 filter_chains:\n\ + \x20 - main\n\ + filter_chains:\n\ + \x20 - name: main\n\ + \x20 filters:\n\ + \x20 - filter: router\n\ + \x20 routes:\n\ + \x20 - path_prefix: \"/\"\n\ + \x20 cluster: backend\n\ + \x20 - filter: token_rate_limit\n\ + \x20 backend:\n\ + \x20 kind: valkey\n\ + \x20 url: {valkey_url}\n\ + \x20 namespace: {namespace}\n\ + \x20 rules:\n\ + \x20 - name: team-alpha\n\ + \x20 match:\n\ + \x20 headers:\n\ + \x20 x-app-id: alpha\n\ + \x20 algorithm: sliding_window\n\ + \x20 window: 1h\n\ + \x20 capacity: 100\n\ + \x20 reserved_tokens: 100\n\ + \x20 - name: team-beta\n\ + \x20 match:\n\ + \x20 headers:\n\ + \x20 x-app-id: beta\n\ + \x20 algorithm: token_bucket\n\ + \x20 capacity: 100\n\ + \x20 refill_rate: 0.001\n\ + \x20 reserved_tokens: 100\n\ + \x20 - filter: access_log\n\ + \x20 - filter: load_balancer\n\ + \x20 clusters:\n\ + \x20 - name: backend\n\ + \x20 endpoints:\n\ + \x20 - \"127.0.0.1:3000\"\n" + ); + patch_yaml(&yaml, proxy_port, &HashMap::from([("127.0.0.1:3000", backend_port)])) +} + +/// Proves both algorithms get the distributed-state property that's the +/// whole point of the Valkey backend -- not just in isolation (already +/// covered at the unit-test tier in `filters/src/token_rate_limit/ +/// tests.rs`), but through the real gateway pipeline, with two +/// independent proxy processes standing in for two gateway +/// replicas/instances sharing one Valkey namespace behind a load +/// balancer. +#[test] +fn mixed_algorithm_rules_valkey_backend_isolates_budgets_across_gateway_replicas() { + let Ok(valkey_url) = std::env::var("TOKEN_RATE_LIMIT_VALKEY_URL") else { + eprintln!("skipping: TOKEN_RATE_LIMIT_VALKEY_URL not set"); + return; + }; + let namespace = format!("praxis-it-mixed-algorithms-{}", std::process::id()); + + let backend_one = Backend::fixed(PLAIN_TEXT_BODY) + .header("content-type", "text/plain") + .start_with_shutdown(); + let backend_two = Backend::fixed(PLAIN_TEXT_BODY) + .header("content-type", "text/plain") + .start_with_shutdown(); + + // Two independent proxy processes, each built from the same rules + // config and pointed at the same Valkey namespace -- exactly as two + // gateway replicas behind a load balancer would be. + let proxy_one_port = free_port(); + let config_one = praxis_core::config::Config::from_yaml(&mixed_algorithm_rules_config( + proxy_one_port, + backend_one.port(), + &valkey_url, + &namespace, + )) + .expect("config should parse"); + let proxy_one = start_proxy(&config_one); + + let proxy_two_port = free_port(); + let config_two = praxis_core::config::Config::from_yaml(&mixed_algorithm_rules_config( + proxy_two_port, + backend_two.port(), + &valkey_url, + &namespace, + )) + .expect("config should parse"); + let proxy_two = start_proxy(&config_two); + + // team-alpha's sliding-window rule: admitted on replica one, + // capacity-exhausted (100/100) on replica two via the shared Valkey + // budget. + let alpha_first = http_send( + proxy_one.addr(), + &json_post_with_headers("/v1/chat/completions", "{}", &[("x-app-id", "alpha")]), + ); + assert_eq!( + parse_status(&alpha_first), + 200, + "team-alpha's first request should be admitted on replica one" + ); + let alpha_second = http_send( + proxy_two.addr(), + &json_post_with_headers("/v1/chat/completions", "{}", &[("x-app-id", "alpha")]), + ); + assert_eq!( + parse_status(&alpha_second), + 429, + "team-alpha's exhausted sliding-window budget must be visible on replica two via shared Valkey state" + ); + + // team-beta's token-bucket rule: admitted on replica one, + // capacity-exhausted on replica two -- proving the *second* + // algorithm gets the same cross-replica property, and that it + // doesn't share state with (or get blocked by) team-alpha's budget. + let beta_first = http_send( + proxy_one.addr(), + &json_post_with_headers("/v1/chat/completions", "{}", &[("x-app-id", "beta")]), + ); + assert_eq!( + parse_status(&beta_first), + 200, + "team-beta's first request should be admitted on replica one, unaffected by team-alpha's exhaustion" + ); + let beta_second = http_send( + proxy_two.addr(), + &json_post_with_headers("/v1/chat/completions", "{}", &[("x-app-id", "beta")]), + ); + assert_eq!( + parse_status(&beta_second), + 429, + "team-beta's exhausted token-bucket budget must be visible on replica two via shared Valkey state" + ); +} From 0eb7651c2522b615988d519a081b363428e4f11d Mon Sep 17 00:00:00 2001 From: Jordi Gil Date: Fri, 28 Aug 2026 21:46:04 -0400 Subject: [PATCH 3/5] ci(token_rate_limit): add Valkey-backed workflow and Makefile targets The Valkey-backed tests were silently skipping in CI because nothing set TOKEN_RATE_LIMIT_VALKEY_URL. Adds a dedicated workflow (path-scoped so it only runs when relevant files change) plus Makefile targets that spin up Valkey and run the gated test suite against it. Signed-off-by: Jordi Gil --- .github/workflows/valkey.yaml | 155 ++++++++++++++++++++++++++++++++++ Makefile | 18 +++- 2 files changed, 170 insertions(+), 3 deletions(-) create mode 100644 .github/workflows/valkey.yaml diff --git a/.github/workflows/valkey.yaml b/.github/workflows/valkey.yaml new file mode 100644 index 0000000000..4df4337615 --- /dev/null +++ b/.github/workflows/valkey.yaml @@ -0,0 +1,155 @@ +name: Tests (Token Rate Limit Valkey Backend) + +# ------------------------------------------------------------------------------ +# Workflow Settings +# ------------------------------------------------------------------------------ +# +# Every `token_rate_limit` test that needs a live Valkey/Redis (cross-instance +# shared state, background-worker reconciliation, EVAL fault injection) is +# gated on TOKEN_RATE_LIMIT_VALKEY_URL: unset, it early-returns rather than +# failing, so `make test-unit`/`make test-integration` stay green on a +# contributor's laptop with no Valkey running. That gate was never wired to +# an actual Valkey anywhere in CI either, so it silently took the same +# early-return path there too -- this workflow is that wiring, split out +# (like postgres.yaml) so the common fast test/integration jobs don't all +# pay for a service container on every unrelated change. + +on: + push: + branches: [main] + paths: + - "filters/src/token_rate_limit/**" + - "tests/integration/tests/suite/examples/token_rate_limit.rs" + - "Cargo.toml" + - "Cargo.lock" + - "Makefile" + - ".github/workflows/valkey.yaml" + pull_request: + branches: [main] + types: [opened, synchronize, reopened] + paths: + - "filters/src/token_rate_limit/**" + - "tests/integration/tests/suite/examples/token_rate_limit.rs" + - "Cargo.toml" + - "Cargo.lock" + - "Makefile" + - ".github/workflows/valkey.yaml" + merge_group: + branches: [main] + workflow_dispatch: + inputs: + debug: + description: "Enable verbose test output (V=1)" + required: false + type: boolean + default: false + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +permissions: {} + +env: + CARGO_TERM_COLOR: always + V: ${{ inputs.debug && '1' || '' }} + +jobs: + # ---------------------------------------------------------------------------- + # Path filter for merge_group (which does not support on.paths) + # ---------------------------------------------------------------------------- + + changes: + if: github.event_name == 'merge_group' + runs-on: ubuntu-24.04 + permissions: + contents: read + outputs: + relevant: ${{ steps.filter.outputs.relevant }} + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + fetch-depth: 0 + + - name: Detect relevant path changes + id: filter + run: | + CHANGED=$(git diff --name-only "${{ github.event.merge_group.base_sha }}" \ + "${{ github.event.merge_group.head_sha }}" -- \ + 'filters/src/token_rate_limit/' \ + 'tests/integration/tests/suite/examples/token_rate_limit.rs' \ + 'Cargo.toml' \ + 'Cargo.lock' \ + 'Makefile' \ + '.github/workflows/valkey.yaml') + if [ -n "$CHANGED" ]; then + echo "relevant=true" >> "$GITHUB_OUTPUT" + else + echo "relevant=false" >> "$GITHUB_OUTPUT" + fi + + # ---------------------------------------------------------------------------- + # token_rate_limit Valkey unit tests (praxis-ai-filters) + # ---------------------------------------------------------------------------- + + unit: + needs: [changes] + if: | + always() && + (needs.changes.result == 'skipped' || needs.changes.outputs.relevant == 'true') + runs-on: ubuntu-24.04 + permissions: + contents: read + services: + valkey: + image: docker.io/valkey/valkey:8-alpine + ports: + - 6379:6379 + options: >- + --health-cmd "valkey-cli ping" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + TOKEN_RATE_LIMIT_VALKEY_URL: redis://127.0.0.1:6379 + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Setup Rust + uses: praxis-proxy/conventions/.github/actions/setup-rust@7e1e8d97c2dc820d24b31f9a65b119c4d0e5342c # v0.1.0 + + - name: token_rate_limit Valkey unit tests + run: make test-token-rate-limit-valkey-unit + + # ---------------------------------------------------------------------------- + # token_rate_limit Valkey integration test (cross-instance shared state) + # ---------------------------------------------------------------------------- + + integration: + needs: [changes] + if: | + always() && + (needs.changes.result == 'skipped' || needs.changes.outputs.relevant == 'true') + runs-on: ubuntu-24.04 + permissions: + contents: read + services: + valkey: + image: docker.io/valkey/valkey:8-alpine + ports: + - 6379:6379 + options: >- + --health-cmd "valkey-cli ping" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + TOKEN_RATE_LIMIT_VALKEY_URL: redis://127.0.0.1:6379 + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Setup Rust + uses: praxis-proxy/conventions/.github/actions/setup-rust@7e1e8d97c2dc820d24b31f9a65b119c4d0e5342c # v0.1.0 + + - name: token_rate_limit Valkey integration test + run: make test-token-rate-limit-valkey-integration diff --git a/Makefile b/Makefile index 5051000076..ab5a987197 100644 --- a/Makefile +++ b/Makefile @@ -11,7 +11,7 @@ V ?= # Experimental filter features are off by default in builds, so lint and # test explicitly enable them — otherwise the gated filter code is never # compiled, linted, or tested by CI. -EXPERIMENTAL_FEATURES := azure-ad-filter,gcp-adc-filter,http-callout-filter +EXPERIMENTAL_FEATURES := azure-ad-filter,gcp-adc-filter,http-callout-filter,token-rate-limit-filter ifneq ($(V),) _NOCAPTURE := -- --nocapture @@ -20,6 +20,7 @@ endif .PHONY: all build release check clean \ test test-unit test-schema test-integration test-inference-fixtures \ test-postgres-unit test-postgres-integration test-environment \ + test-token-rate-limit-valkey-unit test-token-rate-limit-valkey-integration \ openai-conformance check-openai-conformance-reference test-openai-conformance \ lint fmt doc audit coverage-check \ require-container-engine \ @@ -84,7 +85,8 @@ test-schema: test-integration: cargo test -p praxis-tests-integration $(_NOCAPTURE) cargo test -p praxis-tests-integration --features $(EXPERIMENTAL_FEATURES) --test suite \ - -- examples::azure_ad examples::gcp_adc examples::lakera_guard $(if $(V),--nocapture) + -- examples::azure_ad examples::gcp_adc examples::lakera_guard examples::token_rate_limit \ + $(if $(V),--nocapture) test-inference-fixtures: cargo test -p praxis-test-utils $(_NOCAPTURE) @@ -97,6 +99,13 @@ test-postgres-unit: test-postgres-integration: cargo test -p praxis-tests-integration --test suite openai_response_store_postgres -- --ignored $(_NOCAPTURE) +test-token-rate-limit-valkey-unit: + cargo test -p praxis-ai-filters --features token-rate-limit-filter valkey $(_NOCAPTURE) + +test-token-rate-limit-valkey-integration: + cargo test -p praxis-tests-integration --features token-rate-limit-filter --test suite \ + mixed_algorithm_rules_valkey_backend_isolates_budgets_across_gateway_replicas $(_NOCAPTURE) + openai-conformance: cargo xtask openai-conformance $(OPENAI_CONFORMANCE_ARGS) @@ -105,6 +114,7 @@ check-openai-conformance-reference: test-openai-conformance: openai-conformance + test-environment: cargo test -p praxis-ai-llmd-ext-proc $(_NOCAPTURE) cargo test -p praxis-tests-integration --features llmd-ext-proc llmd_ext_proc $(_NOCAPTURE) @@ -117,7 +127,7 @@ test-environment: lint: cargo clippy --workspace --all-targets -- -D warnings cargo clippy --workspace --all-targets \ - --features praxis-ai-proxy/azure-ad-filter,praxis-ai-proxy/gcp-adc-filter,praxis-ai-proxy/http-callout-filter,praxis-tests-integration/azure-ad-filter,praxis-tests-integration/gcp-adc-filter,praxis-tests-integration/http-callout-filter \ + --features praxis-ai-proxy/azure-ad-filter,praxis-ai-proxy/gcp-adc-filter,praxis-ai-proxy/http-callout-filter,praxis-ai-proxy/token-rate-limit-filter,praxis-tests-integration/azure-ad-filter,praxis-tests-integration/gcp-adc-filter,praxis-tests-integration/http-callout-filter,praxis-tests-integration/token-rate-limit-filter \ -- -D warnings cargo +nightly fmt --all -- --check cargo machete --with-metadata . @@ -215,6 +225,8 @@ help: @echo " test-inference-fixtures inference fixture and replay tests" @echo " test-postgres-unit postgres store unit tests (needs DATABASE_URL)" @echo " test-postgres-integration postgres store integration tests (needs container engine)" + @echo " test-token-rate-limit-valkey-unit token_rate_limit Valkey unit tests (needs TOKEN_RATE_LIMIT_VALKEY_URL)" + @echo " test-token-rate-limit-valkey-integration token_rate_limit Valkey integration test (needs TOKEN_RATE_LIMIT_VALKEY_URL)" @echo " test-environment llm-d ext_proc environment tests" @echo " openai-conformance compare registered API areas with OpenAI's OpenAPI spec" @echo " check-openai-conformance-reference verify the pinned complete OpenAI reference" From b4e798df5b42abfcb00a31ef3768f139d7698b3b Mon Sep 17 00:00:00 2001 From: Jordi Gil Date: Fri, 28 Aug 2026 21:46:11 -0400 Subject: [PATCH 4/5] fix(xtask): render tag field for flattened, internally-tagged enums generate-filter-docs dropped the discriminator field for any #[serde(flatten)]'d struct whose type was an internally-tagged enum (#[serde(tag = "...")]) -- e.g. rules[].algorithm on the new token_rate_limit filter never made it into the generated table. Synthesizes the missing row instead of silently omitting it, with a regression test. Signed-off-by: Jordi Gil --- xtask/src/filter_docs.rs | 126 +++++++++++++++++++++++++++++++++++---- 1 file changed, 116 insertions(+), 10 deletions(-) diff --git a/xtask/src/filter_docs.rs b/xtask/src/filter_docs.rs index 96df890aa1..07842ae587 100644 --- a/xtask/src/filter_docs.rs +++ b/xtask/src/filter_docs.rs @@ -239,6 +239,9 @@ struct EnumInfo { variants: Vec, /// Whether serde tries variants by shape instead of by variant tag. untagged: bool, + /// `#[serde(tag = "...")]` discriminator key name, for internally + /// tagged enums. `None` for untagged/externally tagged enums. + tag: Option, /// Source shape for each variant. variant_shapes: Vec, /// Named fields from struct-like variants. @@ -858,6 +861,8 @@ fn append_rendered_fields( doc: field.doc.clone(), required: required_kind(field), }); + } else if let Some(tag_field) = flattened_tag_field(prefix, field, items) { + out.push(tag_field); } let nested_prefix = if field.flatten { @@ -897,6 +902,26 @@ fn append_nested_fields( } } +/// Synthesize the discriminator field row for a `#[serde(flatten)]` +/// field whose type is an internally tagged enum (`#[serde(tag = +/// "...")]`). The tag itself (e.g. `algorithm`) is a real, required +/// YAML key, but has no corresponding field on any of the enum's +/// variants for [`append_rendered_fields`] to otherwise pick up -- +/// without this, it's silently missing from the generated table. +/// Returns `None` for untagged/externally tagged/non-enum flattened +/// fields, which have no single discriminator key to document here. +fn flattened_tag_field(prefix: &str, field: &RawField, items: &ModuleItems) -> Option { + let type_name = nested_type_name(&field.ty)?; + let info = items.enums.get(&type_name)?; + let tag = info.tag.as_ref()?; + Some(FieldInfo { + name: field_path(prefix, tag), + type_str: render_enum_type(info, &items.enums), + doc: field.doc.clone(), + required: required_kind(field), + }) +} + /// Return the rendered requirement kind for a raw field. fn required_kind(field: &RawField) -> RequiredKind { if matches!(field.requirement_hint, RequirementHint::OneOf) { @@ -1031,11 +1056,20 @@ fn serde_field_name(field: &syn::Field, rename_all: Option<&str>) -> String { field .ident .as_ref() - .map(|ident| apply_rename(&ident.to_string(), rename_all)) + .map(|ident| apply_rename(strip_raw_ident_prefix(&ident.to_string()), rename_all)) }) .unwrap_or_default() } +/// Strip a raw identifier's `r#` prefix (e.g. `r#match` -> `match`), which +/// `syn::Ident::to_string()` preserves but serde's own field-name +/// resolution does not -- a raw identifier is only needed to use a +/// reserved keyword as a Rust binding, it has no effect on the +/// serialized/deserialized field name. +fn strip_raw_ident_prefix(ident: &str) -> &str { + ident.strip_prefix("r#").unwrap_or(ident) +} + /// Return whether a serde attribute contains a given nested key. fn serde_attr_contains(attr: &syn::Attribute, name: &str) -> bool { if !attr.path().is_ident("serde") { @@ -1087,6 +1121,7 @@ fn serde_lit_value(attr: &syn::Attribute, name: &str) -> Option { fn extract_enum_info(e: &syn::ItemEnum) -> EnumInfo { let rename_all = detect_rename_all(&e.attrs); let untagged = has_serde_attr(&e.attrs, "untagged"); + let tag = e.attrs.iter().find_map(|attr| serde_lit_value(attr, "tag")); let variants = e .variants .iter() @@ -1098,7 +1133,23 @@ fn extract_enum_info(e: &syn::ItemEnum) -> EnumInfo { }) .collect(); let variant_shapes = e.variants.iter().map(enum_variant_shape).collect(); - let mut variant_fields: Vec> = e.variants.iter().map(parse_variant_fields).collect(); + let variant_fields: Vec> = e.variants.iter().map(parse_variant_fields).collect(); + let fields = flatten_variant_fields_marking_one_of(variant_fields); + + EnumInfo { + variants, + untagged, + tag, + variant_shapes, + fields, + } +} + +/// Flatten every variant's fields into one list, marking each as +/// [`RequirementHint::OneOf`] when more than one variant has its own +/// named fields (a struct-like tagged/untagged enum), since only one +/// variant's fields are ever present in a given YAML document. +fn flatten_variant_fields_marking_one_of(mut variant_fields: Vec>) -> Vec { let named_variant_count = variant_fields.iter().filter(|fields| !fields.is_empty()).count(); if named_variant_count > 1 { for fields in &mut variant_fields { @@ -1107,14 +1158,7 @@ fn extract_enum_info(e: &syn::ItemEnum) -> EnumInfo { } } } - let fields = variant_fields.into_iter().flatten().collect(); - - EnumInfo { - variants, - untagged, - variant_shapes, - fields, - } + variant_fields.into_iter().flatten().collect() } /// Return the source shape for an enum variant. @@ -1967,6 +2011,33 @@ mod tests { assert_eq!(capitalize(""), "", "empty string"); } + #[test] + fn strip_raw_ident_prefix_removes_the_r_hash_prefix() { + assert_eq!(strip_raw_ident_prefix("r#match"), "match"); + assert_eq!(strip_raw_ident_prefix("r#type"), "type"); + assert_eq!( + strip_raw_ident_prefix("estimate_tokens"), + "estimate_tokens", + "non-raw idents pass through" + ); + } + + #[test] + fn parse_config_fields_renders_a_raw_ident_field_without_its_r_hash_prefix() { + let file: syn::File = syn::parse_str( + "#[derive(Deserialize)]\n#[serde(deny_unknown_fields)]\nstruct Cfg { r#match: Option }", + ) + .unwrap(); + let syn::Item::Struct(s) = &file.items[0] else { + panic!("expected a struct item"); + }; + let fields = parse_config_fields(s).unwrap(); + assert_eq!( + fields[0].name, "match", + "the YAML field name must be the bare keyword, not the Rust raw-identifier spelling" + ); + } + #[test] fn first_paragraph_extracts_before_blank_line() { let doc = "First line.\nSecond line.\n\nSecond paragraph."; @@ -2341,6 +2412,41 @@ mod tests { ); } + #[test] + fn flattened_internally_tagged_enum_synthesizes_a_tag_field() { + let source = " + #[derive(Debug, Deserialize)] + #[serde(deny_unknown_fields)] + struct OuterConfig { rules: Vec } + #[derive(Debug, Deserialize)] + struct RuleConfig { + /// Which admission algorithm this rule enforces. + #[serde(flatten)] + algorithm: RuleAlgorithm, + } + #[derive(Debug, Deserialize)] + #[serde(tag = \"algorithm\", rename_all = \"snake_case\")] + enum RuleAlgorithm { SlidingWindow { window: String }, TokenBucket { refill_rate: f64 } } + "; + let file: syn::File = syn::parse_str(source).unwrap(); + let mut items = ModuleItems::new(); + parse_file_items(&file, &mut items); + let filter = build_filter(&items, "test", Some("OuterConfig")); + + let tag_field = filter + .fields + .iter() + .find(|field| field.name == "rules[].algorithm") + .expect("the tag discriminator itself must render as a documented field"); + assert_eq!(tag_field.type_str, "`sliding_window` \\| `token_bucket`"); + assert_eq!(tag_field.required, RequiredKind::Yes); + assert_eq!(tag_field.doc, "Which admission algorithm this rule enforces."); + assert!( + filter.fields.iter().any(|field| field.name == "rules[].window"), + "tagged variants' own fields must still render alongside the tag" + ); + } + #[test] fn module_level_yaml_examples_are_included() { let source = " From 28ce5c3d8a78a7d445dd4517dc9e2987a712634b Mon Sep 17 00:00:00 2001 From: Jordi Gil Date: Fri, 28 Aug 2026 21:46:17 -0400 Subject: [PATCH 5/5] docs(token_rate_limit): generate filter reference documentation cargo xtask generate-filter-docs output for the new filter, including the algorithm field now that the generator handles flattened, internally-tagged enums correctly. Signed-off-by: Jordi Gil --- docs/filters/reference.md | 6 ++++ docs/filters/token_rate_limit.md | 61 ++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+) create mode 100644 docs/filters/token_rate_limit.md diff --git a/docs/filters/reference.md b/docs/filters/reference.md index 988669fdab..8a95a5382b 100644 --- a/docs/filters/reference.md +++ b/docs/filters/reference.md @@ -110,6 +110,12 @@ see the [Praxis core filter reference][core-ref]. |--------|-------------| | [`time_to_first_token`](time_to_first_token.md) | Measures time-to-first-token for streaming AI responses. | +### Token Rate Limit + +| Filter | Description | +|--------|-------------| +| [`token_rate_limit`](token_rate_limit.md) | Token-denominated rate limiter: reserves an estimated cost at admission, reconciles against actual usage after the response completes. Evaluates an ordered list of rules, each with its own optional match condition, algorithm choice, and budget. | + ### Token Usage | Filter | Description | diff --git a/docs/filters/token_rate_limit.md b/docs/filters/token_rate_limit.md new file mode 100644 index 0000000000..58b69bbba5 --- /dev/null +++ b/docs/filters/token_rate_limit.md @@ -0,0 +1,61 @@ + + + +# `token_rate_limit` + +Token-denominated rate limiter: reserves an estimated cost at admission, reconciles against actual usage after the response completes. Evaluates an ordered list of rules, each with its own optional match condition, algorithm choice, and budget. + +## Configuration Notes + +Experimental: requires the `token-rate-limit-filter` cargo feature, which is off by default and activates the `experimental` marker. This filter delivers the agreed M1/M2/M6 milestone scope, but its parent proposal is not yet `accepted` and open questions remain (HA/clustered-Valkey failure modes, and the relationship to Kuadrant's `TokenRateLimitPolicy` -- see `ai#127`). The configuration surface may change between releases. + +Mirrors the `rules:`/`match:` shape from the `00121_token-rate-limiting` proposal in `praxis-proxy/enhancements`, scoped to this milestone's static header-value matchers and per-rule algorithm choice. CEL matchers, soft-limit tiers, weighted per-type accounting, and configurable estimation strategies are still out of scope (see the module doc comment) -- upstream itself defers the first two; the latter two are deferred to a separate follow-up by design, not by upstream mandate. + +Assumes request identity has already been resolved upstream (this filter doesn't authenticate callers) -- a catch-all rule (no `match:`) reserves quota for every request that reaches it, including probes and health checks. Scope rules with explicit `match:` conditions, or place an identity/auth filter earlier in the pipeline. Tracked as follow-on integration work in `grid#101`. + +## Configuration + +| Field | Type | Required | Description | +|-------|------|---------|-------------| +| `rules` | RuleConfig[] | yes | Evaluated in order; the first rule whose `match` is satisfied (or which has no `match` at all) applies to a given request. A request satisfying no rule's `match` is not rate limited by this filter instance -- add a trailing rule with no `match` to enforce a catch-all budget instead. | +| `rules[].name` | string | yes | Human-readable rule identifier, folded into Valkey key namespacing so distinct rules sharing one backend never collide. Renaming a live `valkey`-backed rule is therefore not a no-op for operators: it changes the Valkey key hash, so the old name's tracked budget is orphaned (left to expire on its own TTL) and the new name starts with a fresh budget. There's no migration/rename path today -- routine config hygiene (e.g. renaming `"gold"` to `"gold-tier"`) silently resets that rule's state. | +| `rules[].match` | MatchConfig | no | Static header-value match condition. Every listed header must be present on the request with an exact value match (`ANDed`) for this rule to apply. Omit entirely for a catch-all rule. | +| `rules[].match.headers` | object | yes | Every header must be present on the request with this exact value for the rule to match (`ANDed` across all entries). | +| `rules[].algorithm` | `sliding_window` \| `token_bucket` | yes | Which admission algorithm this rule enforces, and that algorithm's own parameters. | +| `rules[].window` | string | one of | Sliding window duration (e.g. `"1h"`, `"60s"`). | +| `rules[].capacity` | integer | one of | Maximum tokens admitted within `window`. | +| `rules[].capacity` | integer | one of | Maximum tokens held at once (the bucket's ceiling). | +| `rules[].refill_rate` | number | one of | Tokens refilled per second, up to `capacity`. | +| `rules[].reserved_tokens` | integer | yes | Fixed token cost reserved at admission time, before actual usage is known. Placeholder pending M3 (configurable estimation strategies). Real deployments will want this derived from request metadata (e.g. `max_tokens`) rather than a single fixed constant -- that's out of scope for this milestone. | +| `rules[].reservation_timeout` | string | no | How long an admitted-but-never-reconciled reservation (lost request: timeout, connection reset, upstream crash) is tracked as active before that already-reserved-at-admission charge against its estimate becomes irreversibly locked in (sliding-window: folded into the settled total so it survives the window's normal aging-out; token-bucket: the tokens were already decremented at reserve time regardless, this only bounds how long the reservation is tracked as pending). This does **not** defer when the charge first applies -- it applies immediately at admission, same as any other reservation. Answers the proposal's still-open "lost request handling" question for this milestone. Defaults to [`DEFAULT_RESERVATION_TIMEOUT`] when unset. | +| `backend` | BackendConfig | no | Where every rule's admission state lives: in-process (default, one budget per gateway instance) or a shared Valkey backend (one budget shared across every gateway instance/replica). One backend for the whole filter, not per rule -- rules already share Valkey key-space isolation via `namespace`/rule-name hashing, so per-rule backend selection bought no isolation benefit, only a separate Valkey connection per rule pointed at the same URL. Revisit if a real deployment ever needs to mix in-process and Valkey rules in one filter instance. | +| `backend.kind` | `memory` \| `valkey` | no | Which backend implementation to use. | +| `backend.url` | string | no | Backend connection URL. Supports one `${ENV_VAR}` reference, so credentials/hostnames don't need to be committed to config. Required when `kind: valkey`, ignored otherwise. | +| `backend.namespace` | string | no | Key namespace prefix, so multiple filter rules or deployments can share one Valkey instance without colliding. Ignored for `kind: memory`. Defaults to `"praxis:token_rate_limit"` when unset. | + +## Example + +```yaml +filter: token_rate_limit +backend: # optional: defaults to in-process state, shared by every rule + kind: valkey # memory (default) | valkey + url: "${TOKEN_RATE_LIMIT_VALKEY_URL}" + namespace: praxis:token_rate_limit +rules: + - name: team-alpha # human-readable, unique per filter instance + match: # optional: omit for a catch-all rule + headers: + x-app-id: alpha + algorithm: sliding_window # sliding_window | token_bucket + window: 1h # sliding_window only: window duration + capacity: 100000 # max tokens admitted (sliding_window) or held (token_bucket) + reserved_tokens: 500 # fixed cost reserved per request at admission + - name: team-beta + match: + headers: + x-app-id: beta + algorithm: token_bucket + capacity: 50000 + refill_rate: 50 # token_bucket only: tokens refilled per second + reserved_tokens: 200 +```