use std::collections::HashMap;
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use secrecy::SecretString;
use serde::Serialize;
use crate::config::{Config, SelectionStrategy};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CredentialSource {
Platform,
Byok,
}
#[derive(Clone)]
pub struct CredentialLease {
pub id: String,
pub secret: SecretString,
health_key: String,
}
#[cfg(test)]
impl CredentialLease {
pub(crate) fn test(id: &str) -> Self {
Self {
id: id.to_owned(),
secret: SecretString::new(id.to_owned().into_boxed_str()),
health_key: id.to_owned(),
}
}
}
pub struct CredentialPlan {
pub source: CredentialSource,
pub attempts: Vec<CredentialLease>,
pub parked: Vec<CredentialSkip>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CredentialSkip {
pub id: String,
pub reason: &'static str,
}
#[derive(Debug, thiserror::Error)]
pub enum CredentialError {
#[error(
"credential `{id}` for namespace `{namespace}` provider `{provider}` references env var `{env}`, which is unset or empty"
)]
MissingEnv {
namespace: String,
provider: String,
id: String,
env: String,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CredentialState {
Healthy,
Parked,
Probe,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CredentialStatusView<'a> {
Namespace(&'a str),
All,
}
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
pub struct CredentialStatus {
pub namespace: String,
pub provider: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub credential_id: Option<String>,
pub source: &'static str,
pub state: CredentialState,
}
impl Serialize for CredentialState {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
serializer.serialize_str(match self {
Self::Healthy => "healthy",
Self::Parked => "parked",
Self::Probe => "probe",
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Eligibility {
Healthy,
Probe,
Parked,
}
struct PoolEntry {
id: String,
status_id: Option<String>,
secret: SecretString,
weight: u32,
health_key: String,
}
struct Pool {
entries: Vec<PoolEntry>,
cursor: AtomicUsize,
total_weight: u64,
}
impl Pool {
fn new(entries: Vec<PoolEntry>) -> Self {
let total_weight: u64 = entries.iter().map(|e| u64::from(e.weight)).sum();
Self {
entries,
cursor: AtomicUsize::new(0),
total_weight: total_weight.max(1),
}
}
fn order(&self, strategy: SelectionStrategy) -> Vec<usize> {
let n = self.entries.len();
if n == 0 {
return Vec::new();
}
let tick = self.cursor.fetch_add(1, Ordering::Relaxed) as u64;
let start = match strategy {
SelectionStrategy::RoundRobin => (tick % n as u64) as usize,
SelectionStrategy::Weighted => {
let mut offset = tick % self.total_weight;
let mut chosen = n - 1;
for (i, entry) in self.entries.iter().enumerate() {
let weight = u64::from(entry.weight);
if offset < weight {
chosen = i;
break;
}
offset -= weight;
}
chosen
}
};
(0..n).map(|i| (start + i) % n).collect()
}
}
struct CredentialHealth {
threshold: u32,
cooldown: Duration,
circuits: Mutex<HashMap<String, Circuit>>,
}
#[derive(Debug, Clone, Copy, Default)]
struct Circuit {
failures: u32,
parked_at: Option<Instant>,
}
impl CredentialHealth {
fn new(threshold: u32, cooldown: Duration) -> Self {
Self {
threshold: threshold.max(1),
cooldown,
circuits: Mutex::new(HashMap::new()),
}
}
fn classify_at(&self, key: &str, now: Instant) -> Eligibility {
let mut circuits = self.lock();
let Some(circuit) = circuits.get_mut(key) else {
return Eligibility::Healthy;
};
match circuit.parked_at {
None => Eligibility::Healthy,
Some(parked_at) if now.saturating_duration_since(parked_at) >= self.cooldown => {
circuit.parked_at = Some(now);
Eligibility::Probe
}
Some(_) => Eligibility::Parked,
}
}
fn record_success(&self, key: &str) {
self.lock().remove(key);
}
fn record_failure_at(&self, key: &str, now: Instant) {
let mut circuits = self.lock();
let circuit = circuits.entry(key.to_owned()).or_default();
circuit.failures = circuit.failures.saturating_add(1);
if circuit.failures >= self.threshold {
circuit.parked_at = Some(now);
}
}
fn state_at(&self, key: &str, now: Instant) -> CredentialState {
let circuits = self.lock();
let Some(circuit) = circuits.get(key) else {
return CredentialState::Healthy;
};
match circuit.parked_at {
None => CredentialState::Healthy,
Some(parked_at) if now.saturating_duration_since(parked_at) < self.cooldown => {
CredentialState::Parked
}
Some(_) => CredentialState::Probe,
}
}
fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<String, Circuit>> {
self.circuits
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
pub struct Credentials {
pools: HashMap<(String, String), Pool>,
platform_ns: String,
strategy: SelectionStrategy,
health: CredentialHealth,
}
impl Credentials {
pub fn from_env(
config: &Config,
env: &HashMap<String, String>,
) -> Result<Self, CredentialError> {
let mut pools: HashMap<(String, String), Vec<PoolEntry>> = HashMap::new();
for c in &config.credential {
let secret = env.get(&c.env).filter(|v| !v.is_empty()).ok_or_else(|| {
CredentialError::MissingEnv {
namespace: c.namespace.clone(),
provider: c.provider.clone(),
id: c.label().to_string(),
env: c.env.clone(),
}
})?;
pools
.entry((c.namespace.clone(), c.provider.clone()))
.or_default()
.push(PoolEntry {
id: c.label().to_string(),
status_id: c.id.clone(),
secret: SecretString::from(secret.clone()),
weight: c.weight,
health_key: health_key(&c.namespace, &c.provider, c.label()),
});
}
Ok(Self {
pools: pools
.into_iter()
.map(|(key, entries)| (key, Pool::new(entries)))
.collect(),
platform_ns: config.default_namespace().to_string(),
strategy: config.credential_pool.strategy,
health: CredentialHealth::new(
config.credential_pool.failure_threshold,
Duration::from_secs(config.credential_pool.cooldown_seconds),
),
})
}
pub fn plan(&self, config: &Config, namespace: &str, provider: &str) -> Option<CredentialPlan> {
self.plan_at(config, namespace, provider, Instant::now())
}
pub fn plan_pinned(
&self,
config: &Config,
namespace: &str,
provider: &str,
) -> Option<CredentialPlan> {
let (pool, source) = self.resolve_pool(config, namespace, provider)?;
let entry = pool.entries.first()?;
Some(CredentialPlan {
source,
attempts: vec![lease(entry)],
parked: Vec::new(),
})
}
fn plan_at(
&self,
config: &Config,
namespace: &str,
provider: &str,
now: Instant,
) -> Option<CredentialPlan> {
let (pool, source) = self.resolve_pool(config, namespace, provider)?;
let order = pool.order(self.strategy);
let entries = order.iter().map(|&i| &pool.entries[i]);
let mut attempts: Vec<CredentialLease> = Vec::new();
let mut parked: Vec<&PoolEntry> = Vec::new();
for entry in entries {
match self.health.classify_at(&entry.health_key, now) {
Eligibility::Healthy => attempts.push(lease(entry)),
Eligibility::Probe => attempts.insert(0, lease(entry)),
Eligibility::Parked => parked.push(entry),
}
}
let forced = if attempts.is_empty() {
attempts = parked
.first()
.map(|entry| vec![lease(entry)])
.into_iter()
.flatten()
.collect();
attempts.first().map(|lease| lease.id.clone())
} else {
None
};
if attempts.is_empty() {
return None;
}
let parked = parked
.into_iter()
.filter(|entry| forced.as_deref() != Some(entry.id.as_str()))
.map(|entry| CredentialSkip {
id: entry.id.clone(),
reason: "parked",
})
.collect();
Some(CredentialPlan {
source,
attempts,
parked,
})
}
fn resolve_pool<'a>(
&'a self,
config: &Config,
namespace: &str,
provider: &str,
) -> Option<(&'a Pool, CredentialSource)> {
let own = self
.pools
.get(&(namespace.to_string(), provider.to_string()));
match own {
Some(pool) => {
let source = if namespace == self.platform_ns {
CredentialSource::Platform
} else {
CredentialSource::Byok
};
Some((pool, source))
}
None => {
let allow_fallback = config
.namespace(namespace)
.is_some_and(|n| n.allow_platform_fallback);
if !allow_fallback || namespace == self.platform_ns {
return None;
}
let pool = self
.pools
.get(&(self.platform_ns.clone(), provider.to_string()))?;
Some((pool, CredentialSource::Platform))
}
}
}
pub fn record_success(&self, lease: &CredentialLease) {
self.health.record_success(&lease.health_key);
}
pub fn record_failure(&self, lease: &CredentialLease) {
self.health
.record_failure_at(&lease.health_key, Instant::now());
}
pub fn is_present(&self, config: &Config, namespace: &str, provider: &str) -> bool {
if self
.pools
.contains_key(&(namespace.to_string(), provider.to_string()))
{
return true;
}
namespace != self.platform_ns
&& config
.namespace(namespace)
.is_some_and(|n| n.allow_platform_fallback)
&& self
.pools
.contains_key(&(self.platform_ns.clone(), provider.to_string()))
}
pub fn status(&self, config: &Config, view: CredentialStatusView<'_>) -> Vec<CredentialStatus> {
self.status_at(config, view, Instant::now())
}
fn status_at(
&self,
config: &Config,
view: CredentialStatusView<'_>,
now: Instant,
) -> Vec<CredentialStatus> {
let mut statuses = Vec::new();
let mut add_pool =
|owner: &str, provider: &str, pool: &Pool, source: &'static str, hide_default_id| {
for entry in &pool.entries {
statuses.push(CredentialStatus {
namespace: owner.to_owned(),
provider: provider.to_owned(),
credential_id: if hide_default_id {
entry.status_id.clone()
} else {
Some(entry.id.clone())
},
source,
state: self.health.state_at(&entry.health_key, now),
});
}
};
match view {
CredentialStatusView::All => {
for ((namespace, provider), pool) in &self.pools {
add_pool(
namespace,
provider,
pool,
if namespace == &self.platform_ns {
"platform"
} else {
"byok"
},
false,
);
}
}
CredentialStatusView::Namespace(namespace) => {
let allow_platform_fallback = namespace != self.platform_ns
&& config
.namespace(namespace)
.is_some_and(|n| n.allow_platform_fallback);
for ((owner, provider), pool) in &self.pools {
if owner == namespace {
add_pool(
owner,
provider,
pool,
if owner == &self.platform_ns {
"platform"
} else {
"byok"
},
false,
);
}
}
if allow_platform_fallback {
for ((owner, provider), pool) in &self.pools {
if owner == &self.platform_ns
&& !self
.pools
.contains_key(&(namespace.to_owned(), provider.clone()))
{
add_pool(owner, provider, pool, "platform", true);
}
}
}
}
}
statuses.sort_by(|a, b| {
(&a.namespace, &a.provider, &a.credential_id).cmp(&(
&b.namespace,
&b.provider,
&b.credential_id,
))
});
statuses
}
}
fn lease(entry: &PoolEntry) -> CredentialLease {
CredentialLease {
id: entry.id.clone(),
secret: entry.secret.clone(),
health_key: entry.health_key.clone(),
}
}
fn health_key(namespace: &str, provider: &str, id: &str) -> String {
format!("{namespace}/{provider}/{id}")
}
#[cfg(test)]
mod tests {
use super::*;
fn config(extra: &str) -> Config {
let toml = format!(
r#"
[[namespace]]
id = "platform"
default = true
[[namespace]]
id = "acme"
[[provider]]
id = "openai"
kind = "openai"
base_url = "https://api.openai.com/v1"
[[gateway_key]]
env = "AXOND_INBOUND_KEY"
namespace = "platform"
{extra}
"#
);
Config::from_toml_str(&toml).expect("valid config")
}
fn env(pairs: &[(&str, &str)]) -> HashMap<String, String> {
pairs
.iter()
.map(|(k, v)| ((*k).to_string(), (*v).to_string()))
.collect()
}
const TWO_PLATFORM_KEYS: &str = r#"
[[credential]]
namespace = "platform"
provider = "openai"
env = "K1"
id = "openai-a"
[[credential]]
namespace = "platform"
provider = "openai"
env = "K2"
id = "openai-b"
"#;
fn two_key_credentials(cfg: &Config) -> Credentials {
Credentials::from_env(cfg, &env(&[("K1", "sk-a"), ("K2", "sk-b")])).expect("credentials")
}
#[test]
fn fallback_status_hides_default_platform_label_but_keeps_explicit_id() {
let cfg = config(
r#"
[[namespace]]
id = "tenant"
allow_platform_fallback = true
[[credential]]
namespace = "platform"
provider = "openai"
env = "K1"
[[credential]]
namespace = "platform"
provider = "openai"
env = "K2"
id = "public-platform"
"#,
);
let creds = Credentials::from_env(&cfg, &env(&[("K1", "sk-a"), ("K2", "sk-b")]))
.expect("credentials");
let statuses = creds.status(&cfg, CredentialStatusView::Namespace("tenant"));
let body = serde_json::to_value(statuses).expect("status JSON");
let entries = body.as_array().expect("status list");
assert_eq!(entries.len(), 2);
assert_eq!(entries[0]["namespace"], "platform");
assert_eq!(entries[0]["source"], "platform");
assert_eq!(entries[0]["state"], "healthy");
assert!(
!entries[0]
.as_object()
.unwrap()
.contains_key("credential_id")
);
assert_eq!(entries[1]["credential_id"], "public-platform");
}
#[test]
fn round_robin_spreads_requests_across_the_pool() {
let cfg = config(TWO_PLATFORM_KEYS);
let creds = two_key_credentials(&cfg);
let first: Vec<String> = (0..4)
.map(|_| {
creds
.plan(&cfg, "platform", "openai")
.expect("plan")
.attempts[0]
.id
.clone()
})
.collect();
assert_eq!(first, ["openai-a", "openai-b", "openai-a", "openai-b"]);
}
#[test]
fn is_present_is_a_pure_query_that_does_not_perturb_rotation_or_health() {
let cfg = config(TWO_PLATFORM_KEYS);
let creds = two_key_credentials(&cfg);
for _ in 0..5 {
assert!(creds.is_present(&cfg, "platform", "openai"));
}
assert_eq!(
creds
.plan(&cfg, "platform", "openai")
.expect("plan")
.attempts[0]
.id,
"openai-a",
);
let now = Instant::now();
let (head_id, head_key) = {
let plan = creds
.plan_at(&cfg, "platform", "openai", now)
.expect("plan");
let head = &plan.attempts[0];
(head.id.clone(), head.health_key.clone())
};
creds.health.record_failure_at(&head_key, now);
creds.health.record_failure_at(&head_key, now);
let after_cooldown = now + Duration::from_secs(31);
for _ in 0..5 {
assert!(creds.is_present(&cfg, "platform", "openai"));
}
let plan = creds
.plan_at(&cfg, "platform", "openai", after_cooldown)
.expect("plan");
assert_eq!(
plan.attempts[0].id, head_id,
"the recovery probe must not be consumed by presence checks",
);
}
#[test]
fn plan_lists_every_credential_so_a_rate_limit_falls_to_the_next() {
let cfg = config(TWO_PLATFORM_KEYS);
let creds = two_key_credentials(&cfg);
let plan = creds.plan(&cfg, "platform", "openai").expect("plan");
let ids: Vec<&str> = plan.attempts.iter().map(|a| a.id.as_str()).collect();
assert_eq!(ids, ["openai-a", "openai-b"]);
assert_eq!(plan.source, CredentialSource::Platform);
}
#[test]
fn weighted_selection_follows_configured_shares() {
let cfg = config(
r#"
[credential_pool]
strategy = "weighted"
[[credential]]
namespace = "platform"
provider = "openai"
env = "K1"
id = "openai-a"
weight = 3
[[credential]]
namespace = "platform"
provider = "openai"
env = "K2"
id = "openai-b"
weight = 1
"#,
);
let creds = two_key_credentials(&cfg);
let firsts: Vec<String> = (0..8)
.map(|_| {
creds
.plan(&cfg, "platform", "openai")
.expect("plan")
.attempts[0]
.id
.clone()
})
.collect();
assert_eq!(firsts.iter().filter(|id| *id == "openai-a").count(), 6);
assert_eq!(firsts.iter().filter(|id| *id == "openai-b").count(), 2);
}
#[test]
fn repeated_failures_park_one_credential_and_a_probe_recovers_it() {
let cfg = config(TWO_PLATFORM_KEYS);
let creds = two_key_credentials(&cfg);
let now = Instant::now();
let plan = creds
.plan_at(&cfg, "platform", "openai", now)
.expect("plan");
let bad = &plan.attempts[0];
assert_eq!(bad.id, "openai-a");
creds.health.record_failure_at(&bad.health_key, now);
creds.health.record_failure_at(&bad.health_key, now);
for _ in 0..4 {
let plan = creds
.plan_at(&cfg, "platform", "openai", now)
.expect("plan");
let ids: Vec<&str> = plan.attempts.iter().map(|a| a.id.as_str()).collect();
assert_eq!(ids, ["openai-b"], "parked credential must be skipped");
assert_eq!(
plan.parked
.iter()
.map(|entry| entry.id.as_str())
.collect::<Vec<_>>(),
["openai-a"],
);
}
let after_cooldown = now + Duration::from_secs(31);
let plan = creds
.plan_at(&cfg, "platform", "openai", after_cooldown)
.expect("plan");
let ids: Vec<&str> = plan.attempts.iter().map(|a| a.id.as_str()).collect();
assert_eq!(ids[0], "openai-a", "the probe leads the plan");
assert!(plan.parked.is_empty());
let next = creds
.plan_at(&cfg, "platform", "openai", after_cooldown)
.expect("plan");
let ids: Vec<&str> = next.attempts.iter().map(|a| a.id.as_str()).collect();
assert_eq!(
ids,
["openai-b"],
"the probe is single-shot: a concurrent request must not retry the parked key"
);
let probed = plan
.attempts
.iter()
.find(|a| a.id == "openai-a")
.expect("probe lease");
creds.record_success(probed);
assert!(
creds
.status_at(&cfg, CredentialStatusView::All, Instant::now())
.iter()
.all(|status| status.state == CredentialState::Healthy)
);
}
#[test]
fn status_reports_probe_and_is_pure() {
let cfg = config(TWO_PLATFORM_KEYS);
let creds = two_key_credentials(&cfg);
let now = Instant::now();
let plan = creds
.plan_at(&cfg, "platform", "openai", now)
.expect("plan");
let head = &plan.attempts[0];
creds.health.record_failure_at(&head.health_key, now);
creds.health.record_failure_at(&head.health_key, now);
let after_cooldown = now + Duration::from_secs(31);
let statuses = creds.status_at(&cfg, CredentialStatusView::All, after_cooldown);
assert_eq!(
statuses
.iter()
.find(|status| status.credential_id.as_deref() == Some(head.id.as_str()))
.expect("head status")
.state,
CredentialState::Probe
);
for _ in 0..5 {
let _ = creds.status_at(&cfg, CredentialStatusView::All, after_cooldown);
}
assert_eq!(
creds
.plan_at(&cfg, "platform", "openai", after_cooldown)
.expect("plan")
.attempts[0]
.id,
head.id,
"status reads must not consume the half-open probe"
);
}
#[test]
fn a_fully_parked_pool_still_serves_a_forced_probe() {
let cfg = config(TWO_PLATFORM_KEYS);
let creds = two_key_credentials(&cfg);
let now = Instant::now();
for key in ["platform/openai/openai-a", "platform/openai/openai-b"] {
creds.health.record_failure_at(key, now);
creds.health.record_failure_at(key, now);
}
let plan = creds
.plan_at(&cfg, "platform", "openai", now)
.expect("a parked pool still yields one attempt");
assert_eq!(plan.attempts.len(), 1);
}
#[test]
fn byok_namespace_uses_its_own_pool_and_never_borrows_by_default() {
let cfg = config(
r#"
[[credential]]
namespace = "platform"
provider = "openai"
env = "K1"
[[credential]]
namespace = "acme"
provider = "openai"
env = "K2"
"#,
);
let creds = two_key_credentials(&cfg);
let plan = creds.plan(&cfg, "acme", "openai").expect("plan");
assert_eq!(plan.source, CredentialSource::Byok);
let ids: Vec<&str> = plan.attempts.iter().map(|a| a.id.as_str()).collect();
assert_eq!(ids, ["K2"], "BYOK pool must not include platform keys");
let no_key = config(
r#"
[[credential]]
namespace = "platform"
provider = "openai"
env = "K1"
"#,
);
let creds = Credentials::from_env(&no_key, &env(&[("K1", "sk-a")])).expect("credentials");
assert!(creds.plan(&no_key, "acme", "openai").is_none());
}
#[test]
fn platform_fallback_yields_the_whole_platform_pool_attributed_to_platform() {
let cfg = Config::from_toml_str(&format!(
r#"
[[namespace]]
id = "platform"
default = true
[[namespace]]
id = "acme"
allow_platform_fallback = true
[[provider]]
id = "openai"
kind = "openai"
base_url = "https://api.openai.com/v1"
[[gateway_key]]
env = "AXOND_INBOUND_KEY"
namespace = "platform"
{TWO_PLATFORM_KEYS}
"#
))
.expect("valid config");
let creds = two_key_credentials(&cfg);
let plan = creds.plan(&cfg, "acme", "openai").expect("plan");
assert_eq!(plan.source, CredentialSource::Platform);
assert_eq!(plan.attempts.len(), 2);
}
#[test]
fn a_dangling_credential_reference_fails_at_boot() {
let cfg = config(TWO_PLATFORM_KEYS);
let Err(err) = Credentials::from_env(&cfg, &env(&[("K1", "sk-a"), ("K2", "")])) else {
panic!("an unset credential env var must refuse to boot");
};
assert!(matches!(err, CredentialError::MissingEnv { .. }), "{err:?}");
}
}