1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
4#[serde(rename_all = "lowercase")]
5pub enum TenantTier {
6 Free,
7 Dev,
8 Pro,
9 Enterprise,
10}
11
12impl TenantTier {
13 pub fn from_str(s: &str) -> Option<Self> {
14 match s.to_lowercase().as_str() {
15 "free" => Some(TenantTier::Free),
16 "dev" => Some(TenantTier::Dev),
17 "pro" => Some(TenantTier::Pro),
18 "enterprise" => Some(TenantTier::Enterprise),
19 _ => None,
20 }
21 }
22
23 pub fn as_str(&self) -> &'static str {
24 match self {
25 TenantTier::Free => "free",
26 TenantTier::Dev => "dev",
27 TenantTier::Pro => "pro",
28 TenantTier::Enterprise => "enterprise",
29 }
30 }
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct TenantQuota {
35 pub tier: TenantTier,
36 pub requests_per_month: u64,
37 pub tokens_per_month: u64,
38 pub max_agents: u32,
39 pub requests_per_day: u64,
40}
41
42impl Default for TenantQuota {
43 fn default() -> Self {
44 Self::free()
45 }
46}
47
48impl TenantQuota {
49 pub fn free() -> Self {
50 Self {
51 tier: TenantTier::Free,
52 requests_per_month: 1_000,
53 tokens_per_month: 100_000,
54 max_agents: 1,
55 requests_per_day: 50,
56 }
57 }
58
59 pub fn dev() -> Self {
60 Self {
61 tier: TenantTier::Dev,
62 requests_per_month: 50_000,
63 tokens_per_month: 5_000_000,
64 max_agents: 10,
65 requests_per_day: 2_000,
66 }
67 }
68
69 pub fn pro() -> Self {
70 Self {
71 tier: TenantTier::Pro,
72 requests_per_month: 500_000,
73 tokens_per_month: 50_000_000,
74 max_agents: u32::MAX,
75 requests_per_day: 20_000,
76 }
77 }
78
79 pub fn enterprise() -> Self {
80 Self {
81 tier: TenantTier::Enterprise,
82 requests_per_month: u64::MAX,
83 tokens_per_month: u64::MAX,
84 max_agents: u32::MAX,
85 requests_per_day: u64::MAX,
86 }
87 }
88
89 pub fn from_tier(tier: &TenantTier) -> Self {
90 match tier {
91 TenantTier::Free => Self::free(),
92 TenantTier::Dev => Self::dev(),
93 TenantTier::Pro => Self::pro(),
94 TenantTier::Enterprise => Self::enterprise(),
95 }
96 }
97}
98
99#[derive(Debug, Clone, Serialize, Deserialize)]
100pub struct Tenant {
101 pub id: String,
102 pub name: String,
103 pub tier: TenantTier,
104 pub created_at: i64,
105 pub updated_at: i64,
106}
107
108impl Tenant {
109 pub fn new(id: String, name: String, tier: TenantTier) -> Self {
110 let now = chrono::Utc::now().timestamp();
111 Self {
112 id,
113 name,
114 tier,
115 created_at: now,
116 updated_at: now,
117 }
118 }
119}
120
121#[derive(Debug, Clone, Serialize, Deserialize)]
122pub struct ApiKey {
123 pub id: String,
124 pub tenant_id: String,
125 pub key_hash: String,
126 pub key_prefix: String,
127 pub name: String,
128 pub is_active: bool,
129 pub created_at: i64,
130 pub expires_at: Option<i64>,
131}
132
133impl ApiKey {
134 pub fn new(
135 id: String,
136 tenant_id: String,
137 key_hash: String,
138 key_prefix: String,
139 name: String,
140 ) -> Self {
141 Self {
142 id,
143 tenant_id,
144 key_hash,
145 key_prefix,
146 name,
147 is_active: true,
148 created_at: chrono::Utc::now().timestamp(),
149 expires_at: None,
150 }
151 }
152}
153
154#[derive(Debug, Clone, Serialize, Deserialize)]
155pub struct TenantContext {
156 pub tenant_id: String,
157 pub tier: TenantTier,
158 pub quota: TenantQuota,
159}
160
161impl TenantContext {
162 pub fn new(tenant_id: String, tier: TenantTier) -> Self {
163 Self {
164 tenant_id,
165 tier,
166 quota: TenantQuota::from_tier(&tier),
167 }
168 }
169
170 pub fn can_make_request(&self, monthly_requests: u64, daily_requests: u64) -> bool {
171 if monthly_requests >= self.quota.requests_per_month {
172 return false;
173 }
174 if daily_requests >= self.quota.requests_per_day {
175 return false;
176 }
177 true
178 }
179
180 pub fn can_use_tokens(&self, monthly_tokens: u64, additional_tokens: u64) -> bool {
181 let Some(new_total) = monthly_tokens.checked_add(additional_tokens) else {
182 return false;
183 };
184 new_total <= self.quota.tokens_per_month
185 }
186}
187
188#[cfg(test)]
189mod tests {
190 use super::*;
191
192 #[test]
193 fn test_tier_from_str() {
194 assert_eq!(TenantTier::from_str("free"), Some(TenantTier::Free));
195 assert_eq!(TenantTier::from_str("dev"), Some(TenantTier::Dev));
196 assert_eq!(TenantTier::from_str("pro"), Some(TenantTier::Pro));
197 assert_eq!(
198 TenantTier::from_str("enterprise"),
199 Some(TenantTier::Enterprise)
200 );
201 assert_eq!(TenantTier::from_str("unknown"), None);
202 }
203
204 #[test]
205 fn test_tier_as_str() {
206 assert_eq!(TenantTier::Free.as_str(), "free");
207 assert_eq!(TenantTier::Dev.as_str(), "dev");
208 assert_eq!(TenantTier::Pro.as_str(), "pro");
209 assert_eq!(TenantTier::Enterprise.as_str(), "enterprise");
210 }
211
212 #[test]
213 fn test_free_quota() {
214 let quota = TenantQuota::free();
215 assert_eq!(quota.tier, TenantTier::Free);
216 assert_eq!(quota.requests_per_month, 1_000);
217 assert_eq!(quota.tokens_per_month, 100_000);
218 assert_eq!(quota.max_agents, 1);
219 assert_eq!(quota.requests_per_day, 50);
220 }
221
222 #[test]
223 fn test_dev_quota() {
224 let quota = TenantQuota::dev();
225 assert_eq!(quota.tier, TenantTier::Dev);
226 assert_eq!(quota.requests_per_month, 50_000);
227 assert_eq!(quota.tokens_per_month, 5_000_000);
228 assert_eq!(quota.max_agents, 10);
229 assert_eq!(quota.requests_per_day, 2_000);
230 }
231
232 #[test]
233 fn test_pro_quota() {
234 let quota = TenantQuota::pro();
235 assert_eq!(quota.tier, TenantTier::Pro);
236 assert_eq!(quota.requests_per_month, 500_000);
237 assert_eq!(quota.tokens_per_month, 50_000_000);
238 assert_eq!(quota.max_agents, u32::MAX);
239 assert_eq!(quota.requests_per_day, 20_000);
240 }
241
242 #[test]
243 fn test_enterprise_quota() {
244 let quota = TenantQuota::enterprise();
245 assert_eq!(quota.tier, TenantTier::Enterprise);
246 assert_eq!(quota.requests_per_month, u64::MAX);
247 assert_eq!(quota.tokens_per_month, u64::MAX);
248 }
249
250 #[test]
251 fn test_quota_from_tier() {
252 assert_eq!(
253 TenantQuota::from_tier(&TenantTier::Free).requests_per_month,
254 1_000
255 );
256 assert_eq!(
257 TenantQuota::from_tier(&TenantTier::Dev).requests_per_month,
258 50_000
259 );
260 assert_eq!(
261 TenantQuota::from_tier(&TenantTier::Pro).requests_per_month,
262 500_000
263 );
264 }
265
266 #[test]
267 fn test_tenant_context_can_make_request() {
268 let ctx = TenantContext::new("test".to_string(), TenantTier::Free);
269 assert!(ctx.can_make_request(0, 0));
270 assert!(ctx.can_make_request(999, 0));
271 assert!(ctx.can_make_request(0, 49));
272 assert!(!ctx.can_make_request(1000, 0));
273 assert!(!ctx.can_make_request(0, 50));
274 }
275
276 #[test]
277 fn test_tenant_context_can_use_tokens() {
278 let ctx = TenantContext::new("test".to_string(), TenantTier::Free);
279 assert!(ctx.can_use_tokens(0, 100_000));
280 assert!(ctx.can_use_tokens(50_000, 50_000));
281 assert!(!ctx.can_use_tokens(50_000, 50_001));
282 assert!(!ctx.can_use_tokens(100_000, 1));
283 }
284
285 #[test]
286 fn test_tenant_creation() {
287 let tenant = Tenant::new("t1".to_string(), "Test Tenant".to_string(), TenantTier::Dev);
288 assert_eq!(tenant.id, "t1");
289 assert_eq!(tenant.name, "Test Tenant");
290 assert_eq!(tenant.tier, TenantTier::Dev);
291 assert!(tenant.created_at > 0);
292 }
293
294 #[test]
295 fn test_api_key_creation() {
296 let key = ApiKey::new(
297 "k1".to_string(),
298 "t1".to_string(),
299 "hash123".to_string(),
300 "ares_abc".to_string(),
301 "Test Key".to_string(),
302 );
303 assert_eq!(key.id, "k1");
304 assert_eq!(key.tenant_id, "t1");
305 assert!(key.is_active);
306 assert!(key.created_at > 0);
307 }
308}