usage_monitor_cli/provider/
abacus.rs

1use async_trait::async_trait;
2use chrono::{DateTime, Utc};
3
4use crate::error::SpendPanelError;
5use crate::model::{CreditsSnapshot, PlanInfo, RateWindow, UsageSnapshot};
6use crate::provider::{ProviderContext, ProviderMetadata, UsageProvider};
7
8/// Abacus AI compute-points provider (apps.abacus.ai, browser cookie auth).
9pub struct AbacusProvider {
10    metadata: ProviderMetadata,
11    base_url: Option<String>,
12}
13
14impl AbacusProvider {
15    pub fn new() -> Self {
16        Self {
17            metadata: ProviderMetadata {
18                id: "abacus",
19                name: "Abacus AI",
20                description: "Abacus AI compute-points monitor (browser cookie)",
21                auth_methods: &["cookie", "env"],
22                website: Some("https://abacus.ai"),
23            },
24            base_url: None,
25        }
26    }
27
28    pub fn with_base_url(url: &str) -> Self {
29        let mut p = Self::new();
30        p.base_url = Some(url.to_string());
31        p
32    }
33
34    fn api_base(&self) -> &str {
35        self.base_url.as_deref().unwrap_or("https://apps.abacus.ai")
36    }
37
38    fn clean(raw: &str) -> String {
39        let mut v = raw.trim();
40        if v.len() >= 2
41            && ((v.starts_with('"') && v.ends_with('"'))
42                || (v.starts_with('\'') && v.ends_with('\'')))
43        {
44            v = &v[1..v.len() - 1];
45        }
46        v.trim().to_string()
47    }
48
49    fn resolve_cookie(ctx: &ProviderContext) -> Result<String, SpendPanelError> {
50        for key in ["cookie", "token"] {
51            if let Some(v) = ctx.config.get(key) {
52                let c = Self::clean(v);
53                if !c.is_empty() {
54                    return Ok(c);
55                }
56            }
57        }
58        if let Ok(v) = std::env::var("ABACUS_COOKIE") {
59            let c = Self::clean(&v);
60            if !c.is_empty() {
61                return Ok(c);
62            }
63        }
64        Err(SpendPanelError::AuthFailed(
65            "abacus".into(),
66            "no session cookie in cookie config or ABACUS_COOKIE".into(),
67        ))
68    }
69
70    fn build_client(ctx: &ProviderContext) -> Result<reqwest::Client, SpendPanelError> {
71        reqwest::Client::builder()
72            .timeout(std::time::Duration::from_secs(ctx.timeout_secs))
73            .build()
74            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))
75    }
76
77    async fn get_json(
78        client: &reqwest::Client,
79        url: String,
80        cookie: &str,
81    ) -> Result<serde_json::Value, SpendPanelError> {
82        let resp = client
83            .get(url)
84            .header("Accept", "application/json")
85            .header("Cookie", cookie)
86            .send()
87            .await
88            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
89        let status = resp.status();
90        let body = resp
91            .text()
92            .await
93            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
94        if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
95            return Err(SpendPanelError::AuthFailed(
96                "abacus".into(),
97                format!("session cookie rejected (HTTP {})", status.as_u16()),
98            ));
99        }
100        if !status.is_success() {
101            return Err(SpendPanelError::ProviderError(
102                "abacus".into(),
103                format!("HTTP {}: {}", status, body),
104            ));
105        }
106        serde_json::from_str(&body)
107            .map_err(|e| SpendPanelError::ParseError("abacus".into(), e.to_string()))
108    }
109
110    fn snapshot_from(
111        compute: &serde_json::Value,
112        billing: Option<&serde_json::Value>,
113    ) -> Result<UsageSnapshot, SpendPanelError> {
114        let total = num(compute, "totalComputePoints");
115        let left = num(compute, "computePointsLeft");
116        let (Some(total), Some(left)) = (total, left) else {
117            return Err(SpendPanelError::ParseError(
118                "abacus".into(),
119                "missing totalComputePoints / computePointsLeft".into(),
120            ));
121        };
122        let used = (total - left).max(0.0);
123
124        let mut snapshot = UsageSnapshot::new("abacus");
125        let mut window = RateWindow::new(
126            used.round() as u64,
127            total.round() as u64,
128            "Compute points",
129            0,
130        );
131        let resets = billing
132            .and_then(|b| b.get("nextBillingDate"))
133            .and_then(parse_date);
134        window.resets_at = resets;
135        snapshot.primary_rate_window = Some(window);
136
137        let mut credits = CreditsSnapshot::new(left, "credits");
138        credits.total = Some(total);
139        credits.used = Some(used);
140        credits.renews_at = resets;
141        snapshot.credits = Some(credits);
142
143        if let Some(tier) = billing
144            .and_then(|b| b.get("currentTier"))
145            .and_then(|v| v.as_str())
146            .filter(|s| !s.is_empty())
147        {
148            snapshot.plan = Some(PlanInfo {
149                name: tier.to_string(),
150                tier: None,
151                features: Vec::new(),
152                price: None,
153                currency: None,
154                billing_period: None,
155            });
156        }
157        Ok(snapshot)
158    }
159}
160
161fn num(v: &serde_json::Value, key: &str) -> Option<f64> {
162    v.get(key).and_then(|x| {
163        x.as_f64()
164            .or_else(|| x.as_str().and_then(|s| s.parse().ok()))
165    })
166}
167
168fn parse_date(v: &serde_json::Value) -> Option<DateTime<Utc>> {
169    if let Some(s) = v.as_str() {
170        if let Ok(secs) = s.parse::<i64>() {
171            let secs = if secs > 1_000_000_000_000 {
172                secs / 1000
173            } else {
174                secs
175            };
176            return chrono::TimeZone::timestamp_opt(&Utc, secs, 0).single();
177        }
178        return DateTime::parse_from_rfc3339(s)
179            .ok()
180            .map(|d| d.with_timezone(&Utc));
181    }
182    if let Some(secs) = v.as_i64() {
183        let secs = if secs > 1_000_000_000_000 {
184            secs / 1000
185        } else {
186            secs
187        };
188        return chrono::TimeZone::timestamp_opt(&Utc, secs, 0).single();
189    }
190    None
191}
192
193impl Default for AbacusProvider {
194    fn default() -> Self {
195        Self::new()
196    }
197}
198
199#[async_trait]
200impl UsageProvider for AbacusProvider {
201    fn metadata(&self) -> &ProviderMetadata {
202        &self.metadata
203    }
204
205    fn detect_credentials(&self) -> bool {
206        std::env::var("ABACUS_COOKIE")
207            .map(|v| !v.trim().is_empty())
208            .unwrap_or(false)
209    }
210
211    async fn fetch_usage(&self, ctx: &ProviderContext) -> Result<UsageSnapshot, SpendPanelError> {
212        let cookie = Self::resolve_cookie(ctx)?;
213        let client = Self::build_client(ctx)?;
214        let base = self.api_base().trim_end_matches('/');
215        let compute = Self::get_json(
216            &client,
217            format!("{}/api/_getOrganizationComputePoints", base),
218            &cookie,
219        )
220        .await?;
221        // Billing is optional — plan/reset are best-effort.
222        let billing = Self::get_json(&client, format!("{}/api/_getBillingInfo", base), &cookie)
223            .await
224            .ok();
225        Self::snapshot_from(&compute, billing.as_ref())
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    use pretty_assertions::assert_eq;
233    use wiremock::matchers::{method, path};
234    use wiremock::{Mock, MockServer, ResponseTemplate};
235
236    #[test]
237    fn test_metadata() {
238        assert_eq!(AbacusProvider::new().metadata().id, "abacus");
239    }
240
241    #[test]
242    fn test_snapshot_from_compute_and_billing() {
243        let compute = serde_json::json!({"totalComputePoints": 1000, "computePointsLeft": 400});
244        let billing =
245            serde_json::json!({"currentTier": "PRO", "nextBillingDate": "2026-07-01T00:00:00Z"});
246        let snap = AbacusProvider::snapshot_from(&compute, Some(&billing)).unwrap();
247        let w = snap.primary_rate_window.unwrap();
248        assert_eq!(w.used, Some(600));
249        assert_eq!(w.limit, Some(1000));
250        let c = snap.credits.unwrap();
251        assert_eq!(c.balance, 400.0);
252        assert_eq!(snap.plan.unwrap().name, "PRO");
253    }
254
255    #[test]
256    fn test_snapshot_without_billing() {
257        // Billing is optional — no plan, no reset, but credits still resolve.
258        let compute = serde_json::json!({"totalComputePoints": 800, "computePointsLeft": 800});
259        let snap = AbacusProvider::snapshot_from(&compute, None).unwrap();
260        assert_eq!(snap.primary_rate_window.unwrap().used, Some(0));
261        assert_eq!(snap.credits.unwrap().balance, 800.0);
262        assert!(snap.plan.is_none());
263    }
264
265    #[test]
266    fn test_missing_fields_error() {
267        let compute = serde_json::json!({"foo": 1});
268        assert!(matches!(
269            AbacusProvider::snapshot_from(&compute, None).unwrap_err(),
270            SpendPanelError::ParseError(_, _)
271        ));
272    }
273
274    #[tokio::test]
275    async fn test_fetch_success() {
276        let server = MockServer::start().await;
277        Mock::given(method("GET"))
278            .and(path("/api/_getOrganizationComputePoints"))
279            .respond_with(ResponseTemplate::new(200).set_body_raw(
280                r#"{"totalComputePoints": 500, "computePointsLeft": 200}"#,
281                "application/json",
282            ))
283            .mount(&server)
284            .await;
285        Mock::given(method("GET"))
286            .and(path("/api/_getBillingInfo"))
287            .respond_with(
288                ResponseTemplate::new(200)
289                    .set_body_raw(r#"{"currentTier": "Free"}"#, "application/json"),
290            )
291            .mount(&server)
292            .await;
293        let provider = AbacusProvider::with_base_url(&server.uri());
294        let mut ctx = ProviderContext::new();
295        ctx.config.insert("cookie".into(), "sid=abc".into());
296        let snap = provider.fetch_usage(&ctx).await.unwrap();
297        assert_eq!(snap.primary_rate_window.unwrap().used, Some(300));
298        assert_eq!(snap.plan.unwrap().name, "Free");
299    }
300
301    #[tokio::test]
302    async fn test_fetch_401() {
303        let server = MockServer::start().await;
304        Mock::given(method("GET"))
305            .and(path("/api/_getOrganizationComputePoints"))
306            .respond_with(ResponseTemplate::new(401))
307            .mount(&server)
308            .await;
309        let provider = AbacusProvider::with_base_url(&server.uri());
310        let mut ctx = ProviderContext::new();
311        ctx.config.insert("cookie".into(), "bad".into());
312        assert!(matches!(
313            provider.fetch_usage(&ctx).await.unwrap_err(),
314            SpendPanelError::AuthFailed(_, _)
315        ));
316    }
317}
usage_monitor_cli/provider/abacus.rs

usage_monitor_cli/provider/
abacus.rs
