Skip to main content

ai_usagebar/copilot/
types.rs

1//! Defensive wire parsing for GitHub Copilot's VS Code quota response.
2
3use chrono::{DateTime, NaiveDate, Utc};
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6
7use crate::error::{AppError, Result};
8
9#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
10pub struct Quota {
11    pub percent_remaining: i32,
12    pub entitlement: Option<u64>,
13    pub remaining: Option<u64>,
14    pub unlimited: bool,
15}
16
17impl Quota {
18    pub fn used_pct(&self) -> i32 {
19        if self.unlimited {
20            0
21        } else {
22            100 - self.percent_remaining
23        }
24    }
25
26    pub fn used_and_entitlement(&self) -> Option<(u64, u64)> {
27        Some((
28            self.entitlement?.saturating_sub(self.remaining?),
29            self.entitlement?,
30        ))
31    }
32}
33
34/// Normalized and cacheable fields only. Raw GitHub responses can include
35/// account metadata, which is neither needed for display nor retained.
36#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
37pub struct Snapshot {
38    pub plan: String,
39    pub premium: Option<Quota>,
40    pub chat: Option<Quota>,
41    pub completions: Option<Quota>,
42    pub reset_at: Option<DateTime<Utc>>,
43}
44
45impl Snapshot {
46    pub fn quotas(&self) -> impl Iterator<Item = (&'static str, &Quota)> {
47        [
48            ("Premium requests", self.premium.as_ref()),
49            ("Chat", self.chat.as_ref()),
50            ("Completions", self.completions.as_ref()),
51        ]
52        .into_iter()
53        .filter_map(|(label, quota)| quota.map(|quota| (label, quota)))
54    }
55
56    pub fn worst_pct(&self) -> i32 {
57        self.quotas()
58            .map(|(_, quota)| quota.used_pct())
59            .max()
60            .unwrap_or(0)
61    }
62}
63
64#[derive(Debug, Default, Deserialize)]
65#[serde(default)]
66pub struct Response {
67    pub copilot_plan: Option<String>,
68    pub quota_reset_date: Option<String>,
69    pub quota_reset_date_utc: Option<String>,
70    pub quota_snapshots: Option<QuotaSnapshots>,
71}
72
73#[derive(Debug, Default, Deserialize)]
74#[serde(default)]
75pub struct QuotaSnapshots {
76    pub premium_interactions: Option<Value>,
77    pub chat: Option<Value>,
78    pub completions: Option<Value>,
79}
80
81pub fn to_snapshot(response: Response) -> Result<Snapshot> {
82    let quotas = response.quota_snapshots.unwrap_or_default();
83    let premium = parse_quota(quotas.premium_interactions.as_ref());
84    let chat = parse_quota(quotas.chat.as_ref());
85    let completions = parse_quota(quotas.completions.as_ref());
86    if premium.is_none() && chat.is_none() && completions.is_none() {
87        return Err(AppError::Schema(
88            "GitHub Copilot response contains no usable quota snapshots".into(),
89        ));
90    }
91    let plan = response
92        .copilot_plan
93        .map(|value| value.trim().to_string())
94        .filter(|value| !value.is_empty())
95        .unwrap_or_else(|| "GitHub Copilot".to_string());
96    let reset_at = response
97        .quota_reset_date_utc
98        .as_deref()
99        .and_then(parse_reset)
100        .or_else(|| response.quota_reset_date.as_deref().and_then(parse_reset));
101    Ok(Snapshot {
102        plan,
103        premium,
104        chat,
105        completions,
106        reset_at,
107    })
108}
109
110/// One malformed optional bucket must not hide usable data in another. The
111/// caller still rejects a response where every known bucket is unusable.
112fn parse_quota(value: Option<&Value>) -> Option<Quota> {
113    let object = value?.as_object()?;
114    let unlimited = object
115        .get("unlimited")
116        .and_then(Value::as_bool)
117        .unwrap_or(false);
118    let entitlement = object.get("entitlement").and_then(nonnegative_integer);
119    let remaining = object.get("remaining").and_then(nonnegative_integer);
120    let percent_remaining = object
121        .get("percent_remaining")
122        .and_then(percent)
123        .or_else(|| {
124            entitlement
125                .zip(remaining)
126                .and_then(|(entitlement, remaining)| {
127                    (entitlement > 0).then(|| ((remaining * 100) / entitlement).min(100) as i32)
128                })
129        });
130    (unlimited || percent_remaining.is_some()).then_some(Quota {
131        percent_remaining: percent_remaining.unwrap_or(100),
132        entitlement,
133        remaining,
134        unlimited,
135    })
136}
137
138fn nonnegative_integer(value: &Value) -> Option<u64> {
139    value
140        .as_u64()
141        .or_else(|| value.as_str()?.trim().parse().ok())
142}
143
144fn percent(value: &Value) -> Option<i32> {
145    let value = value
146        .as_f64()
147        .or_else(|| value.as_str()?.trim().parse::<f64>().ok())?;
148    value
149        .is_finite()
150        .then(|| value.round().clamp(0.0, 100.0) as i32)
151}
152
153/// GitHub has returned both RFC3339 `quota_reset_date_utc` and date-only
154/// `quota_reset_date`. A date-only reset means midnight UTC of that date.
155fn parse_reset(value: &str) -> Option<DateTime<Utc>> {
156    DateTime::parse_from_rfc3339(value)
157        .ok()
158        .map(|value| value.with_timezone(&Utc))
159        .or_else(|| {
160            NaiveDate::parse_from_str(value, "%Y-%m-%d")
161                .ok()?
162                .and_hms_opt(0, 0, 0)
163                .map(|value| value.and_utc())
164        })
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170
171    #[test]
172    fn parses_all_quota_buckets_and_date_only_reset() {
173        let response: Response = serde_json::from_str(
174            r#"{
175                "copilot_plan":"business",
176                "quota_reset_date":"2026-09-15",
177                "quota_snapshots":{
178                    "premium_interactions":{"entitlement":300,"remaining":45,"percent_remaining":15},
179                    "chat":{"entitlement":"1000","remaining":"250"},
180                    "completions":{"unlimited":true}
181                }
182            }"#,
183        )
184        .unwrap();
185        let snapshot = to_snapshot(response).unwrap();
186        assert_eq!(snapshot.plan, "business");
187        assert_eq!(snapshot.premium.unwrap().used_pct(), 85);
188        assert_eq!(
189            snapshot.chat.unwrap().used_and_entitlement(),
190            Some((750, 1000))
191        );
192        assert!(snapshot.completions.unwrap().unlimited);
193        assert_eq!(
194            snapshot.reset_at.unwrap().to_rfc3339(),
195            "2026-09-15T00:00:00+00:00"
196        );
197    }
198
199    #[test]
200    fn accepts_one_good_bucket_and_rejects_an_empty_schema() {
201        let partial: Response = serde_json::from_str(
202            r#"{"quota_snapshots":{"chat":{"percent_remaining":"not-a-number"},"completions":{"remaining":4,"entitlement":8}}}"#,
203        )
204        .unwrap();
205        let snapshot = to_snapshot(partial).unwrap();
206        assert!(snapshot.chat.is_none());
207        assert_eq!(snapshot.completions.unwrap().used_pct(), 50);
208
209        let empty: Response = serde_json::from_str(r#"{"quota_snapshots":{}}"#).unwrap();
210        assert!(to_snapshot(empty).is_err());
211    }
212}