Skip to main content

car_inference/
key_pool.rs

1//! API key pool — load-balanced multi-key management for remote endpoints.
2//!
3//! Supports multiple API keys per endpoint with:
4//! - **Round-robin** selection with rate-limit awareness
5//! - **Per-key stats** — request counts, errors, 429s, latency
6//! - **Automatic exclusion** — rate-limited keys are skipped until cooldown expires
7//! - **Persistence** — stats survive across calls for informed selection
8//!
9//! ## Resolution order
10//!
11//! Delegates to [`car_secrets::resolve_env_or_keychain`] (the single
12//! source of truth across the workspace). Priority:
13//!
14//! 1. **Process env var** — `std::env::var(name)`. Wins everything.
15//!    Preserves the universal config pattern that works in
16//!    containers, CI, K8s pods, systemd units, and headless Linux.
17//! 2. **The secret store** under service `"car"`, account = env-var name:
18//!    the OS keychain, or on Linux with no Secret Service an owner-only file
19//!    under `~/.car/secrets/`. Never exported into the process env.
20//! 3. **Missing** — endpoint registers with no keys, requests fail
21//!    at lease time with a clear "no API key configured" error.
22//!
23//! Multi-key semantics are preserved across every source: a
24//! comma-separated value in either env or keychain
25//! (`sk-a,sk-b,sk-c`) becomes a load-balanced pool with per-sub-key
26//! rate-limit awareness.
27//!
28//! The `keychain` cargo feature was removed in v0.6.x — keychain
29//! lookup is unconditional now (per CLAUDE.md §1, no feature
30//! flags). Embedders that don't want any keychain interaction
31//! simply don't store keys there; `is_available()` short-circuits
32//! before any DBus / Security.framework round-trip.
33
34use std::collections::HashMap;
35use std::time::{SystemTime, UNIX_EPOCH};
36
37use serde::{Deserialize, Serialize};
38use tokio::sync::RwLock;
39use tracing::{debug, warn};
40
41/// Cooldown period after a 429 before retrying a key (seconds).
42const RATE_LIMIT_COOLDOWN_SECS: u64 = 60;
43
44/// Split a raw key value (which may be comma-separated for multi-key pools)
45/// into `(sub_var_name, key_value)` pairs.
46///
47/// **Every key is trimmed**, single or pooled. The comma-separated branch
48/// always did; the single-key branch passed `raw` through verbatim, so one
49/// stray newline — the normal outcome of pasting a key into a GUI field, or
50/// `echo $KEY | car secrets put` — went into the `x-api-key` /
51/// `Authorization` header as part of the credential and the provider answered
52/// `401 provider rejected the API key`. Nothing upstream trims it either:
53/// CarHost's secrets pane trims the service and key NAMES and stores the value
54/// as typed, and `car_secrets` stores bytes faithfully, which is correct for a
55/// general secret store — an arbitrary secret may legitimately contain
56/// surrounding whitespace.
57///
58/// So the trim belongs here, at the point of use, where the value is known to
59/// be an API key rather than an arbitrary secret. An API key with meaningful
60/// leading or trailing whitespace does not exist; a stored password with it
61/// might, and trimming those at the storage layer would corrupt them.
62fn 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/// Per-key performance statistics.
81#[derive(Debug, Clone, Serialize, Deserialize)]
82pub struct KeyStats {
83    /// Which env var this key came from.
84    pub env_var: String,
85    /// Total requests made with this key.
86    pub total_requests: u64,
87    /// Total successful requests.
88    pub successes: u64,
89    /// Total failed requests (non-429).
90    pub failures: u64,
91    /// Total 429 (rate-limit) responses.
92    pub rate_limits: u64,
93    /// Cumulative latency in ms (for computing average).
94    pub total_latency_ms: u64,
95    /// Total input tokens processed.
96    pub total_input_tokens: u64,
97    /// Total output tokens produced.
98    pub total_output_tokens: u64,
99    /// Currently in-flight requests (not serialized).
100    #[serde(skip)]
101    pub active_requests: u64,
102    /// Unix timestamp of last rate-limit (0 = never).
103    #[serde(default)]
104    pub last_rate_limit_at: u64,
105    /// Unix timestamp of last successful request.
106    #[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    /// Average latency in ms. Returns 0 if no requests completed.
128    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    /// Success rate (0.0 to 1.0). Returns 0.5 prior if no data.
136    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    /// Whether this key is in cooldown from a recent rate-limit.
145    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    /// Estimated cost in USD based on token counts and per-million-token rates.
154    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/// A resolved key ready for use — contains the actual secret and a handle for reporting back.
162#[derive(Debug, Clone)]
163pub struct KeyLease {
164    /// The actual API key value.
165    pub api_key: String,
166    /// The env var it came from (for stats reporting).
167    pub env_var: String,
168}
169
170/// Pool of API keys for a single endpoint.
171struct EndpointKeys {
172    /// Ordered list of (env_var_name, resolved_key_value).
173    keys: Vec<(String, String)>,
174    /// Per-key stats, indexed by env var name.
175    stats: HashMap<String, KeyStats>,
176    /// Round-robin cursor.
177    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    /// Select the best available key using token-aware load balancing.
202    ///
203    /// Strategy: pick the key with the lowest recent token load, breaking ties
204    /// with fewest active in-flight requests. This spreads usage evenly across
205    /// keys, which matters because provider rate limits are typically token-based
206    /// (TPM) rather than request-based.
207    ///
208    /// For cold-start (no stats yet), falls back to round-robin so keys get
209    /// initial data points.
210    fn lease(&mut self) -> Option<KeyLease> {
211        if self.keys.is_empty() {
212            return None;
213        }
214
215        // Collect candidates: non-rate-limited keys with their scores
216        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            // Skip rate-limited keys
223            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            // Score = total tokens + penalty for in-flight and pending requests.
233            // Lower is better. In-flight requests haven't reported tokens yet,
234            // so we estimate them from that key's average. Keys that have been
235            // leased but not yet reported get a request-based penalty to avoid
236            // piling on the same key before results come back.
237            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                        // Leased but no results yet — use request count as proxy
247                        s.active_requests as f64 * 1000.0
248                    }
249                }
250                _ => 0.0, // Cold key — lowest possible score to encourage exploration
251            };
252
253            candidates.push((idx, score));
254        }
255
256        // If all candidates are cold (no stats), use round-robin for even warm-up
257        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            // All keys are rate-limited — pick the one closest to cooldown expiry
266            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        // Pick the candidate with the lowest token score
282        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    /// Issue a lease for a key at the given index, updating stats.
288    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
341/// Thread-safe key pool managing multiple API keys across endpoints.
342pub struct KeyPool {
343    /// Keyed by a canonical endpoint identifier (e.g., `"https://api.openai.com"`).
344    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    /// Register keys for an endpoint. Merges with any existing keys.
355    /// Resolution order: process env → OS keychain (see module docs).
356    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            // Skip if the bare env var name is already registered (single-key
369            // path). For multi-key sources we also dedupe each [N] sub-name
370            // below before insertion.
371            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    /// Lease a key for the given endpoint. Returns None if no keys available.
399    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    /// Lease a key, falling back to a single env var (or keychain) if the
405    /// endpoint isn't registered. This provides backward compatibility with
406    /// the single-key model and triggers the same env→keychain resolution
407    /// order as `register_endpoint`.
408    pub async fn lease_or_env(&self, endpoint: &str, fallback_env: &str) -> Option<KeyLease> {
409        // Try pool first.
410        if let Some(lease) = self.lease(endpoint).await {
411            return Some(lease);
412        }
413
414        // Fallback: probe the same sources `register_endpoint` would. Auto-register
415        // for future calls so subsequent leases don't re-probe.
416        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    /// Report a successful request.
426    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    /// Report a failed request.
441    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    /// Book one failed attempt (a failure or a rate limit) against a key
449    /// whose lease is still open because the caller is about to retry on it.
450    /// Unlike [`Self::report_failure`] it does not end the lease: one lease
451    /// that retries is one in-flight request, however many attempts it makes.
452    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    /// End a lease that neither succeeded nor failed — the caller hung up
460    /// before the provider finished. Frees the in-flight slot without
461    /// booking the request against the key's health.
462    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    /// Get stats for all keys on an endpoint.
470    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    /// Get stats for all endpoints.
479    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    /// Total key count across all endpoints.
488    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    /// Get the number of available (non-rate-limited) keys for an endpoint.
494    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    /// Save key stats to a JSON file.
513    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    /// Load key stats from a JSON file.
527    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        // Cold start: no stats yet, should round-robin
603        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        // Two different keys should have been used
612        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        // Warm up both keys with cold-start round-robin
631        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        // Now both keys have stats. Key 2 has far fewer tokens (150 vs 1500).
640        // Token-aware selection should prefer key 2.
641        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        // Report heavy usage on key 2 to flip the balance
648        pool.report_success("https://api.test.com", &l3.env_var, 100, 5000, 5000)
649            .await;
650
651        // Now key 1 has 1500 total tokens, key 2 has 10150 — should pick key 1
652        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        // First key gets rate-limited
697        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        // Next lease should skip the rate-limited key
702        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    /// **A pasted key with a trailing newline must not reach the provider.**
710    ///
711    /// The comma-separated branch always trimmed; the single-key branch did
712    /// not, so `sk-ant-xxx\n` was sent as the credential verbatim and the
713    /// provider answered `401 provider rejected the API key`. That is the
714    /// normal outcome of pasting into a GUI field or piping `echo $KEY`, and
715    /// nothing upstream trims it: CarHost's secrets pane trims the NAMES and
716    /// stores the value as typed, and the secret store keeps bytes faithfully
717    /// on purpose.
718    #[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        // The pooled branch's existing behaviour is unchanged.
728        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        // No pre-registered endpoint — should fall back to env var
745        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        // Should be auto-registered now
752        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        // car-secrets has a test-process tripwire before SystemSecurityCli can
775        // inspect or spawn against the login Keychain. Reaching `None` here
776        // with both overrides absent proves this dependency build chose files.
777        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}