1use std::collections::HashMap;
35use std::time::{SystemTime, UNIX_EPOCH};
36
37use serde::{Deserialize, Serialize};
38use tokio::sync::RwLock;
39use tracing::{debug, warn};
40
41const RATE_LIMIT_COOLDOWN_SECS: u64 = 60;
43
44fn expand_multi_key(env_var: &str, raw: String) -> Vec<(String, String)> {
63 if raw.contains(',') {
64 raw.split(',')
65 .enumerate()
66 .filter_map(|(i, k)| {
67 let k = k.trim().to_string();
68 if k.is_empty() {
69 None
70 } else {
71 Some((format!("{}[{}]", env_var, i), k))
72 }
73 })
74 .collect()
75 } else {
76 vec![(env_var.to_string(), raw.trim().to_string())]
77 }
78}
79
80#[derive(Debug, Clone, Serialize, Deserialize)]
82pub struct KeyStats {
83 pub env_var: String,
85 pub total_requests: u64,
87 pub successes: u64,
89 pub failures: u64,
91 pub rate_limits: u64,
93 pub total_latency_ms: u64,
95 pub total_input_tokens: u64,
97 pub total_output_tokens: u64,
99 #[serde(skip)]
101 pub active_requests: u64,
102 #[serde(default)]
104 pub last_rate_limit_at: u64,
105 #[serde(default)]
107 pub last_success_at: u64,
108}
109
110impl KeyStats {
111 fn new(env_var: String) -> Self {
112 Self {
113 env_var,
114 total_requests: 0,
115 successes: 0,
116 failures: 0,
117 rate_limits: 0,
118 total_latency_ms: 0,
119 total_input_tokens: 0,
120 total_output_tokens: 0,
121 active_requests: 0,
122 last_rate_limit_at: 0,
123 last_success_at: 0,
124 }
125 }
126
127 pub fn avg_latency_ms(&self) -> f64 {
129 if self.successes == 0 {
130 return 0.0;
131 }
132 self.total_latency_ms as f64 / self.successes as f64
133 }
134
135 pub fn success_rate(&self) -> f64 {
137 let total = self.successes + self.failures;
138 if total == 0 {
139 return 0.5;
140 }
141 self.successes as f64 / total as f64
142 }
143
144 pub fn is_rate_limited(&self) -> bool {
146 if self.last_rate_limit_at == 0 {
147 return false;
148 }
149 let now = now_unix();
150 now.saturating_sub(self.last_rate_limit_at) < RATE_LIMIT_COOLDOWN_SECS
151 }
152
153 pub fn estimated_cost(&self, input_per_mtok: f64, output_per_mtok: f64) -> f64 {
155 let input_cost = (self.total_input_tokens as f64 / 1_000_000.0) * input_per_mtok;
156 let output_cost = (self.total_output_tokens as f64 / 1_000_000.0) * output_per_mtok;
157 input_cost + output_cost
158 }
159}
160
161#[derive(Debug, Clone)]
163pub struct KeyLease {
164 pub api_key: String,
166 pub env_var: String,
168}
169
170struct EndpointKeys {
172 keys: Vec<(String, String)>,
174 stats: HashMap<String, KeyStats>,
176 next_index: usize,
178}
179
180impl EndpointKeys {
181 fn new(env_vars: Vec<String>) -> Self {
182 let mut keys = Vec::new();
183 let mut stats = HashMap::new();
184
185 for env_var in env_vars {
186 if let Some(raw) = car_secrets::resolve_env_or_keychain(&env_var) {
187 for (sub_var, key) in expand_multi_key(&env_var, raw) {
188 stats.insert(sub_var.clone(), KeyStats::new(sub_var.clone()));
189 keys.push((sub_var, key));
190 }
191 }
192 }
193
194 Self {
195 keys,
196 stats,
197 next_index: 0,
198 }
199 }
200
201 fn lease(&mut self) -> Option<KeyLease> {
211 if self.keys.is_empty() {
212 return None;
213 }
214
215 let mut candidates: Vec<(usize, f64)> = Vec::new();
217 let mut all_cold = true;
218
219 for (idx, (ref env_var, _)) in self.keys.iter().enumerate() {
220 let stats = self.stats.get(env_var);
221
222 if let Some(s) = stats {
224 if s.is_rate_limited() {
225 continue;
226 }
227 if s.total_requests > 0 {
228 all_cold = false;
229 }
230 }
231
232 let score = match stats {
238 Some(s) if s.total_requests > 0 => {
239 let total_tokens = s.total_input_tokens + s.total_output_tokens;
240 let completed = s.successes + s.failures + s.rate_limits;
241 if completed > 0 {
242 let avg_tokens_per_req = total_tokens as f64 / completed as f64;
243 let inflight_estimate = s.active_requests as f64 * avg_tokens_per_req;
244 total_tokens as f64 + inflight_estimate
245 } else {
246 s.active_requests as f64 * 1000.0
248 }
249 }
250 _ => 0.0, };
252
253 candidates.push((idx, score));
254 }
255
256 if all_cold && !candidates.is_empty() {
258 let start = self.next_index % candidates.len();
259 let (idx, _) = candidates[start];
260 self.next_index = start + 1;
261 return self.issue_lease(idx);
262 }
263
264 if candidates.is_empty() {
265 let mut best_idx = 0;
267 let mut oldest_rl = u64::MAX;
268 for (idx, (ref env_var, _)) in self.keys.iter().enumerate() {
269 if let Some(stats) = self.stats.get(env_var) {
270 if stats.last_rate_limit_at < oldest_rl {
271 oldest_rl = stats.last_rate_limit_at;
272 best_idx = idx;
273 }
274 }
275 }
276 let env_var = &self.keys[best_idx].0;
277 warn!(env_var = %env_var, "all keys rate-limited, using oldest-cooldown key");
278 return self.issue_lease(best_idx);
279 }
280
281 candidates.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
283 let (best_idx, _score) = candidates[0];
284 self.issue_lease(best_idx)
285 }
286
287 fn issue_lease(&mut self, idx: usize) -> Option<KeyLease> {
289 let (ref env_var, ref key) = self.keys[idx];
290 let env_var = env_var.clone();
291 let api_key = key.clone();
292
293 if let Some(stats) = self.stats.get_mut(&env_var) {
294 stats.active_requests += 1;
295 stats.total_requests += 1;
296 }
297
298 Some(KeyLease { api_key, env_var })
299 }
300
301 fn report_success(
302 &mut self,
303 env_var: &str,
304 latency_ms: u64,
305 input_tokens: u64,
306 output_tokens: u64,
307 ) {
308 if let Some(stats) = self.stats.get_mut(env_var) {
309 stats.successes += 1;
310 stats.active_requests = stats.active_requests.saturating_sub(1);
311 stats.total_latency_ms += latency_ms;
312 stats.total_input_tokens += input_tokens;
313 stats.total_output_tokens += output_tokens;
314 stats.last_success_at = now_unix();
315 }
316 }
317
318 fn release(&mut self, env_var: &str) {
319 if let Some(stats) = self.stats.get_mut(env_var) {
320 stats.active_requests = stats.active_requests.saturating_sub(1);
321 }
322 }
323
324 fn report_failure(&mut self, env_var: &str, is_rate_limit: bool) {
325 self.release(env_var);
326 self.record_failed_attempt(env_var, is_rate_limit);
327 }
328
329 fn record_failed_attempt(&mut self, env_var: &str, is_rate_limit: bool) {
330 if let Some(stats) = self.stats.get_mut(env_var) {
331 if is_rate_limit {
332 stats.rate_limits += 1;
333 stats.last_rate_limit_at = now_unix();
334 } else {
335 stats.failures += 1;
336 }
337 }
338 }
339}
340
341pub struct KeyPool {
343 endpoints: RwLock<HashMap<String, EndpointKeys>>,
345}
346
347impl KeyPool {
348 pub fn new() -> Self {
349 Self {
350 endpoints: RwLock::new(HashMap::new()),
351 }
352 }
353
354 pub async fn register_endpoint(&self, endpoint: &str, env_vars: Vec<String>) {
357 let mut endpoints = self.endpoints.write().await;
358 let entry = endpoints
359 .entry(endpoint.to_string())
360 .or_insert_with(|| EndpointKeys::new(vec![]));
361
362 let existing_vars: std::collections::HashSet<String> =
363 entry.keys.iter().map(|(v, _)| v.clone()).collect();
364
365 let mut new_keys: Vec<(String, String)> = Vec::new();
366
367 for env_var in env_vars {
368 if existing_vars.contains(&env_var) {
372 continue;
373 }
374 if let Some(raw) = car_secrets::resolve_env_or_keychain(&env_var) {
375 for (sub_var, key) in expand_multi_key(&env_var, raw) {
376 if !existing_vars.contains(&sub_var) {
377 new_keys.push((sub_var, key));
378 }
379 }
380 }
381 }
382
383 for (var, key) in new_keys {
384 entry
385 .stats
386 .entry(var.clone())
387 .or_insert_with(|| KeyStats::new(var.clone()));
388 entry.keys.push((var, key));
389 }
390
391 debug!(
392 endpoint = %endpoint,
393 key_count = entry.keys.len(),
394 "registered endpoint keys"
395 );
396 }
397
398 pub async fn lease(&self, endpoint: &str) -> Option<KeyLease> {
400 let mut endpoints = self.endpoints.write().await;
401 endpoints.get_mut(endpoint)?.lease()
402 }
403
404 pub async fn lease_or_env(&self, endpoint: &str, fallback_env: &str) -> Option<KeyLease> {
409 if let Some(lease) = self.lease(endpoint).await {
411 return Some(lease);
412 }
413
414 if car_secrets::resolve_env_or_keychain(fallback_env).is_some() {
417 self.register_endpoint(endpoint, vec![fallback_env.to_string()])
418 .await;
419 return self.lease(endpoint).await;
420 }
421
422 None
423 }
424
425 pub async fn report_success(
427 &self,
428 endpoint: &str,
429 env_var: &str,
430 latency_ms: u64,
431 input_tokens: u64,
432 output_tokens: u64,
433 ) {
434 let mut endpoints = self.endpoints.write().await;
435 if let Some(ep) = endpoints.get_mut(endpoint) {
436 ep.report_success(env_var, latency_ms, input_tokens, output_tokens);
437 }
438 }
439
440 pub async fn report_failure(&self, endpoint: &str, env_var: &str, is_rate_limit: bool) {
442 let mut endpoints = self.endpoints.write().await;
443 if let Some(ep) = endpoints.get_mut(endpoint) {
444 ep.report_failure(env_var, is_rate_limit);
445 }
446 }
447
448 pub async fn record_failed_attempt(&self, endpoint: &str, env_var: &str, is_rate_limit: bool) {
453 let mut endpoints = self.endpoints.write().await;
454 if let Some(ep) = endpoints.get_mut(endpoint) {
455 ep.record_failed_attempt(env_var, is_rate_limit);
456 }
457 }
458
459 pub async fn release(&self, endpoint: &str, env_var: &str) {
463 let mut endpoints = self.endpoints.write().await;
464 if let Some(ep) = endpoints.get_mut(endpoint) {
465 ep.release(env_var);
466 }
467 }
468
469 pub async fn endpoint_stats(&self, endpoint: &str) -> Vec<KeyStats> {
471 let endpoints = self.endpoints.read().await;
472 endpoints
473 .get(endpoint)
474 .map(|ep| ep.stats.values().cloned().collect())
475 .unwrap_or_default()
476 }
477
478 pub async fn all_stats(&self) -> HashMap<String, Vec<KeyStats>> {
480 let endpoints = self.endpoints.read().await;
481 endpoints
482 .iter()
483 .map(|(ep, keys)| (ep.clone(), keys.stats.values().cloned().collect()))
484 .collect()
485 }
486
487 pub async fn total_keys(&self) -> usize {
489 let endpoints = self.endpoints.read().await;
490 endpoints.values().map(|ep| ep.keys.len()).sum()
491 }
492
493 pub async fn available_keys(&self, endpoint: &str) -> usize {
495 let endpoints = self.endpoints.read().await;
496 endpoints
497 .get(endpoint)
498 .map(|ep| {
499 ep.keys
500 .iter()
501 .filter(|(env_var, _)| {
502 ep.stats
503 .get(env_var)
504 .map(|s| !s.is_rate_limited())
505 .unwrap_or(true)
506 })
507 .count()
508 })
509 .unwrap_or(0)
510 }
511
512 pub async fn save_stats(&self, path: &std::path::Path) -> Result<(), std::io::Error> {
514 let endpoints = self.endpoints.read().await;
515 let stats: HashMap<String, Vec<KeyStats>> = endpoints
516 .iter()
517 .map(|(ep, keys)| (ep.clone(), keys.stats.values().cloned().collect()))
518 .collect();
519 let json = serde_json::to_string_pretty(&stats).map_err(std::io::Error::other)?;
520 if let Some(parent) = path.parent() {
521 std::fs::create_dir_all(parent)?;
522 }
523 std::fs::write(path, json)
524 }
525
526 pub async fn load_stats(&self, path: &std::path::Path) -> Result<usize, std::io::Error> {
528 if !path.exists() {
529 return Ok(0);
530 }
531 let json = std::fs::read_to_string(path)?;
532 let saved: HashMap<String, Vec<KeyStats>> = serde_json::from_str(&json)
533 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
534
535 let mut endpoints = self.endpoints.write().await;
536 let mut count = 0;
537 for (endpoint, stats_list) in saved {
538 let ep = endpoints
539 .entry(endpoint)
540 .or_insert_with(|| EndpointKeys::new(vec![]));
541 for stats in stats_list {
542 ep.stats.insert(stats.env_var.clone(), stats);
543 count += 1;
544 }
545 }
546 Ok(count)
547 }
548}
549
550impl Default for KeyPool {
551 fn default() -> Self {
552 Self::new()
553 }
554}
555
556fn now_unix() -> u64 {
557 SystemTime::now()
558 .duration_since(UNIX_EPOCH)
559 .unwrap_or_default()
560 .as_secs()
561}
562
563#[cfg(test)]
564mod tests {
565 use super::*;
566
567 #[tokio::test]
568 async fn single_key_round_trip() {
569 std::env::set_var("TEST_KEY_POOL_1", "sk-test-111");
570
571 let pool = KeyPool::new();
572 pool.register_endpoint("https://api.test.com", vec!["TEST_KEY_POOL_1".into()])
573 .await;
574
575 let lease = pool.lease("https://api.test.com").await.unwrap();
576 assert_eq!(lease.api_key, "sk-test-111");
577 assert_eq!(lease.env_var, "TEST_KEY_POOL_1");
578
579 pool.report_success("https://api.test.com", &lease.env_var, 500, 100, 50)
580 .await;
581
582 let stats = pool.endpoint_stats("https://api.test.com").await;
583 assert_eq!(stats.len(), 1);
584 assert_eq!(stats[0].successes, 1);
585 assert_eq!(stats[0].total_latency_ms, 500);
586
587 std::env::remove_var("TEST_KEY_POOL_1");
588 }
589
590 #[tokio::test]
591 async fn multi_key_cold_start_round_robin() {
592 std::env::set_var("TEST_KEY_POOL_A", "sk-aaa");
593 std::env::set_var("TEST_KEY_POOL_B", "sk-bbb");
594
595 let pool = KeyPool::new();
596 pool.register_endpoint(
597 "https://api.test.com",
598 vec!["TEST_KEY_POOL_A".into(), "TEST_KEY_POOL_B".into()],
599 )
600 .await;
601
602 let l1 = pool.lease("https://api.test.com").await.unwrap();
604 pool.report_success("https://api.test.com", &l1.env_var, 100, 10, 5)
605 .await;
606
607 let l2 = pool.lease("https://api.test.com").await.unwrap();
608 pool.report_success("https://api.test.com", &l2.env_var, 100, 10, 5)
609 .await;
610
611 assert_ne!(l1.env_var, l2.env_var);
613
614 std::env::remove_var("TEST_KEY_POOL_A");
615 std::env::remove_var("TEST_KEY_POOL_B");
616 }
617
618 #[tokio::test]
619 async fn token_aware_prefers_least_used() {
620 std::env::set_var("TEST_KEY_POOL_TA1", "sk-ta1");
621 std::env::set_var("TEST_KEY_POOL_TA2", "sk-ta2");
622
623 let pool = KeyPool::new();
624 pool.register_endpoint(
625 "https://api.test.com",
626 vec!["TEST_KEY_POOL_TA1".into(), "TEST_KEY_POOL_TA2".into()],
627 )
628 .await;
629
630 let l1 = pool.lease("https://api.test.com").await.unwrap();
632 pool.report_success("https://api.test.com", &l1.env_var, 100, 1000, 500)
633 .await;
634
635 let l2 = pool.lease("https://api.test.com").await.unwrap();
636 pool.report_success("https://api.test.com", &l2.env_var, 100, 100, 50)
637 .await;
638
639 let l3 = pool.lease("https://api.test.com").await.unwrap();
642 assert_eq!(
643 l3.env_var, l2.env_var,
644 "should pick the key with fewer tokens"
645 );
646
647 pool.report_success("https://api.test.com", &l3.env_var, 100, 5000, 5000)
649 .await;
650
651 let l4 = pool.lease("https://api.test.com").await.unwrap();
653 assert_eq!(
654 l4.env_var, l1.env_var,
655 "should pick key with fewer tokens after rebalance"
656 );
657
658 std::env::remove_var("TEST_KEY_POOL_TA1");
659 std::env::remove_var("TEST_KEY_POOL_TA2");
660 }
661
662 #[tokio::test]
663 async fn comma_separated_keys() {
664 std::env::set_var("TEST_KEY_POOL_CSV", "sk-one, sk-two, sk-three");
665
666 let pool = KeyPool::new();
667 pool.register_endpoint("https://api.test.com", vec!["TEST_KEY_POOL_CSV".into()])
668 .await;
669
670 assert_eq!(pool.total_keys().await, 3);
671
672 let l1 = pool.lease("https://api.test.com").await.unwrap();
673 assert_eq!(l1.api_key, "sk-one");
674
675 let l2 = pool.lease("https://api.test.com").await.unwrap();
676 assert_eq!(l2.api_key, "sk-two");
677
678 let l3 = pool.lease("https://api.test.com").await.unwrap();
679 assert_eq!(l3.api_key, "sk-three");
680
681 std::env::remove_var("TEST_KEY_POOL_CSV");
682 }
683
684 #[tokio::test]
685 async fn rate_limited_key_skipped() {
686 std::env::set_var("TEST_KEY_POOL_RL1", "sk-rl1");
687 std::env::set_var("TEST_KEY_POOL_RL2", "sk-rl2");
688
689 let pool = KeyPool::new();
690 pool.register_endpoint(
691 "https://api.test.com",
692 vec!["TEST_KEY_POOL_RL1".into(), "TEST_KEY_POOL_RL2".into()],
693 )
694 .await;
695
696 let l1 = pool.lease("https://api.test.com").await.unwrap();
698 pool.report_failure("https://api.test.com", &l1.env_var, true)
699 .await;
700
701 let l2 = pool.lease("https://api.test.com").await.unwrap();
703 assert_ne!(l1.env_var, l2.env_var);
704
705 std::env::remove_var("TEST_KEY_POOL_RL1");
706 std::env::remove_var("TEST_KEY_POOL_RL2");
707 }
708
709 #[test]
719 fn a_single_key_is_trimmed_like_a_pooled_one() {
720 let pasted = expand_multi_key("PROVIDER_API_KEY", " sk-ant-abc123\n".to_string());
721 assert_eq!(
722 pasted,
723 vec![("PROVIDER_API_KEY".to_string(), "sk-ant-abc123".to_string())],
724 "a single key must be trimmed before it becomes an auth header"
725 );
726
727 let pooled = expand_multi_key("PROVIDER_API_KEY", " k1 , k2\n".to_string());
729 assert_eq!(
730 pooled,
731 vec![
732 ("PROVIDER_API_KEY[0]".to_string(), "k1".to_string()),
733 ("PROVIDER_API_KEY[1]".to_string(), "k2".to_string()),
734 ]
735 );
736 }
737
738 #[tokio::test]
739 async fn lease_or_env_fallback() {
740 std::env::set_var("TEST_KEY_POOL_FB", "sk-fallback");
741
742 let pool = KeyPool::new();
743
744 let lease = pool
746 .lease_or_env("https://api.new.com", "TEST_KEY_POOL_FB")
747 .await
748 .unwrap();
749 assert_eq!(lease.api_key, "sk-fallback");
750
751 assert_eq!(pool.total_keys().await, 1);
753
754 std::env::remove_var("TEST_KEY_POOL_FB");
755 }
756
757 #[tokio::test]
758 async fn missing_key_resolution_uses_the_test_file_backend_by_default() {
759 if !crate::run_in_isolated_test_process(
760 "key_pool::tests::missing_key_resolution_uses_the_test_file_backend_by_default",
761 "CAR_KEY_POOL_TEST_FILE_BACKEND_CHILD",
762 ) {
763 return;
764 }
765
766 const MISSING_KEY: &str = "CAR_KEY_POOL_MISSING_TEST_KEY";
767 unsafe {
768 std::env::remove_var(MISSING_KEY);
769 std::env::remove_var("CAR_SECRETS_FILE_DIR");
770 std::env::remove_var("CAR_TEST_NATIVE_KEYCHAIN");
771 std::env::remove_var("CAR_AGENT_ID");
772 std::env::remove_var("CAR_AGENT_TOKEN");
773 }
774 let before = car_secrets::secret_store_activity();
778 assert_eq!(car_secrets::resolve_env_or_keychain(MISSING_KEY), None);
779 let after = car_secrets::secret_store_activity();
780 assert_eq!(after.get_attempts, before.get_attempts + 1);
781 }
782
783 #[tokio::test]
784 async fn no_keys_returns_none() {
785 let pool = KeyPool::new();
786 assert!(pool.lease("https://nonexistent.com").await.is_none());
787 }
788}