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