1use crate::error::{AvError, Result};
2
3pub struct BudgetMonitor {
5 limit: f64,
6 ce_client: aws_sdk_costexplorer::Client,
7}
8
9impl BudgetMonitor {
10 pub async fn new(limit: f64) -> Result<Self> {
11 if limit.is_nan() || limit.is_infinite() || limit <= 0.0 {
12 return Err(AvError::InvalidPolicy(format!(
13 "Budget limit must be a positive finite number, got {}",
14 limit
15 )));
16 }
17 let config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
18 Ok(Self {
19 limit,
20 ce_client: aws_sdk_costexplorer::Client::new(&config),
21 })
22 }
23
24 pub async fn check_spend(&self) -> Result<f64> {
30 tracing::warn!(
31 "Budget check uses Cost Explorer which has 8-24h data delay — \
32 recent session costs are not reflected"
33 );
34 let today = chrono::Utc::now().format("%Y-%m-%d").to_string();
35 let tomorrow = (chrono::Utc::now() + chrono::Duration::days(1))
36 .format("%Y-%m-%d")
37 .to_string();
38
39 let result = self
40 .ce_client
41 .get_cost_and_usage()
42 .time_period(
43 aws_sdk_costexplorer::types::DateInterval::builder()
44 .start(&today)
45 .end(&tomorrow)
46 .build()
47 .map_err(|e| AvError::Sts(format!("Failed to build date interval: {}", e)))?,
48 )
49 .granularity(aws_sdk_costexplorer::types::Granularity::Daily)
50 .metrics("UnblendedCost")
51 .send()
52 .await
53 .map_err(|e| AvError::Sts(format!("Cost Explorer error: {}", e)))?;
54
55 let spend = result
56 .results_by_time()
57 .first()
58 .and_then(|r| r.total())
59 .and_then(|t| t.get("UnblendedCost"))
60 .and_then(|m| m.amount())
61 .and_then(|a| a.parse::<f64>().ok())
62 .unwrap_or(0.0);
63
64 Ok(spend)
65 }
66
67 pub fn limit(&self) -> f64 {
68 self.limit
69 }
70
71 pub fn is_exceeded(&self, spend: f64) -> bool {
72 spend > self.limit
73 }
74
75 pub fn is_warning(&self, spend: f64) -> bool {
76 spend > self.limit * 0.8
77 }
78}