litellm_rs/core/keys/
manager.rs

1//! API Key Manager
2//!
3//! This module provides the main API key management functionality.
4
5use super::repository::KeyRepository;
6use super::types::{
7    CreateKeyConfig, KeyInfo, KeyStatus, KeyUsageStats, ManagedApiKey, UpdateKeyConfig,
8    UsageRecord, VerifyKeyResult,
9};
10use crate::utils::auth::crypto::keys::{extract_api_key_prefix, generate_api_key, hash_api_key};
11use crate::utils::error::gateway_error::{GatewayError, Result};
12use chrono::Utc;
13use dashmap::DashMap;
14use std::sync::Arc;
15use std::time::{Duration, Instant};
16use tokio::sync::{Semaphore, oneshot};
17use tracing::{debug, info, warn};
18use uuid::Uuid;
19
20/// Minimum interval between DB writes for the same key's last_used timestamp.
21const LAST_USED_THROTTLE: Duration = Duration::from_secs(5 * 60);
22/// Maximum number of entries retained in the last-used throttle cache.
23const LAST_USED_CACHE_MAX_ENTRIES: usize = 10_000;
24/// Bound accepted usage writes, including writes whose request waiter was cancelled.
25const MAX_USAGE_WRITES: usize = 1024;
26
27/// API Key Manager for handling all key operations
28#[derive(Clone)]
29pub struct KeyManager {
30    /// Repository for key storage
31    repository: Arc<dyn KeyRepository>,
32    /// Tracks when each key's `last_used_at` was last persisted.
33    last_used_cache: Arc<DashMap<Uuid, Instant>>,
34    usage_write_slots: Arc<Semaphore>,
35    /// Optional HMAC secret for key hashing.
36    hmac_secret: Option<Arc<str>>,
37}
38
39impl std::fmt::Debug for KeyManager {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        f.debug_struct("KeyManager")
42            .field("repository", &"<KeyRepository>")
43            .field("last_used_cache_size", &self.last_used_cache.len())
44            .field(
45                "hmac_secret",
46                &self.hmac_secret.as_ref().map(|_| "[REDACTED]"),
47            )
48            .finish()
49    }
50}
51
52impl KeyManager {
53    /// Create a new KeyManager with the given repository
54    pub fn new<R: KeyRepository + 'static>(repository: R) -> Self {
55        Self {
56            repository: Arc::new(repository),
57            last_used_cache: Arc::new(DashMap::new()),
58            usage_write_slots: Arc::new(Semaphore::new(MAX_USAGE_WRITES)),
59            hmac_secret: None,
60        }
61    }
62
63    /// Create a new KeyManager with an Arc repository
64    pub fn with_arc_repository(repository: Arc<dyn KeyRepository>) -> Self {
65        Self {
66            repository,
67            last_used_cache: Arc::new(DashMap::new()),
68            usage_write_slots: Arc::new(Semaphore::new(MAX_USAGE_WRITES)),
69            hmac_secret: None,
70        }
71    }
72
73    /// Set the HMAC secret for key hashing
74    pub fn with_hmac_secret(mut self, secret: Option<String>) -> Self {
75        self.hmac_secret = secret.map(Arc::from);
76        self
77    }
78
79    /// Get the HMAC secret as Option<&str>
80    fn hmac_secret(&self) -> Option<&str> {
81        self.hmac_secret.as_deref()
82    }
83
84    /// Bound the last-used throttle cache when the manager is process-scoped.
85    fn prune_last_used_cache(&self, now: Instant) {
86        self.last_used_cache
87            .retain(|_, last_persisted| now.duration_since(*last_persisted) < LAST_USED_THROTTLE);
88
89        while self.last_used_cache.len() >= LAST_USED_CACHE_MAX_ENTRIES {
90            let Some(oldest_key) = self
91                .last_used_cache
92                .iter()
93                .min_by_key(|entry| *entry.value())
94                .map(|entry| *entry.key())
95            else {
96                break;
97            };
98
99            self.last_used_cache.remove(&oldest_key);
100        }
101    }
102
103    #[cfg(test)]
104    #[cfg(feature = "gateway")]
105    pub(crate) fn has_recorded_key_usage(&self, key_id: Uuid) -> bool {
106        self.last_used_cache.contains_key(&key_id)
107    }
108
109    /// Generate a new API key
110    ///
111    /// Returns a tuple of (key_id, raw_key). The raw_key should be shown to the user
112    /// only once and never stored. Only the hash is stored.
113    pub async fn generate_key(&self, config: CreateKeyConfig) -> Result<(Uuid, String)> {
114        info!("Generating new API key: {}", config.name);
115
116        // Validate configuration
117        self.validate_create_config(&config)?;
118
119        // Generate the raw key
120        let raw_key = generate_api_key();
121        let key_hash = hash_api_key(&raw_key, self.hmac_secret());
122        let key_prefix = extract_api_key_prefix(&raw_key);
123
124        let now = Utc::now();
125
126        // Create the managed key
127        let managed_key = ManagedApiKey {
128            id: Uuid::new_v4(),
129            key_hash,
130            key_prefix,
131            name: config.name,
132            description: config.description,
133            user_id: config.user_id,
134            team_id: config.team_id,
135            budget_id: config.budget_id,
136            permissions: config.permissions,
137            rate_limits: config.rate_limits,
138            status: KeyStatus::Active,
139            expires_at: config.expires_at,
140            created_at: now,
141            updated_at: now,
142            last_used_at: None,
143            usage_stats: KeyUsageStats::new(),
144            metadata: config.metadata,
145        };
146
147        let key_id = managed_key.id;
148
149        // Store the key
150        self.repository.create(managed_key).await?;
151
152        info!("API key generated successfully: {}", key_id);
153
154        // Return key_id and raw_key (raw_key should only be shown once)
155        Ok((key_id, raw_key))
156    }
157
158    /// Validate a raw API key
159    ///
160    /// Returns verification result with key info if valid.
161    pub async fn validate_key(&self, raw_key: &str) -> Result<VerifyKeyResult> {
162        let result = self.validate_key_without_usage_update(raw_key).await?;
163        self.record_validated_key_usage(&result);
164        Ok(result)
165    }
166
167    /// Validate a raw API key without recording it as used.
168    ///
169    /// Authorization-sensitive callers must use this method and call
170    /// [`Self::record_validated_key_usage`] only after access is granted.
171    pub(crate) async fn validate_key_without_usage_update(
172        &self,
173        raw_key: &str,
174    ) -> Result<VerifyKeyResult> {
175        debug!("Validating API key");
176
177        // Hash the provided key
178        let key_hash = hash_api_key(raw_key, self.hmac_secret());
179
180        // Find key by hash
181        let key = match self.repository.find_by_hash(&key_hash).await? {
182            Some(k) => k,
183            None => {
184                debug!("API key not found");
185                return Ok(VerifyKeyResult {
186                    valid: false,
187                    key: None,
188                    invalid_reason: Some("API key not found".to_string()),
189                });
190            }
191        };
192
193        // Check if key is active
194        if key.status == KeyStatus::Revoked {
195            debug!("API key is revoked");
196            return Ok(VerifyKeyResult {
197                valid: false,
198                key: Some(KeyInfo::from(&key)),
199                invalid_reason: Some("API key has been revoked".to_string()),
200            });
201        }
202
203        // Check expiration
204        if let Some(expires_at) = key.expires_at
205            && Utc::now() > expires_at
206        {
207            debug!("API key is expired");
208            return Ok(VerifyKeyResult {
209                valid: false,
210                key: Some(KeyInfo::from(&key)),
211                invalid_reason: Some("API key has expired".to_string()),
212            });
213        }
214
215        debug!("API key validated successfully");
216        Ok(VerifyKeyResult {
217            valid: true,
218            key: Some(KeyInfo::from(&key)),
219            invalid_reason: None,
220        })
221    }
222
223    /// Record usage for a successfully validated key, throttled to once per
224    /// five minutes per key.
225    pub(crate) fn record_validated_key_usage(&self, result: &VerifyKeyResult) {
226        if !result.valid {
227            return;
228        }
229        let Some(key_id) = result.key.as_ref().map(|key| key.id) else {
230            return;
231        };
232
233        let now = Instant::now();
234        self.prune_last_used_cache(now);
235        let should_update = match self.last_used_cache.get(&key_id) {
236            Some(last_persisted) => now.duration_since(*last_persisted) >= LAST_USED_THROTTLE,
237            None => true,
238        };
239
240        if should_update {
241            self.last_used_cache.insert(key_id, now);
242            let repo = self.repository.clone();
243            tokio::spawn(async move {
244                if let Err(e) = repo.update_last_used(key_id).await {
245                    warn!("Failed to update last used timestamp: {}", e);
246                }
247            });
248        }
249    }
250
251    /// Revoke an API key
252    pub async fn revoke_key(&self, key_id: Uuid) -> Result<()> {
253        info!("Revoking API key: {}", key_id);
254
255        // Verify key exists
256        let key = self
257            .repository
258            .find_by_id(key_id)
259            .await?
260            .ok_or_else(|| GatewayError::not_found("API key not found"))?;
261
262        if key.status == KeyStatus::Revoked {
263            return Err(GatewayError::conflict("API key is already revoked"));
264        }
265
266        self.repository
267            .update_status(key_id, KeyStatus::Revoked)
268            .await?;
269
270        info!("API key revoked successfully: {}", key_id);
271        Ok(())
272    }
273
274    /// Rotate an API key (generate new key, revoke old one)
275    ///
276    /// Returns a tuple of (new_key_id, new_raw_key). The new_raw_key should be shown
277    /// to the user only once.
278    pub async fn rotate_key(&self, key_id: Uuid) -> Result<(Uuid, String)> {
279        info!("Rotating API key: {}", key_id);
280
281        // Get existing key
282        let old_key = self
283            .repository
284            .find_by_id(key_id)
285            .await?
286            .ok_or_else(|| GatewayError::not_found("API key not found"))?;
287
288        if old_key.status == KeyStatus::Revoked {
289            return Err(GatewayError::conflict("Cannot rotate a revoked key"));
290        }
291
292        // Create new key with same configuration
293        let config = CreateKeyConfig {
294            name: format!("{} (rotated)", old_key.name),
295            description: old_key.description,
296            user_id: old_key.user_id,
297            team_id: old_key.team_id,
298            budget_id: old_key.budget_id,
299            permissions: old_key.permissions,
300            rate_limits: old_key.rate_limits,
301            expires_at: old_key.expires_at,
302            metadata: old_key.metadata,
303        };
304
305        let (new_key_id, new_raw_key) = self.generate_key(config).await?;
306
307        // Revoke the old key
308        self.repository
309            .update_status(key_id, KeyStatus::Revoked)
310            .await?;
311
312        info!("API key rotated successfully: {} -> {}", key_id, new_key_id);
313        Ok((new_key_id, new_raw_key))
314    }
315
316    /// Update an API key's configuration
317    pub async fn update_key(&self, key_id: Uuid, config: UpdateKeyConfig) -> Result<KeyInfo> {
318        info!("Updating API key: {}", key_id);
319
320        // Verify key exists and is not revoked
321        let existing = self
322            .repository
323            .find_by_id(key_id)
324            .await?
325            .ok_or_else(|| GatewayError::not_found("API key not found"))?;
326
327        if existing.status == KeyStatus::Revoked {
328            return Err(GatewayError::conflict("Cannot update a revoked key"));
329        }
330
331        let updated = self.repository.update(key_id, config).await?;
332
333        info!("API key updated successfully: {}", key_id);
334        Ok(KeyInfo::from(&updated))
335    }
336
337    /// Get key information by ID
338    pub async fn get_key(&self, key_id: Uuid) -> Result<Option<KeyInfo>> {
339        let key = self.repository.find_by_id(key_id).await?;
340        Ok(key.map(|k| KeyInfo::from(&k)))
341    }
342
343    /// Get key usage statistics
344    pub async fn get_usage_stats(&self, key_id: Uuid) -> Result<KeyUsageStats> {
345        let key = self
346            .repository
347            .find_by_id(key_id)
348            .await?
349            .ok_or_else(|| GatewayError::not_found("API key not found"))?;
350
351        Ok(key.usage_stats)
352    }
353
354    /// Record usage for a key
355    pub async fn record_usage(&self, key_id: Uuid, tokens: u64, cost: f64) -> Result<()> {
356        self.record_usage_record(key_id, UsageRecord::priced(tokens, cost))
357            .await
358    }
359
360    /// Record usage for a key from an explicit usage record.
361    /// On Tokio, a write accepted on the first poll survives cancellation of its waiter.
362    /// Other executors directly await the repository under the same capacity bound,
363    /// preserving their existing execution and cancellation behavior.
364    /// Capacity exhaustion and repository failures remain best-effort write errors;
365    /// runtime shutdown does not provide durable retry.
366    pub async fn record_usage_record(&self, key_id: Uuid, record: UsageRecord) -> Result<()> {
367        let permit = self
368            .usage_write_slots
369            .clone()
370            .try_acquire_owned()
371            .map_err(|_| GatewayError::unavailable("Key usage write capacity exhausted"))?;
372        let Ok(runtime) = tokio::runtime::Handle::try_current() else {
373            let _permit = permit;
374            return self.repository.update_usage(key_id, record).await;
375        };
376        let repository = self.repository.clone();
377        let (sender, receiver) = oneshot::channel();
378        runtime.spawn(async move {
379            let _permit = permit;
380            let result = repository.update_usage(key_id, record).await;
381            if let Err(Err(error)) = sender.send(result) {
382                tracing::error!(%key_id, %error, "failed to record key usage after waiter cancellation");
383            }
384        });
385        receiver.await.map_err(|_| {
386            GatewayError::internal("Key usage write task stopped before returning its result")
387        })?
388    }
389
390    /// List keys for a user
391    pub async fn list_user_keys(&self, user_id: Uuid) -> Result<Vec<KeyInfo>> {
392        let keys = self.repository.list_by_user(user_id).await?;
393        Ok(keys.iter().map(KeyInfo::from).collect())
394    }
395
396    /// List keys for a team
397    pub async fn list_team_keys(&self, team_id: Uuid) -> Result<Vec<KeyInfo>> {
398        let keys = self.repository.list_by_team(team_id).await?;
399        Ok(keys.iter().map(KeyInfo::from).collect())
400    }
401
402    /// List all keys with optional filtering
403    pub async fn list_keys(
404        &self,
405        status: Option<KeyStatus>,
406        limit: Option<usize>,
407        offset: Option<usize>,
408    ) -> Result<Vec<KeyInfo>> {
409        let keys = self.repository.list_all(status, limit, offset).await?;
410        Ok(keys.iter().map(KeyInfo::from).collect())
411    }
412
413    /// Delete a key permanently (use with caution)
414    pub async fn delete_key(&self, key_id: Uuid) -> Result<()> {
415        info!("Deleting API key permanently: {}", key_id);
416
417        // Verify key exists
418        self.repository
419            .find_by_id(key_id)
420            .await?
421            .ok_or_else(|| GatewayError::not_found("API key not found"))?;
422
423        self.repository.delete(key_id).await?;
424
425        info!("API key deleted permanently: {}", key_id);
426        Ok(())
427    }
428
429    /// Cleanup expired keys
430    pub async fn cleanup_expired_keys(&self) -> Result<u64> {
431        info!("Cleaning up expired API keys");
432        let deleted = self.repository.delete_expired().await?;
433        info!("Cleaned up {} expired API keys", deleted);
434        Ok(deleted)
435    }
436
437    /// Count keys
438    pub async fn count_keys(&self, status: Option<KeyStatus>) -> Result<u64> {
439        self.repository.count(status).await
440    }
441
442    /// Validate create configuration
443    fn validate_create_config(&self, config: &CreateKeyConfig) -> Result<()> {
444        if config.name.is_empty() {
445            return Err(GatewayError::validation("Key name cannot be empty"));
446        }
447
448        if config.name.len() > 255 {
449            return Err(GatewayError::validation(
450                "Key name cannot exceed 255 characters",
451            ));
452        }
453
454        if let Some(ref desc) = config.description
455            && desc.len() > 1000
456        {
457            return Err(GatewayError::validation(
458                "Key description cannot exceed 1000 characters",
459            ));
460        }
461
462        // Check expiration is in the future
463        if let Some(expires_at) = config.expires_at
464            && expires_at <= Utc::now()
465        {
466            return Err(GatewayError::validation(
467                "Expiration date must be in the future",
468            ));
469        }
470
471        Ok(())
472    }
473}
474
475#[cfg(test)]
476mod manager_tests {
477    use super::*;
478    use crate::core::keys::repository::InMemoryKeyRepository;
479    use crate::core::keys::types::{KeyPermissions, KeyRateLimits};
480
481    fn create_manager() -> KeyManager {
482        KeyManager::new(InMemoryKeyRepository::new())
483    }
484
485    #[test]
486    fn test_prune_last_used_cache_enforces_hard_cap() {
487        let manager = create_manager();
488        let now = Instant::now();
489
490        for _ in 0..=LAST_USED_CACHE_MAX_ENTRIES {
491            manager.last_used_cache.insert(Uuid::new_v4(), now);
492        }
493
494        manager.prune_last_used_cache(now);
495
496        assert!(manager.last_used_cache.len() < LAST_USED_CACHE_MAX_ENTRIES);
497    }
498
499    #[test]
500    fn test_clone_reuses_hmac_secret_allocation() {
501        let manager = create_manager().with_hmac_secret(Some("shared-secret".to_string()));
502        let cloned = manager.clone();
503
504        let (Some(original_secret), Some(cloned_secret)) =
505            (manager.hmac_secret.as_ref(), cloned.hmac_secret.as_ref())
506        else {
507            panic!("secret should be set on both managers");
508        };
509
510        assert!(Arc::ptr_eq(original_secret, cloned_secret));
511        assert_eq!(cloned.hmac_secret(), Some("shared-secret"));
512    }
513
514    #[tokio::test]
515    async fn test_generate_key() {
516        let manager = create_manager();
517
518        let config = CreateKeyConfig {
519            name: "Test Key".to_string(),
520            ..Default::default()
521        };
522
523        let (key_id, raw_key) = manager.generate_key(config).await.unwrap();
524
525        assert!(!raw_key.is_empty());
526        assert!(raw_key.starts_with("gw-"));
527
528        // Verify key was stored
529        let key_info = manager.get_key(key_id).await.unwrap();
530        assert!(key_info.is_some());
531        assert_eq!(key_info.unwrap().name, "Test Key");
532    }
533
534    #[tokio::test]
535    async fn test_validate_key() {
536        let manager = create_manager();
537
538        let config = CreateKeyConfig {
539            name: "Validation Test".to_string(),
540            ..Default::default()
541        };
542
543        let (_, raw_key) = manager.generate_key(config).await.unwrap();
544
545        let result = manager.validate_key(&raw_key).await.unwrap();
546        assert!(result.valid);
547        assert!(result.key.is_some());
548    }
549
550    #[tokio::test]
551    async fn test_validate_invalid_key() {
552        let manager = create_manager();
553
554        let result = manager.validate_key("invalid-key").await.unwrap();
555        assert!(!result.valid);
556        assert!(result.invalid_reason.is_some());
557    }
558
559    #[tokio::test]
560    async fn test_revoke_key() {
561        let manager = create_manager();
562
563        let config = CreateKeyConfig {
564            name: "Revoke Test".to_string(),
565            ..Default::default()
566        };
567
568        let (key_id, raw_key) = manager.generate_key(config).await.unwrap();
569
570        manager.revoke_key(key_id).await.unwrap();
571
572        // Key should no longer be valid
573        let result = manager.validate_key(&raw_key).await.unwrap();
574        assert!(!result.valid);
575        assert!(result.invalid_reason.as_ref().unwrap().contains("revoked"));
576    }
577
578    #[tokio::test]
579    async fn test_rotate_key() {
580        let manager = create_manager();
581
582        let config = CreateKeyConfig {
583            name: "Rotate Test".to_string(),
584            permissions: KeyPermissions::full_access(),
585            rate_limits: KeyRateLimits::standard(),
586            ..Default::default()
587        };
588
589        let (old_key_id, old_raw_key) = manager.generate_key(config).await.unwrap();
590
591        let (new_key_id, new_raw_key) = manager.rotate_key(old_key_id).await.unwrap();
592
593        // Old key should be revoked
594        let old_result = manager.validate_key(&old_raw_key).await.unwrap();
595        assert!(!old_result.valid);
596
597        // New key should be valid
598        let new_result = manager.validate_key(&new_raw_key).await.unwrap();
599        assert!(new_result.valid);
600
601        // New key should have same permissions
602        let new_key = manager.get_key(new_key_id).await.unwrap().unwrap();
603        assert!(new_key.name.contains("rotated"));
604    }
605
606    #[tokio::test]
607    async fn test_update_key() {
608        let manager = create_manager();
609
610        let config = CreateKeyConfig {
611            name: "Update Test".to_string(),
612            ..Default::default()
613        };
614
615        let (key_id, _) = manager.generate_key(config).await.unwrap();
616
617        let update = UpdateKeyConfig {
618            name: Some("Updated Name".to_string()),
619            ..Default::default()
620        };
621
622        let updated = manager.update_key(key_id, update).await.unwrap();
623        assert_eq!(updated.name, "Updated Name");
624    }
625
626    #[tokio::test]
627    async fn test_list_user_keys() {
628        let manager = create_manager();
629        let user_id = Uuid::new_v4();
630
631        for i in 0..3 {
632            let config = CreateKeyConfig {
633                name: format!("User Key {}", i),
634                user_id: Some(user_id),
635                ..Default::default()
636            };
637            manager.generate_key(config).await.unwrap();
638        }
639
640        let keys = manager.list_user_keys(user_id).await.unwrap();
641        assert_eq!(keys.len(), 3);
642    }
643
644    #[tokio::test]
645    async fn test_expired_key_validation() {
646        let manager = create_manager();
647
648        // We need to create a key and then manually expire it
649        // Since we can't create an already-expired key, we'll update it
650
651        let config = CreateKeyConfig {
652            name: "Expiring Key".to_string(),
653            expires_at: Some(Utc::now() + chrono::Duration::seconds(1)),
654            ..Default::default()
655        };
656
657        let (_, raw_key) = manager.generate_key(config).await.unwrap();
658
659        // Wait for expiration
660        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
661
662        let result = manager.validate_key(&raw_key).await.unwrap();
663        assert!(!result.valid);
664        assert!(result.invalid_reason.as_ref().unwrap().contains("expired"));
665    }
666
667    #[tokio::test]
668    async fn test_validation_empty_name() {
669        let manager = create_manager();
670
671        let config = CreateKeyConfig {
672            name: "".to_string(),
673            ..Default::default()
674        };
675
676        let result = manager.generate_key(config).await;
677        assert!(result.is_err());
678    }
679
680    #[tokio::test]
681    async fn test_delete_key() {
682        let manager = create_manager();
683
684        let config = CreateKeyConfig {
685            name: "Delete Test".to_string(),
686            ..Default::default()
687        };
688
689        let (key_id, _) = manager.generate_key(config).await.unwrap();
690
691        manager.delete_key(key_id).await.unwrap();
692
693        let key = manager.get_key(key_id).await.unwrap();
694        assert!(key.is_none());
695    }
696
697    #[tokio::test]
698    async fn test_count_keys() {
699        let manager = create_manager();
700
701        for i in 0..5 {
702            let config = CreateKeyConfig {
703                name: format!("Count Key {}", i),
704                ..Default::default()
705            };
706            manager.generate_key(config).await.unwrap();
707        }
708
709        let count = manager.count_keys(None).await.unwrap();
710        assert_eq!(count, 5);
711    }
712
713    #[tokio::test]
714    async fn cancelled_waiter_retains_write_capacity_until_repository_finishes() {
715        let mut manager = KeyManager::new(InMemoryKeyRepository::new());
716        manager.usage_write_slots = Arc::new(Semaphore::new(1));
717        let clone = manager.clone();
718        let (key_id, _) = manager
719            .generate_key(CreateKeyConfig {
720                name: "Cancelled usage waiter".to_string(),
721                ..Default::default()
722            })
723            .await
724            .unwrap();
725
726        let mut write = Box::pin(manager.record_usage(key_id, 100, 0.25));
727        assert!(futures::poll!(&mut write).is_pending());
728        drop(write);
729        // On this current-thread runtime the owned task has not run yet.
730        assert!(matches!(
731            clone.record_usage(key_id, 200, 0.50).await,
732            Err(GatewayError::Unavailable(message))
733                if message == "Key usage write capacity exhausted"
734        ));
735        tokio::time::timeout(Duration::from_secs(1), async {
736            loop {
737                if manager
738                    .get_usage_stats(key_id)
739                    .await
740                    .unwrap()
741                    .total_requests
742                    == 1
743                {
744                    break;
745                }
746                tokio::task::yield_now().await;
747            }
748        })
749        .await
750        .unwrap();
751
752        assert_eq!(manager.usage_write_slots.available_permits(), 1);
753        clone.record_usage(key_id, 200, 0.50).await.unwrap();
754        let stats = manager.get_usage_stats(key_id).await.unwrap();
755        assert_eq!(stats.total_requests, 2);
756        assert_eq!(stats.total_tokens, 300);
757        assert_eq!(stats.total_cost, 0.75);
758    }
759
760    #[test]
761    fn usage_recording_preserves_non_tokio_executor_compatibility() {
762        assert!(tokio::runtime::Handle::try_current().is_err());
763        for local_pool in [false, true] {
764            let exercise = async {
765                let mut manager = create_manager();
766                manager.usage_write_slots = Arc::new(Semaphore::new(1));
767                let (key_id, _) = manager
768                    .generate_key(CreateKeyConfig {
769                        name: "Executor-independent usage".to_string(),
770                        ..Default::default()
771                    })
772                    .await
773                    .unwrap();
774
775                drop(manager.record_usage(key_id, 999, 9.0));
776                assert_eq!(
777                    manager
778                        .get_usage_stats(key_id)
779                        .await
780                        .unwrap()
781                        .total_requests,
782                    0
783                );
784                manager.record_usage(key_id, 100, 0.25).await.unwrap();
785                manager
786                    .record_usage_record(key_id, UsageRecord::unpriced(50, 0.125, "allow_unpriced"))
787                    .await
788                    .unwrap();
789                let stats = manager.get_usage_stats(key_id).await.unwrap();
790                assert_eq!(stats.total_requests, 2);
791                assert_eq!(stats.total_tokens, 150);
792                assert!((stats.total_cost - 0.375).abs() < f64::EPSILON);
793                assert_eq!(stats.unpriced_requests, 1);
794                assert_eq!(stats.unpriced_tokens, 50);
795                assert!((stats.unpriced_cost - 0.125).abs() < f64::EPSILON);
796                assert_eq!(manager.usage_write_slots.available_permits(), 1);
797
798                let permit = manager
799                    .usage_write_slots
800                    .clone()
801                    .try_acquire_owned()
802                    .unwrap();
803                assert!(matches!(
804                    manager.record_usage(key_id, 200, 0.5).await,
805                    Err(GatewayError::Unavailable(message))
806                        if message == "Key usage write capacity exhausted"
807                ));
808                drop(permit);
809                assert_eq!(
810                    manager
811                        .get_usage_stats(key_id)
812                        .await
813                        .unwrap()
814                        .total_requests,
815                    2
816                );
817                // Preserve the in-memory repository's missing-key no-op contract.
818                let missing_id = Uuid::new_v4();
819                manager.record_usage(missing_id, 1, 0.01).await.unwrap();
820                assert!(manager.get_key(missing_id).await.unwrap().is_none());
821                assert_eq!(manager.usage_write_slots.available_permits(), 1);
822            };
823            if local_pool {
824                futures::executor::LocalPool::new().run_until(exercise);
825            } else {
826                futures::executor::block_on(exercise);
827            }
828        }
829    }
830}
litellm_rs/core/keys/manager.rs

litellm_rs/core/keys/
manager.rs
