Skip to main content

systemprompt_models/config/
rate_limits.rs

1//! Rate limit configuration.
2
3use crate::auth::RateLimitTier;
4use crate::profile::{RateLimitsConfig, TierMultipliers};
5
6#[derive(Debug, Clone, Copy)]
7pub struct RateLimitConfig {
8    pub oauth_public_per_second: u64,
9    pub oauth_auth_per_second: u64,
10    pub contexts_per_second: u64,
11    pub tasks_per_second: u64,
12    pub artifacts_per_second: u64,
13    pub agent_registry_per_second: u64,
14    pub agents_per_second: u64,
15    pub mcp_registry_per_second: u64,
16    pub mcp_per_second: u64,
17    pub stream_per_second: u64,
18    pub content_per_second: u64,
19    pub burst_multiplier: u64,
20    pub disabled: bool,
21    pub tier_multipliers: TierMultipliers,
22}
23
24impl Default for RateLimitConfig {
25    fn default() -> Self {
26        Self {
27            oauth_public_per_second: 10,
28            oauth_auth_per_second: 10,
29            contexts_per_second: 100,
30            tasks_per_second: 50,
31            artifacts_per_second: 50,
32            agent_registry_per_second: 50,
33            agents_per_second: 20,
34            mcp_registry_per_second: 50,
35            mcp_per_second: 200,
36            stream_per_second: 100,
37            content_per_second: 50,
38            burst_multiplier: 3,
39            disabled: false,
40            tier_multipliers: TierMultipliers::default(),
41        }
42    }
43}
44
45impl RateLimitConfig {
46    pub fn production() -> Self {
47        Self::default()
48    }
49
50    pub fn testing() -> Self {
51        Self {
52            oauth_public_per_second: 10000,
53            oauth_auth_per_second: 10000,
54            contexts_per_second: 10000,
55            tasks_per_second: 10000,
56            artifacts_per_second: 10000,
57            agent_registry_per_second: 10000,
58            agents_per_second: 10000,
59            mcp_registry_per_second: 10000,
60            mcp_per_second: 10000,
61            stream_per_second: 10000,
62            content_per_second: 10000,
63            burst_multiplier: 100,
64            disabled: false,
65            tier_multipliers: TierMultipliers::default(),
66        }
67    }
68
69    pub fn disabled() -> Self {
70        let mut config = Self::testing();
71        config.disabled = true;
72        config
73    }
74
75    #[must_use]
76    pub fn effective_limit(&self, base_rate: u64, tier: RateLimitTier) -> u64 {
77        let multiplier = self.tier_multiplier(tier);
78        let base_capped = u32::try_from(base_rate).unwrap_or(u32::MAX);
79        let scaled = f64::from(base_capped) * multiplier;
80        let clamped = scaled.clamp(1.0, f64::from(u32::MAX));
81        #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
82        let result = clamped as u64;
83        result
84    }
85
86    pub const fn tier_multiplier(&self, tier: RateLimitTier) -> f64 {
87        match tier {
88            RateLimitTier::Admin => self.tier_multipliers.admin,
89            RateLimitTier::User => self.tier_multipliers.user,
90            RateLimitTier::A2a => self.tier_multipliers.a2a,
91            RateLimitTier::Mcp => self.tier_multipliers.mcp,
92            RateLimitTier::Service => self.tier_multipliers.service,
93            RateLimitTier::Anon => self.tier_multipliers.anon,
94        }
95    }
96}
97
98impl From<&RateLimitsConfig> for RateLimitConfig {
99    fn from(config: &RateLimitsConfig) -> Self {
100        Self {
101            oauth_public_per_second: config.oauth_public_per_second,
102            oauth_auth_per_second: config.oauth_auth_per_second,
103            contexts_per_second: config.contexts_per_second,
104            tasks_per_second: config.tasks_per_second,
105            artifacts_per_second: config.artifacts_per_second,
106            agent_registry_per_second: config.agent_registry_per_second,
107            agents_per_second: config.agents_per_second,
108            mcp_registry_per_second: config.mcp_registry_per_second,
109            mcp_per_second: config.mcp_per_second,
110            stream_per_second: config.stream_per_second,
111            content_per_second: config.content_per_second,
112            burst_multiplier: config.burst_multiplier,
113            disabled: config.disabled,
114            tier_multipliers: config.tier_multipliers,
115        }
116    }
117}