alex_core/
amp_usage.rs

1//! Amp Free / credits usage parsing (CLI display text + API envelope).
2//! Shape matches CodexBar's AmpUsageParser so menu limits stay consistent.
3
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6
7#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
8pub struct AmpWorkspaceBalance {
9    pub name: String,
10    pub remaining: f64,
11}
12
13#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
14pub struct AmpUsageSnapshot {
15    pub free_quota: Option<f64>,
16    pub free_used: Option<f64>,
17    pub free_remaining: Option<f64>,
18    pub hourly_replenishment: Option<f64>,
19    pub window_hours: Option<f64>,
20    pub individual_credits: Option<f64>,
21    pub workspace_balances: Vec<AmpWorkspaceBalance>,
22    pub account_email: Option<String>,
23    pub account_organization: Option<String>,
24    pub display_text: String,
25}
26
27#[derive(Debug, Deserialize)]
28struct ApiEnvelope {
29    ok: bool,
30    result: Option<ApiResult>,
31    error: Option<ApiError>,
32}
33
34#[derive(Debug, Deserialize)]
35struct ApiResult {
36    #[serde(rename = "displayText")]
37    display_text: Option<String>,
38}
39
40#[derive(Debug, Deserialize)]
41struct ApiError {
42    code: Option<String>,
43    message: Option<String>,
44}
45
46/// Parse Amp API JSON: `{ "ok": true, "result": { "displayText": "..." } }`.
47pub fn parse_usage_api_response(body: &str) -> Result<AmpUsageSnapshot, String> {
48    let env: ApiEnvelope =
49        serde_json::from_str(body).map_err(|e| format!("invalid Amp usage API JSON: {e}"))?;
50    if !env.ok {
51        if env.error.as_ref().and_then(|e| e.code.as_deref()) == Some("auth-required") {
52            return Err("amp access token is invalid or expired".into());
53        }
54        return Err(env
55            .error
56            .and_then(|e| e.message)
57            .unwrap_or_else(|| "Amp usage API returned an error".into()));
58    }
59    let text = env
60        .result
61        .and_then(|r| r.display_text)
62        .filter(|s| !s.trim().is_empty())
63        .ok_or_else(|| "missing Amp usage display text".to_string())?;
64    parse_usage_display_text(&text)
65}
66
67/// Parse `amp usage` / API display text.
68pub fn parse_usage_display_text(display_text: &str) -> Result<AmpUsageSnapshot, String> {
69    let text = strip_ansi(display_text);
70    let mut email = None;
71    let mut org = None;
72    for line in text.lines() {
73        let line = line.trim();
74        if let Some(rest) = line.strip_prefix("Signed in as ") {
75            if let Some((e, o)) = rest.split_once(" (") {
76                email = Some(e.trim().to_string());
77                org = Some(o.trim_end_matches(')').trim().to_string());
78            } else {
79                email = Some(rest.trim().to_string());
80            }
81            break;
82        }
83    }
84    if email.is_none() && looks_signed_out(&text) {
85        return Err("not logged in to Amp".into());
86    }
87
88    let free = parse_free_line(&text);
89    let individual = parse_individual_credits(&text);
90    let workspace_balances = parse_workspaces(&text);
91
92    if free.is_none() && individual.is_none() && workspace_balances.is_empty() {
93        return Err("missing Amp usage data".into());
94    }
95
96    let (free_quota, free_remaining, free_used, hourly, window_hours) = match free {
97        Some((remaining, quota, hourly)) => {
98            let used = (quota - remaining).max(0.0);
99            let window = if hourly > 0.0 {
100                Some((quota / hourly).max(1.0).round())
101            } else {
102                None
103            };
104            (
105                Some(quota),
106                Some(remaining),
107                Some(used),
108                Some(hourly),
109                window,
110            )
111        }
112        None => (None, None, None, None, None),
113    };
114
115    Ok(AmpUsageSnapshot {
116        free_quota,
117        free_used,
118        free_remaining,
119        hourly_replenishment: hourly,
120        window_hours,
121        individual_credits: individual,
122        workspace_balances,
123        account_email: email,
124        account_organization: org,
125        display_text: text.trim().to_string(),
126    })
127}
128
129/// Convert a snapshot into the admin `/limits` provider entry shape.
130pub fn usage_to_limits_entry(snap: &AmpUsageSnapshot, plan: Option<&str>) -> Value {
131    let mut windows = Vec::new();
132    if let (Some(quota), Some(used)) = (snap.free_quota, snap.free_used) {
133        let used_pct = if quota > 0.0 {
134            ((used / quota) * 100.0).min(100.0)
135        } else {
136            0.0
137        };
138        let mut w = serde_json::json!({
139            "window": "free",
140            "used_pct": used_pct,
141            "quota_usd": quota,
142            "used_usd": used,
143            "remaining_usd": snap.free_remaining,
144        });
145        if let Some(h) = snap.hourly_replenishment {
146            w["hourly_replenishment_usd"] = serde_json::json!(h);
147            if h > 0.0 && used > 0.0 {
148                let hours = used / h;
149                let resets_at_s = (now_s() as f64 + hours * 3600.0) as i64;
150                w["resets_at_s"] = serde_json::json!(resets_at_s);
151            }
152        }
153        windows.push(w);
154    }
155    if let Some(credits) = snap.individual_credits {
156        // Remaining paid balance — surface as a credits window for the menu bar.
157        windows.push(serde_json::json!({
158            "window": "credits",
159            "used_pct": 0.0,
160            "remaining_usd": credits,
161        }));
162    }
163    for ws in &snap.workspace_balances {
164        windows.push(serde_json::json!({
165            "window": format!("ws:{}", ws.name),
166            "used_pct": 0.0,
167            "remaining_usd": ws.remaining,
168        }));
169    }
170
171    let mut entry = serde_json::json!({
172        "provider": "amp",
173        "source": "amp usage API",
174        "windows": windows,
175        "individual_credits_usd": snap.individual_credits,
176        "workspace_balances": snap.workspace_balances,
177        "account_email": snap.account_email,
178        "display_text": snap.display_text,
179    });
180    if let Some(p) = plan {
181        entry["plan"] = serde_json::json!(p);
182    } else if let Some(email) = &snap.account_email {
183        entry["plan"] = serde_json::json!(email);
184    }
185    entry
186}
187
188fn now_s() -> i64 {
189    std::time::SystemTime::now()
190        .duration_since(std::time::UNIX_EPOCH)
191        .map(|d| d.as_secs() as i64)
192        .unwrap_or(0)
193}
194
195fn parse_free_line(text: &str) -> Option<(f64, f64, f64)> {
196    // Amp Free: $8.00 / $10.00 remaining (replenishes +$0.42 / hour)
197    for line in text.lines() {
198        let line = line.trim();
199        let Some(rest) = line.strip_prefix("Amp Free:") else {
200            continue;
201        };
202        let rest = rest.trim();
203        let (left, right) = rest.split_once('/')?;
204        let remaining = parse_money(left)?;
205        let right = right.trim();
206        let (quota_part, after) = right
207            .split_once(" remaining")
208            .map(|(a, b)| (a.trim(), b))
209            .unwrap_or((right, ""));
210        let quota = parse_money(quota_part)?;
211        let mut hourly = 0.0;
212        if let Some(idx) = after.to_lowercase().find("replenishes") {
213            let slice = &after[idx..];
214            let num: String = slice
215                .chars()
216                .skip_while(|c| !c.is_ascii_digit() && *c != '.')
217                .take_while(|c| c.is_ascii_digit() || *c == '.' || *c == ',')
218                .collect();
219            hourly = parse_money(&num).unwrap_or(0.0);
220        }
221        return Some((remaining, quota, hourly));
222    }
223    None
224}
225
226fn parse_individual_credits(text: &str) -> Option<f64> {
227    for line in text.lines() {
228        let line = line.trim();
229        let Some(rest) = line.strip_prefix("Individual credits:") else {
230            continue;
231        };
232        let rest = rest.trim();
233        let (amt, _) = rest.split_once(" remaining")?;
234        return parse_money(amt);
235    }
236    None
237}
238
239fn parse_workspaces(text: &str) -> Vec<AmpWorkspaceBalance> {
240    let mut out = Vec::new();
241    for line in text.lines() {
242        let line = line.trim();
243        let Some(rest) = line.strip_prefix("Workspace ") else {
244            continue;
245        };
246        let Some((name, tail)) = rest.split_once(':') else {
247            continue;
248        };
249        let tail = tail.trim();
250        let Some((amt, _)) = tail.split_once(" remaining") else {
251            continue;
252        };
253        if let Some(remaining) = parse_money(amt) {
254            let name = name.trim().to_string();
255            if !name.is_empty() {
256                out.push(AmpWorkspaceBalance { name, remaining });
257            }
258        }
259    }
260    out
261}
262
263fn parse_money(s: &str) -> Option<f64> {
264    let cleaned: String = s
265        .trim()
266        .trim_start_matches('$')
267        .chars()
268        .filter(|c| c.is_ascii_digit() || *c == '.' || *c == '-')
269        .collect();
270    if cleaned.is_empty() {
271        return None;
272    }
273    cleaned.parse().ok()
274}
275
276fn looks_signed_out(text: &str) -> bool {
277    let lower = text.to_lowercase();
278    lower.contains("not logged in")
279        || lower.contains("please log in")
280        || lower.contains("ampcode.com/login")
281}
282
283fn strip_ansi(s: &str) -> String {
284    let mut out = String::with_capacity(s.len());
285    let mut chars = s.chars().peekable();
286    while let Some(c) = chars.next() {
287        if c == '\u{1b}' {
288            if chars.peek() == Some(&'[') {
289                chars.next();
290                for d in chars.by_ref() {
291                    if ('@'..='~').contains(&d) {
292                        break;
293                    }
294                }
295            }
296            continue;
297        }
298        out.push(c);
299    }
300    out
301}
302
303#[cfg(test)]
304mod tests {
305    use super::*;
306
307    #[test]
308    fn parse_credits_only_display() {
309        let text = "Signed in as me@example.com\nIndividual credits: $20.95 remaining (set up automatic top-up) - https://ampcode.com/settings\n";
310        let snap = parse_usage_display_text(text).unwrap();
311        assert_eq!(snap.account_email.as_deref(), Some("me@example.com"));
312        assert_eq!(snap.individual_credits, Some(20.95));
313        assert!(snap.free_quota.is_none());
314    }
315
316    #[test]
317    fn parse_free_and_workspace() {
318        let text = "\
319Signed in as dev@amp.dev (Acme)
320Amp Free: $8.00 / $10.00 remaining (replenishes +$0.42 / hour)
321Individual credits: $1.50 remaining
322Workspace Team A: $100.00 remaining
323";
324        let snap = parse_usage_display_text(text).unwrap();
325        assert_eq!(snap.account_email.as_deref(), Some("dev@amp.dev"));
326        assert_eq!(snap.account_organization.as_deref(), Some("Acme"));
327        assert_eq!(snap.free_remaining, Some(8.0));
328        assert_eq!(snap.free_quota, Some(10.0));
329        assert_eq!(snap.free_used, Some(2.0));
330        assert_eq!(snap.hourly_replenishment, Some(0.42));
331        assert_eq!(snap.individual_credits, Some(1.5));
332        assert_eq!(snap.workspace_balances.len(), 1);
333        assert_eq!(snap.workspace_balances[0].name, "Team A");
334        assert_eq!(snap.workspace_balances[0].remaining, 100.0);
335    }
336
337    #[test]
338    fn parse_api_envelope() {
339        let body = r#"{"ok":true,"result":{"displayText":"Signed in as a@b.com\nIndividual credits: $5.00 remaining"}}"#;
340        let snap = parse_usage_api_response(body).unwrap();
341        assert_eq!(snap.individual_credits, Some(5.0));
342        let entry = usage_to_limits_entry(&snap, None);
343        assert_eq!(entry["provider"], "amp");
344        assert!(entry["windows"]
345            .as_array()
346            .unwrap()
347            .iter()
348            .any(|w| w["window"] == "credits"));
349    }
350
351    #[test]
352    fn free_window_used_pct() {
353        let text = "Signed in as a@b.com\nAmp Free: $2.00 / $10.00 remaining\n";
354        let snap = parse_usage_display_text(text).unwrap();
355        let entry = usage_to_limits_entry(&snap, None);
356        let free = entry["windows"].as_array().unwrap()[0].clone();
357        assert_eq!(free["window"], "free");
358        assert!((free["used_pct"].as_f64().unwrap() - 80.0).abs() < 0.01);
359    }
360}
alex_core/amp_usage.rs

alex_core/
amp_usage.rs
