feat(rate_limiter) Adding tenant based rate limiting based on input tokens.

Signed-off-by: Varun Shenoy <varun.vinayak.shenoy@oracle.com>

Author: shenoyvvarun

model_gateway/src/app_context.rs

@@ -22,6 +22,7 @@ use crate::{
     middleware::TokenBucket,
     observability::inflight_tracker::InFlightRequestTracker,
     policies::PolicyRegistry,
+    rate_limit::LocalTokenRateLimiter,
     routers::{
         grpc::multimodal::MultimodalConfigRegistry, openai::realtime::RealtimeRegistry,
         router_manager::RouterManager,
@@ -54,6 +55,7 @@ pub struct AppContext {
     pub client: Client,
     pub router_config: RouterConfig,
     pub rate_limiter: Option<Arc<TokenBucket>>,
+    pub token_rate_limiter: Option<Arc<LocalTokenRateLimiter>>,
     pub tokenizer_registry: Arc<TokenizerRegistry>,
     pub multimodal_config_registry: Arc<MultimodalConfigRegistry>,
     pub reasoning_parser_factory: Option<ReasoningParserFactory>,
@@ -97,6 +99,7 @@ pub struct AppContextBuilder {
     client: Option<Client>,
     router_config: Option<RouterConfig>,
     rate_limiter: Option<Arc<TokenBucket>>,
+    token_rate_limiter: Option<Arc<LocalTokenRateLimiter>>,
     tokenizer_registry: Option<Arc<TokenizerRegistry>>,
     reasoning_parser_factory: Option<ReasoningParserFactory>,
     tool_parser_factory: Option<ToolParserFactory>,
@@ -152,6 +155,7 @@ impl AppContextBuilder {
             client: None,
             router_config: None,
             rate_limiter: None,
+            token_rate_limiter: None,
             tokenizer_registry: None,
             reasoning_parser_factory: None,
             tool_parser_factory: None,
@@ -362,6 +366,7 @@ impl AppContextBuilder {
                 .ok_or(AppContextBuildError::MissingField("client"))?,
             router_config,
             rate_limiter: self.rate_limiter,
+            token_rate_limiter: self.token_rate_limiter,
             tokenizer_registry: self
                 .tokenizer_registry
                 .ok_or(AppContextBuildError::MissingField("tokenizer_registry"))?,
@@ -520,6 +525,11 @@ impl AppContextBuilder {
                 )))
             }
         };
+        self.token_rate_limiter = config.multi_tenant_rate_limit.enabled.then(|| {
+            Arc::new(LocalTokenRateLimiter::new(
+                config.multi_tenant_rate_limit.clone(),
+            ))
+        });
         self
     }
 

model_gateway/src/config/builder.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{hash_map::Entry, HashMap};
 
 use smg_mcp::McpConfig;
 
@@ -580,6 +580,59 @@ impl RouterConfigBuilder {
         self
     }
 
+    pub fn multi_tenant_rate_limit_enabled(mut self, enabled: bool) -> Self {
+        self.config.multi_tenant_rate_limit.enabled = enabled;
+        self
+    }
+
+    pub fn default_tokens_per_minute(mut self, limit: u32) -> Self {
+        self.config
+            .multi_tenant_rate_limit
+            .default_tokens_per_minute = limit;
+        self
+    }
+
+    pub fn default_requests_per_minute(mut self, limit: u32) -> Self {
+        self.config
+            .multi_tenant_rate_limit
+            .default_requests_per_minute = limit;
+        self
+    }
+
+    pub fn tenant_rate_limit<S: Into<String>>(
+        mut self,
+        tenant_key: S,
+        tokens_per_minute: u32,
+        requests_per_minute: u32,
+    ) -> Self {
+        let tenant_key = tenant_key.into();
+        let new_policy = crate::rate_limit::TenantTokenPolicy {
+            tokens_per_minute,
+            requests_per_minute,
+        };
+
+        match self
+            .config
+            .multi_tenant_rate_limit
+            .tenants
+            .entry(tenant_key.clone())
+        {
+            Entry::Vacant(entry) => {
+                entry.insert(new_policy);
+            }
+            Entry::Occupied(mut entry) => {
+                tracing::warn!(
+                    tenant_key = %tenant_key,
+                    previous_policy = ?entry.get(),
+                    new_policy = ?new_policy,
+                    "overwriting duplicate tenant rate limit policy"
+                );
+                entry.insert(new_policy);
+            }
+        }
+        self
+    }
+
     pub fn maybe_model_path(mut self, path: Option<impl Into<String>>) -> Self {
         self.config.model_path = path.map(|p| p.into());
         self
@@ -930,6 +983,56 @@ mod tests {
         assert!(modified.trace_config.is_some());
     }
 
+    #[test]
+    fn test_builder_multi_tenant_rate_limit_round_trip() {
+        let config = RouterConfigBuilder::new()
+            .regular_mode(vec!["http://worker1:8000".to_string()])
+            .multi_tenant_rate_limit_enabled(true)
+            .default_tokens_per_minute(10_000)
+            .default_requests_per_minute(60)
+            .tenant_rate_limit("team-a", 50_000, 600)
+            .tenant_rate_limit("team-b", 100_000, 1_200)
+            .build()
+            .unwrap();
+
+        assert!(config.multi_tenant_rate_limit.enabled);
+        assert_eq!(
+            config.multi_tenant_rate_limit.default_tokens_per_minute,
+            10_000
+        );
+        assert_eq!(
+            config.multi_tenant_rate_limit.default_requests_per_minute,
+            60
+        );
+        let team_a = config
+            .multi_tenant_rate_limit
+            .tenants
+            .get("team-a")
+            .expect("team-a override registered");
+        assert_eq!(team_a.tokens_per_minute, 50_000);
+        assert_eq!(team_a.requests_per_minute, 600);
+        assert_eq!(config.multi_tenant_rate_limit.tenants.len(), 2);
+    }
+
+    #[test]
+    fn test_builder_duplicate_tenant_rate_limit_overwrites_latest_policy() {
+        let config = RouterConfigBuilder::new()
+            .regular_mode(vec!["http://worker1:8000".to_string()])
+            .tenant_rate_limit("team-a", 50_000, 600)
+            .tenant_rate_limit("team-a", 100_000, 1_200)
+            .build()
+            .unwrap();
+
+        let team_a = config
+            .multi_tenant_rate_limit
+            .tenants
+            .get("team-a")
+            .expect("team-a override registered");
+        assert_eq!(team_a.tokens_per_minute, 100_000);
+        assert_eq!(team_a.requests_per_minute, 1_200);
+        assert_eq!(config.multi_tenant_rate_limit.tenants.len(), 1);
+    }
+
     /// Test complex routing mode helper method
     #[test]
     fn test_builder_prefill_decode_mode() {

model_gateway/src/config/types.rs

@@ -110,6 +110,8 @@ pub struct RouterConfig {
     pub background: BackgroundConfig,
     #[serde(default)]
     pub tenant_resolution: TenantResolutionConfig,
+    #[serde(default)]
+    pub multi_tenant_rate_limit: crate::rate_limit::MultiTenantRateLimitConfig,
     /// Set to -1 to disable rate limiting
     pub max_concurrent_requests: i32,
     pub queue_size: usize,
@@ -646,6 +648,7 @@ impl Default for RouterConfig {
             memory_runtime: MemoryRuntimeConfig::default(),
             background: BackgroundConfig::default(),
             tenant_resolution: TenantResolutionConfig::default(),
+            multi_tenant_rate_limit: crate::rate_limit::MultiTenantRateLimitConfig::default(),
             max_concurrent_requests: -1,
             queue_size: 100,
             queue_timeout_secs: 60,

model_gateway/src/lib.rs

@@ -5,6 +5,7 @@ pub mod mesh;
 pub mod middleware;
 pub mod observability;
 pub mod policies;
+pub mod rate_limit;
 pub mod routers;
 pub mod server;
 pub mod service_discovery;

model_gateway/src/main.rs

@@ -355,6 +355,22 @@ struct CliArgs {
     #[arg(long, help_heading = "Rate Limiting")]
     rate_limit_tokens_per_second: Option<i32>,
 
+    /// Enable tenant-aware token rate limiting.
+    #[arg(long, default_value_t = false, help_heading = "Rate Limiting")]
+    multi_tenant_rate_limit_enabled: bool,
+
+    /// Default token budget per minute for tenants without an explicit override.
+    #[arg(long, default_value_t = 0, help_heading = "Rate Limiting")]
+    default_tokens_per_minute: u32,
+
+    /// Default request budget per minute for tenants without an explicit override.
+    #[arg(long, default_value_t = 0, help_heading = "Rate Limiting")]
+    default_requests_per_minute: u32,
+
+    /// Per-tenant override in the form tenant_key:tpm:rpm, e.g. header:team-a:1000:10
+    #[arg(long = "tenant-rate-limit", num_args = 0.., help_heading = "Rate Limiting")]
+    tenant_rate_limits: Vec<String>,
+
     // ==================== Retry Configuration ====================
     /// Maximum number of retry attempts
     #[arg(long, default_value_t = 5, help_heading = "Retry Configuration")]
@@ -1249,6 +1265,9 @@ impl CliArgs {
             .trust_tenant_header(self.trust_tenant_header)
             .tenant_header_name(&self.tenant_header_name)
             .maybe_rate_limit_tokens_per_second(self.rate_limit_tokens_per_second)
+            .multi_tenant_rate_limit_enabled(self.multi_tenant_rate_limit_enabled)
+            .default_tokens_per_minute(self.default_tokens_per_minute)
+            .default_requests_per_minute(self.default_requests_per_minute)
             .maybe_model_path(self.model_path.as_ref())
             .maybe_tokenizer_path(self.tokenizer_path.as_ref())
             .maybe_chat_template(self.chat_template.as_ref())
@@ -1269,6 +1288,23 @@ impl CliArgs {
             .dp_minimum_tokens_scheduler(self.dp_minimum_tokens_scheduler)
             .maybe_server_cert_and_key(self.tls_cert_path.as_ref(), self.tls_key_path.as_ref());
 
+        let mut builder = builder;
+        for spec in &self.tenant_rate_limits {
+            let mut parts = spec.rsplitn(3, ':');
+            let rpm = parts.next().and_then(|s| s.parse::<u32>().ok());
+            let tpm = parts.next().and_then(|s| s.parse::<u32>().ok());
+            let tenant_key = parts.next();
+            if let (Some(tenant_key), Some(tpm), Some(rpm)) = (tenant_key, tpm, rpm) {
+                builder = builder.tenant_rate_limit(tenant_key, tpm, rpm);
+            } else {
+                return Err(ConfigError::ValidationFailed {
+                    reason: format!(
+                        "invalid --tenant-rate-limit '{spec}'; expected tenant_key:tpm:rpm"
+                    ),
+                });
+            }
+        }
+
         builder.build()
     }
 

model_gateway/src/rate_limit/error.rs

@@ -0,0 +1,39 @@
+use axum::{
+    http::{self, header::RETRY_AFTER, HeaderValue},
+    response::Response,
+};
+
+use super::local::TERMINAL_REJECTION_RETRY_AFTER_SECS;
+use crate::routers::error::create_error;
+
+pub fn rate_limit_exceeded_response(retry_after_secs: u64) -> Response {
+    let mut response = create_error(
+        http::StatusCode::TOO_MANY_REQUESTS,
+        "tenant_rate_limit_exceeded",
+        "Tenant rate limit exceeded for this request",
+    );
+
+    if retry_after_secs == TERMINAL_REJECTION_RETRY_AFTER_SECS {
+        return response;
+    }
+
+    if let Ok(v) = HeaderValue::from_str(&retry_after_secs.max(1).to_string()) {
+        response.headers_mut().insert(RETRY_AFTER, v);
+    }
+
+    response
+}
+
+#[cfg(test)]
+mod tests {
+    use axum::http::header::RETRY_AFTER;
+
+    use super::{rate_limit_exceeded_response, TERMINAL_REJECTION_RETRY_AFTER_SECS};
+
+    #[test]
+    fn omits_retry_after_for_terminal_rate_limit_rejection() {
+        let response = rate_limit_exceeded_response(TERMINAL_REJECTION_RETRY_AFTER_SECS);
+
+        assert!(response.headers().get(RETRY_AFTER).is_none());
+    }
+}