Skip to main content

systemprompt_models/config/
rate_limits.rs

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