Skip to main content

tryaudex_core/
budget.rs

1use crate::error::{AvError, Result};
2
3/// Tracks spending for a session and enforces budget limits.
4pub 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        let config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
12        Ok(Self {
13            limit,
14            ce_client: aws_sdk_costexplorer::Client::new(&config),
15        })
16    }
17
18    /// Check current estimated spend for today.
19    /// NOTE: AWS Cost Explorer has a delay of ~8-24 hours, so this is best-effort.
20    /// For real-time cost tracking, a future version will use CloudWatch billing metrics.
21    pub async fn check_spend(&self) -> Result<f64> {
22        let today = chrono::Utc::now().format("%Y-%m-%d").to_string();
23        let tomorrow = (chrono::Utc::now() + chrono::Duration::days(1))
24            .format("%Y-%m-%d")
25            .to_string();
26
27        let result = self
28            .ce_client
29            .get_cost_and_usage()
30            .time_period(
31                aws_sdk_costexplorer::types::DateInterval::builder()
32                    .start(&today)
33                    .end(&tomorrow)
34                    .build()
35                    .map_err(|e| AvError::Sts(format!("Failed to build date interval: {}", e)))?,
36            )
37            .granularity(aws_sdk_costexplorer::types::Granularity::Daily)
38            .metrics("UnblendedCost")
39            .send()
40            .await
41            .map_err(|e| AvError::Sts(format!("Cost Explorer error: {}", e)))?;
42
43        let spend = result
44            .results_by_time()
45            .first()
46            .and_then(|r| r.total())
47            .and_then(|t| t.get("UnblendedCost"))
48            .and_then(|m| m.amount())
49            .and_then(|a| a.parse::<f64>().ok())
50            .unwrap_or(0.0);
51
52        Ok(spend)
53    }
54
55    pub fn limit(&self) -> f64 {
56        self.limit
57    }
58
59    pub fn is_exceeded(&self, spend: f64) -> bool {
60        spend > self.limit
61    }
62
63    pub fn is_warning(&self, spend: f64) -> bool {
64        spend > self.limit * 0.8
65    }
66}