microclaw_engine/llm/
key_pool.rs

1/// Round-robin API-key pool. Providers read `current()` per request and
2/// call `advance()` on auth (401/403) or rate-limit (429) errors so a
3/// multi-account setup keeps working when one key is throttled or revoked.
4#[derive(Clone)]
5pub struct KeyPool {
6    keys: std::sync::Arc<Vec<String>>,
7    index: std::sync::Arc<std::sync::atomic::AtomicUsize>,
8    rotations: std::sync::Arc<std::sync::atomic::AtomicU64>,
9}
10
11impl KeyPool {
12    pub fn new(primary: &str, extra: &[String]) -> Self {
13        let mut keys: Vec<String> = Vec::with_capacity(1 + extra.len());
14        keys.push(primary.to_string());
15        for key in extra {
16            let key = key.trim();
17            if !key.is_empty() && !keys.iter().any(|k| k == key) {
18                keys.push(key.to_string());
19            }
20        }
21        KeyPool {
22            keys: std::sync::Arc::new(keys),
23            index: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
24            rotations: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)),
25        }
26    }
27
28    pub fn current(&self) -> String {
29        let idx = self.index.load(std::sync::atomic::Ordering::Relaxed) % self.keys.len();
30        self.keys[idx].clone()
31    }
32
33    /// Rotate to the next key. Returns `false` (and does nothing) when the
34    /// pool has a single key, so callers can skip a pointless retry.
35    pub fn advance(&self) -> bool {
36        if self.keys.len() <= 1 {
37            return false;
38        }
39        self.index
40            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
41        self.rotations
42            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
43        true
44    }
45
46    pub fn len(&self) -> usize {
47        self.keys.len()
48    }
49
50    pub fn is_empty(&self) -> bool {
51        self.keys.is_empty()
52    }
53
54    pub fn rotations(&self) -> u64 {
55        self.rotations.load(std::sync::atomic::Ordering::Relaxed)
56    }
57}
58
59/// Statuses where rotating to another API key can plausibly help: bad or
60/// revoked credentials and per-account rate limits.
61pub(crate) fn is_key_rotation_status(status: u16) -> bool {
62    matches!(status, 401 | 403 | 429)
63}
64
65#[cfg(test)]
66mod key_pool_tests {
67    use super::*;
68
69    #[test]
70    fn key_pool_rotates_round_robin_and_dedupes() {
71        let pool = KeyPool::new("k1", &["k2".into(), " ".into(), "k1".into(), "k3".into()]);
72        assert_eq!(pool.len(), 3);
73        assert_eq!(pool.current(), "k1");
74        assert!(pool.advance());
75        assert_eq!(pool.current(), "k2");
76        assert!(pool.advance());
77        assert_eq!(pool.current(), "k3");
78        assert!(pool.advance());
79        assert_eq!(pool.current(), "k1");
80        assert_eq!(pool.rotations(), 3);
81    }
82
83    #[test]
84    fn single_key_pool_never_advances() {
85        let pool = KeyPool::new("only", &[]);
86        assert!(!pool.advance());
87        assert_eq!(pool.current(), "only");
88        assert_eq!(pool.rotations(), 0);
89    }
90
91    #[test]
92    fn rotation_statuses() {
93        assert!(is_key_rotation_status(401));
94        assert!(is_key_rotation_status(403));
95        assert!(is_key_rotation_status(429));
96        assert!(!is_key_rotation_status(500));
97    }
98}
microclaw_engine/llm/key_pool.rs

microclaw_engine/llm/
key_pool.rs
