use std::collections::HashMap;
use std::sync::{LazyLock, Mutex};
use rust_decimal::Decimal;
use uuid::Uuid;
use crate::core::models::{CostConfidence, CostEvent, EventType, PricingSource};
#[derive(Debug, Clone)]
pub struct DomainRate {
pub cost_usd: Decimal,
pub per: String,
}
static DOMAIN_RATES: LazyLock<Mutex<HashMap<String, DomainRate>>> =
LazyLock::new(|| Mutex::new(HashMap::new()));
const RECORDED_EVENTS_CAP: usize = 10_000;
static RECORDED_EVENTS: LazyLock<Mutex<Vec<CostEvent>>> = LazyLock::new(|| Mutex::new(Vec::new()));
fn push_recorded_event(event: CostEvent) {
let mut guard = RECORDED_EVENTS.lock().unwrap_or_else(|e| {
eprintln!("[dexcost] mutex poisoned, recovering: {}", e);
e.into_inner()
});
guard.push(event);
if guard.len() > RECORDED_EVENTS_CAP {
let drop_n = RECORDED_EVENTS_CAP / 10;
guard.drain(..drop_n);
}
}
#[cfg(test)]
pub(crate) static GLOBAL_HTTP_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
pub fn register_domain_rate(domain: &str, cost_usd: Decimal, per: &str) {
let mut rates = DOMAIN_RATES.lock().unwrap_or_else(|e| {
eprintln!("[dexcost] mutex poisoned, recovering: {}", e);
e.into_inner()
});
rates.insert(
domain.to_string(),
DomainRate {
cost_usd,
per: per.to_string(),
},
);
}
pub fn get_domain_rates() -> HashMap<String, DomainRate> {
DOMAIN_RATES
.lock()
.unwrap_or_else(|e| {
eprintln!("[dexcost] mutex poisoned, recovering: {}", e);
e.into_inner()
})
.clone()
}
pub fn clear_domain_rates() {
DOMAIN_RATES
.lock()
.unwrap_or_else(|e| {
eprintln!("[dexcost] mutex poisoned, recovering: {}", e);
e.into_inner()
})
.clear();
}
fn extract_hostname(url: &str) -> Option<String> {
let without_scheme = if let Some(pos) = url.find("://") {
&url[pos + 3..]
} else {
url
};
let host_with_port = without_scheme.split(['/', '?', '#']).next()?;
let host = if let Some(bracket_end) = host_with_port.find(']') {
&host_with_port[..=bracket_end]
} else if let Some(colon_pos) = host_with_port.rfind(':') {
&host_with_port[..colon_pos]
} else {
host_with_port
};
if host.is_empty() {
None
} else {
Some(host.to_lowercase())
}
}
pub fn resolve_http_cost_event(
url: &str,
task_id: &str,
catalog: Option<&crate::pricing::service_catalog::ServiceCatalog>,
) -> Option<CostEvent> {
let hostname = extract_hostname(url)?;
let rate = {
let rates = DOMAIN_RATES.lock().unwrap_or_else(|e| {
eprintln!("[dexcost] mutex poisoned, recovering: {}", e);
e.into_inner()
});
rates.get(&hostname).cloned()
};
if let Some(rate) = rate {
let mut event = CostEvent::new(task_id, EventType::ExternalCost);
event.event_id = Uuid::new_v4().to_string();
event.cost_usd = rate.cost_usd;
event.cost_confidence = CostConfidence::Computed;
event.pricing_source = Some(PricingSource::Manual);
event.service_name = Some(hostname);
event.details.insert(
"attribution_usage_quantity".to_string(),
serde_json::Value::Number(1.into()),
);
event.details.insert(
"attribution_usage_per".to_string(),
serde_json::Value::String(rate.per),
);
return Some(event);
}
let catalog = catalog?;
let entry = catalog.lookup(url)?;
catalog
.extract_cost(entry, &HashMap::new(), None)
.map(|extraction| {
let mut event = CostEvent::new(task_id, EventType::ExternalCost);
event.event_id = Uuid::new_v4().to_string();
event.cost_usd = extraction.amount;
event.cost_confidence = match extraction.confidence.as_str() {
"exact" => CostConfidence::Exact,
"computed" => CostConfidence::Computed,
"estimated" => CostConfidence::Estimated,
_ => CostConfidence::Unknown,
};
event.pricing_source = Some(match extraction.pricing_source.as_str() {
"user_override" => PricingSource::UserOverride,
_ => PricingSource::ServiceCatalog,
});
if event.pricing_source == Some(PricingSource::ServiceCatalog) {
event.pricing_version = Some(catalog.catalog_version());
}
event.service_name = Some(extraction.service_name.clone());
event.details.insert(
"attribution_usage_quantity".to_string(),
serde_json::Value::String(extraction.usage_quantity.normalize().to_string()),
);
event.details.insert(
"attribution_usage_metric".to_string(),
serde_json::Value::String(extraction.usage_metric.clone()),
);
event
})
}
pub fn record_http_cost(url: &str, task_id: &str) {
record_http_cost_with_catalog(url, task_id, None);
}
pub fn record_http_cost_with_catalog(
url: &str,
task_id: &str,
catalog: Option<&crate::pricing::service_catalog::ServiceCatalog>,
) {
if let Some(event) = resolve_http_cost_event(url, task_id, catalog) {
push_recorded_event(event);
}
}
pub fn get_recorded_events() -> Vec<CostEvent> {
RECORDED_EVENTS
.lock()
.unwrap_or_else(|e| {
eprintln!("[dexcost] mutex poisoned, recovering: {}", e);
e.into_inner()
})
.clone()
}
pub fn clear_recorded_events() {
RECORDED_EVENTS
.lock()
.unwrap_or_else(|e| {
eprintln!("[dexcost] mutex poisoned, recovering: {}", e);
e.into_inner()
})
.clear();
}
#[cfg(test)]
mod tests {
use super::*;
fn d(s: &str) -> Decimal {
s.parse().expect("invalid decimal literal in test")
}
struct StateGuard<'a>(#[allow(dead_code)] tokio::sync::MutexGuard<'a, ()>);
impl<'a> StateGuard<'a> {
fn new() -> Self {
let guard = GLOBAL_HTTP_TEST_LOCK.blocking_lock();
clear_domain_rates();
clear_recorded_events();
StateGuard(guard)
}
}
impl Drop for StateGuard<'_> {
fn drop(&mut self) {
clear_domain_rates();
clear_recorded_events();
}
}
#[test]
fn test_register_and_get_rates() {
let _g = StateGuard::new();
register_domain_rate("api.openai.com", d("0.002"), "call");
register_domain_rate("api.anthropic.com", d("0.005"), "request");
let rates = get_domain_rates();
assert_eq!(rates.len(), 2);
let openai = rates.get("api.openai.com").expect("openai missing");
assert_eq!(openai.cost_usd, d("0.002"));
assert_eq!(openai.per, "call");
let anthropic = rates.get("api.anthropic.com").expect("anthropic missing");
assert_eq!(anthropic.cost_usd, d("0.005"));
assert_eq!(anthropic.per, "request");
}
#[test]
fn test_clear_domain_rates() {
let _g = StateGuard::new();
register_domain_rate("example.com", d("0.001"), "call");
assert_eq!(get_domain_rates().len(), 1);
clear_domain_rates();
assert!(get_domain_rates().is_empty());
}
#[test]
fn test_record_http_cost_when_domain_matches() {
let _g = StateGuard::new();
register_domain_rate("api.openai.com", d("0.003"), "call");
record_http_cost("https://api.openai.com/v1/chat/completions", "task-abc");
let events = get_recorded_events();
assert_eq!(events.len(), 1);
let ev = &events[0];
assert_eq!(ev.task_id, "task-abc");
assert_eq!(ev.cost_usd, d("0.003"));
assert_eq!(ev.event_type, EventType::ExternalCost);
assert_eq!(ev.cost_confidence, CostConfidence::Computed);
assert_eq!(ev.pricing_source, Some(PricingSource::Manual));
assert_eq!(ev.service_name.as_deref(), Some("api.openai.com"));
}
#[test]
fn test_no_record_when_domain_unmatched() {
let _g = StateGuard::new();
register_domain_rate("api.openai.com", d("0.003"), "call");
record_http_cost("https://unknown.example.com/v1/endpoint", "task-xyz");
assert!(get_recorded_events().is_empty());
}
#[test]
fn test_record_http_cost_falls_back_to_catalog() {
let _g = StateGuard::new();
let catalog = crate::pricing::service_catalog::ServiceCatalog::new();
record_http_cost_with_catalog("https://api.exa.ai/search", "task-cat", Some(&catalog));
let events = get_recorded_events();
assert_eq!(events.len(), 1, "catalog fallback should record one event");
let ev = &events[0];
assert_eq!(ev.task_id, "task-cat");
assert!(ev.cost_usd >= Decimal::ZERO);
assert_eq!(
ev.pricing_source,
Some(PricingSource::ServiceCatalog),
"catalog-derived events use the ServiceCatalog pricing source"
);
}
#[test]
fn test_record_http_cost_domain_rate_wins_over_catalog() {
let _g = StateGuard::new();
register_domain_rate("api.exa.ai", d("0.123"), "call");
let catalog = crate::pricing::service_catalog::ServiceCatalog::new();
record_http_cost_with_catalog("https://api.exa.ai/search", "task-pref", Some(&catalog));
let events = get_recorded_events();
assert_eq!(events.len(), 1);
assert_eq!(events[0].cost_usd, d("0.123"));
assert_eq!(events[0].pricing_source, Some(PricingSource::Manual));
}
#[test]
fn test_get_and_clear_recorded_events() {
let _g = StateGuard::new();
register_domain_rate("api.example.com", d("0.001"), "call");
record_http_cost("https://api.example.com/endpoint", "task-1");
record_http_cost("https://api.example.com/endpoint", "task-2");
assert_eq!(get_recorded_events().len(), 2);
clear_recorded_events();
assert!(get_recorded_events().is_empty());
assert!(!get_domain_rates().is_empty());
}
}