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    // A zero entitlement means the plan does not include this bucket (Copilot
120    // Free premium requests), not that it is exhausted.
121    if !unlimited && entitlement == Some(0) {
122        return None;
123    }
124    let remaining = object.get("remaining").and_then(nonnegative_integer);
125    let percent_remaining = object
126        .get("percent_remaining")
127        .and_then(percent)
128        .or_else(|| {
129            entitlement
130                .zip(remaining)
131                .and_then(|(entitlement, remaining)| {
132                    (entitlement > 0).then(|| ((remaining * 100) / entitlement).min(100) as i32)
133                })
134        });
135    (unlimited || percent_remaining.is_some()).then_some(Quota {
136        percent_remaining: percent_remaining.unwrap_or(100),
137        entitlement,
138        remaining,
139        unlimited,
140    })
141}
142
143fn nonnegative_integer(value: &Value) -> Option<u64> {
144    value
145        .as_u64()
146        .or_else(|| value.as_str()?.trim().parse().ok())
147}
148
149fn percent(value: &Value) -> Option<i32> {
150    let value = value
151        .as_f64()
152        .or_else(|| value.as_str()?.trim().parse::<f64>().ok())?;
153    value
154        .is_finite()
155        .then(|| i32::from(crate::format::clamp_pct(value)))
156}
157
158/// GitHub has returned both RFC3339 `quota_reset_date_utc` and date-only
159/// `quota_reset_date`. A date-only reset means midnight UTC of that date.
160fn parse_reset(value: &str) -> Option<DateTime<Utc>> {
161    DateTime::parse_from_rfc3339(value)
162        .ok()
163        .map(|value| value.with_timezone(&Utc))
164        .or_else(|| {
165            NaiveDate::parse_from_str(value, "%Y-%m-%d")
166                .ok()?
167                .and_hms_opt(0, 0, 0)
168                .map(|value| value.and_utc())
169        })
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175
176    #[test]
177    fn parses_all_quota_buckets_and_date_only_reset() {
178        let response: Response = serde_json::from_str(
179            r#"{
180                "copilot_plan":"business",
181                "quota_reset_date":"2026-09-15",
182                "quota_snapshots":{
183                    "premium_interactions":{"entitlement":300,"remaining":45,"percent_remaining":15},
184                    "chat":{"entitlement":"1000","remaining":"250"},
185                    "completions":{"unlimited":true}
186                }
187            }"#,
188        )
189        .unwrap();
190        let snapshot = to_snapshot(response).unwrap();
191        assert_eq!(snapshot.plan, "business");
192        assert_eq!(snapshot.premium.unwrap().used_pct(), 85);
193        assert_eq!(
194            snapshot.chat.unwrap().used_and_entitlement(),
195            Some((750, 1000))
196        );
197        assert!(snapshot.completions.unwrap().unlimited);
198        assert_eq!(
199            snapshot.reset_at.unwrap().to_rfc3339(),
200            "2026-09-15T00:00:00+00:00"
201        );
202    }
203
204    #[test]
205    fn accepts_one_good_bucket_and_rejects_an_empty_schema() {
206        let partial: Response = serde_json::from_str(
207            r#"{"quota_snapshots":{"chat":{"percent_remaining":"not-a-number"},"completions":{"remaining":4,"entitlement":8}}}"#,
208        )
209        .unwrap();
210        let snapshot = to_snapshot(partial).unwrap();
211        assert!(snapshot.chat.is_none());
212        assert_eq!(snapshot.completions.unwrap().used_pct(), 50);
213
214        let empty: Response = serde_json::from_str(r#"{"quota_snapshots":{}}"#).unwrap();
215        assert!(to_snapshot(empty).is_err());
216    }
217
218    #[test]
219    fn skips_a_bucket_the_plan_does_not_include() {
220        // Copilot Free: premium is not part of the plan, reported as 0 of 0 with 0% remaining.
221        let response: Response = serde_json::from_str(
222            r#"{"copilot_plan":"individual","quota_snapshots":{
223                "premium_interactions":{"entitlement":0,"remaining":0,"percent_remaining":0.0,"unlimited":false,"has_quota":false},
224                "chat":{"entitlement":200,"remaining":200,"percent_remaining":100.0,"unlimited":false,"has_quota":true}
225            }}"#,
226        )
227        .unwrap();
228        let snapshot = to_snapshot(response).unwrap();
229        assert!(snapshot.premium.is_none());
230        assert_eq!(snapshot.worst_pct(), 0);
231    }
232}