Skip to main content

ai_usagebar/muse/
types.rs

1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3
4use crate::error::{AppError, Result};
5use crate::usage::{MuseSnapshot, UsageWindow};
6
7pub const UNAVAILABLE: &str = "Muse Code did not report subscription quota (it may omit both windows while the session is idle); remaining allowance is unknown";
8
9/// Only quota-bearing fields: serde skips api_key, user_email, and every
10/// other account field without retaining them in the parsed record.
11#[derive(Clone, Deserialize, Serialize)]
12pub struct Response {
13    pub is_subs_active: bool,
14    pub subs_tier_name: Option<String>,
15    pub subs_usage: Option<Quota>,
16}
17
18#[derive(Clone, Deserialize, Serialize)]
19pub struct Quota {
20    pub window: Option<Window>,
21    pub weekly: Option<Window>,
22}
23
24#[derive(Clone, Deserialize, Serialize)]
25pub struct Window {
26    pub used_percent: f64,
27    pub resets_at: i64,
28    pub window_duration_mins: Option<i64>,
29}
30
31impl Response {
32    pub fn snapshot(&self, now: DateTime<Utc>) -> Result<MuseSnapshot> {
33        if !self.is_subs_active {
34            return Err(AppError::Credentials(
35                "No active Muse Code subscription on this login".into(),
36            ));
37        }
38        let quota = self
39            .subs_usage
40            .as_ref()
41            .ok_or_else(|| AppError::Other(UNAVAILABLE.into()))?;
42        let window = |raw: &Window, minutes: i64| -> Result<Option<UsageWindow>> {
43            if !raw.used_percent.is_finite()
44                || raw.used_percent < 0.0
45                || !(1..=525600).contains(&minutes)
46            {
47                return Err(AppError::Schema(
48                    "Muse Code quota contains an invalid percentage or duration".into(),
49                ));
50            }
51            let reset = DateTime::from_timestamp(raw.resets_at, 0)
52                .filter(|_| raw.resets_at > 0)
53                .ok_or_else(|| {
54                    AppError::Schema("Muse Code quota contains an invalid reset".into())
55                })?;
56            // An expired observation cannot describe the next window.
57            Ok((reset > now).then_some(UsageWindow {
58                utilization_pct: raw.used_percent.round().min(i32::MAX as f64) as i32,
59                resets_at: Some(reset),
60                window_duration: chrono::Duration::minutes(minutes),
61            }))
62        };
63        let session = quota
64            .window
65            .as_ref()
66            .map(|raw| {
67                let minutes = raw
68                    .window_duration_mins
69                    .ok_or_else(|| AppError::Schema("Muse Code session duration missing".into()))?;
70                window(raw, minutes)
71            })
72            .transpose()?
73            .flatten();
74        let weekly = quota
75            .weekly
76            .as_ref()
77            .map(|raw| window(raw, 7 * 24 * 60))
78            .transpose()?
79            .flatten();
80        if session.is_none() && weekly.is_none() {
81            return Err(AppError::Other(UNAVAILABLE.into()));
82        }
83        Ok(MuseSnapshot {
84            plan: crate::display::sanitize_untrusted_field(
85                self.subs_tier_name.as_deref().unwrap_or("Muse Code"),
86            ),
87            session,
88            weekly,
89        })
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    fn now() -> DateTime<Utc> {
97        DateTime::from_timestamp(1791936000, 0).unwrap()
98    }
99    #[test]
100    fn wire_seconds_and_dynamic_duration_project_to_report_windows() {
101        let response: Response = serde_json::from_value(serde_json::json!({
102            "is_subs_active":true,"subs_tier_name":"Muse Code Power Usage",
103            "api_key":"SecretCanary-9348","user_email":"secret@example.test",
104            "subs_usage":{"window":{"used_percent":38.5,"resets_at":1791942677,"window_duration_mins":300},
105                          "weekly":{"used_percent":34,"resets_at":1792368000}}
106        })).unwrap();
107        let snap = response.snapshot(now()).unwrap();
108        assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 39);
109        assert_eq!(
110            snap.weekly.as_ref().unwrap().window_duration,
111            chrono::Duration::days(7)
112        );
113        assert_eq!(
114            snap.session.unwrap().window_duration,
115            chrono::Duration::hours(5)
116        );
117        let cached = serde_json::to_string(&response).unwrap();
118        assert!(!cached.contains("SecretCanary") && !cached.contains("secret@example"));
119    }
120    #[test]
121    fn absent_quota_and_expired_observations_are_unknown_not_zero() {
122        for usage in [
123            serde_json::Value::Null,
124            serde_json::json!({}),
125            serde_json::json!({
126            "weekly":{"used_percent":34,"resets_at":1791935000}}),
127        ] {
128            let response: Response = serde_json::from_value(
129                serde_json::json!({"is_subs_active":true,"subs_usage":usage}),
130            )
131            .unwrap();
132            assert!(response.snapshot(now()).is_err());
133        }
134    }
135}