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        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    /// Check current estimated spend for today.
25    ///
26    /// WARNING: AWS Cost Explorer data is delayed by 8-24 hours. Costs incurred
27    /// in the current session will NOT appear until hours later. This check is
28    /// best-effort and cannot prevent overspend in real-time.
29    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}