usage_monitor_cli/provider/
copilot.rs

1use async_trait::async_trait;
2
3use crate::error::SpendPanelError;
4use crate::model::{PlanInfo, RateWindow, UsageSnapshot};
5use crate::provider::{ProviderContext, ProviderMetadata, UsageProvider};
6
7#[derive(Debug, serde::Deserialize)]
8struct CopilotUsageResponse {
9    #[serde(default)]
10    quota_snapshots: CopilotQuotaSnapshots,
11    #[serde(default)]
12    copilot_plan: Option<String>,
13    #[serde(default)]
14    token_based_billing: bool,
15}
16
17#[derive(Debug, Default, serde::Deserialize)]
18struct CopilotQuotaSnapshots {
19    #[serde(default)]
20    premium_interactions: Option<CopilotQuotaSnapshot>,
21    #[serde(default)]
22    chat: Option<CopilotQuotaSnapshot>,
23}
24
25#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
26struct CopilotQuotaSnapshot {
27    #[serde(default)]
28    entitlement: f64,
29    #[serde(default)]
30    remaining: f64,
31    #[serde(default)]
32    percent_remaining: Option<f64>,
33    #[serde(default)]
34    unlimited: bool,
35}
36
37impl CopilotQuotaSnapshot {
38    /// Percent remaining, derived from entitlement/remaining when absent.
39    fn percent_remaining(&self) -> Option<f64> {
40        if self.unlimited {
41            return Some(100.0);
42        }
43        if let Some(p) = self.percent_remaining {
44            return Some(p);
45        }
46        if self.entitlement > 0.0 {
47            return Some((self.remaining / self.entitlement) * 100.0);
48        }
49        None
50    }
51
52    /// Zero-entitlement placeholder GitHub returns for token-based seats.
53    fn is_placeholder(&self) -> bool {
54        if self.unlimited {
55            return false;
56        }
57        self.entitlement == 0.0 && self.remaining == 0.0
58    }
59
60    fn to_rate_window(&self, label: &str) -> Option<RateWindow> {
61        if self.is_placeholder() {
62            return None;
63        }
64        let percent_remaining = self.percent_remaining()?;
65        let used_percent = (100.0 - percent_remaining).clamp(0.0, 100.0);
66        // entitlement is the denominator (interactions); used = entitlement - remaining.
67        if self.entitlement > 0.0 {
68            let used = (self.entitlement - self.remaining).max(0.0);
69            Some(RateWindow::new(
70                used.round() as u64,
71                self.entitlement.round() as u64,
72                label.to_string(),
73                0,
74            ))
75        } else {
76            // Unlimited or percent-only: synthesize a 0..100 window from the percent.
77            Some(RateWindow::new(
78                used_percent.round() as u64,
79                100,
80                label.to_string(),
81                0,
82            ))
83        }
84    }
85}
86
87/// GitHub Copilot usage provider (GitHub OAuth/PAT token auth).
88pub struct CopilotProvider {
89    metadata: ProviderMetadata,
90    /// Base URL override for tests (defaults to api.github.com).
91    base_url: Option<String>,
92}
93
94impl CopilotProvider {
95    pub fn new() -> Self {
96        Self {
97            metadata: ProviderMetadata {
98                id: "copilot",
99                name: "GitHub Copilot",
100                description: "GitHub Copilot quota monitor",
101                auth_methods: &["token", "api_key", "env"],
102                website: Some("https://github.com/features/copilot"),
103            },
104            base_url: None,
105        }
106    }
107
108    /// Creates a provider with a custom base URL (for tests).
109    pub fn with_base_url(url: &str) -> Self {
110        let mut p = Self::new();
111        p.base_url = Some(url.to_string());
112        p
113    }
114
115    fn api_base(&self) -> &str {
116        self.base_url.as_deref().unwrap_or("https://api.github.com")
117    }
118
119    fn clean(raw: &str) -> String {
120        let mut value = raw.trim();
121        if value.len() >= 2
122            && ((value.starts_with('"') && value.ends_with('"'))
123                || (value.starts_with('\'') && value.ends_with('\'')))
124        {
125            value = &value[1..value.len() - 1];
126        }
127        value.trim().to_string()
128    }
129
130    fn resolve_token(ctx: &ProviderContext) -> Result<String, SpendPanelError> {
131        for key in ["token", "api_key"] {
132            if let Some(value) = ctx.config.get(key) {
133                let cleaned = Self::clean(value);
134                if !cleaned.is_empty() {
135                    return Ok(cleaned);
136                }
137            }
138        }
139        for env in ["COPILOT_API_TOKEN", "GITHUB_TOKEN", "GH_TOKEN"] {
140            if let Ok(value) = std::env::var(env) {
141                let cleaned = Self::clean(&value);
142                if !cleaned.is_empty() {
143                    return Ok(cleaned);
144                }
145            }
146        }
147        Err(SpendPanelError::AuthFailed(
148            "copilot".into(),
149            "no token found in token/api_key config, COPILOT_API_TOKEN, GITHUB_TOKEN, or GH_TOKEN"
150                .into(),
151        ))
152    }
153
154    fn build_client(ctx: &ProviderContext) -> Result<reqwest::Client, SpendPanelError> {
155        reqwest::Client::builder()
156            .timeout(std::time::Duration::from_secs(ctx.timeout_secs))
157            .build()
158            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))
159    }
160
161    async fn fetch_usage_response(
162        base_url: &str,
163        client: &reqwest::Client,
164        token: &str,
165    ) -> Result<CopilotUsageResponse, SpendPanelError> {
166        let url = format!("{}/copilot_internal/user", base_url.trim_end_matches('/'));
167        let resp = client
168            .get(url)
169            .header("Authorization", format!("token {}", token))
170            .header("Accept", "application/json")
171            .header("X-Github-Api-Version", "2025-04-01")
172            .header("Editor-Version", "vscode/1.96.2")
173            .header("User-Agent", "GitHubCopilotChat/0.26.7")
174            .send()
175            .await
176            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
177
178        let status = resp.status();
179        let body = resp
180            .text()
181            .await
182            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
183
184        if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
185            return Err(SpendPanelError::AuthFailed(
186                "copilot".into(),
187                format!("invalid token (HTTP {})", status.as_u16()),
188            ));
189        }
190        if status == reqwest::StatusCode::NOT_FOUND {
191            return Err(SpendPanelError::ProviderError(
192                "copilot".into(),
193                "no Copilot subscription for this token (HTTP 404)".into(),
194            ));
195        }
196        if !status.is_success() {
197            return Err(SpendPanelError::ProviderError(
198                "copilot".into(),
199                format!("HTTP {}: {}", status, body),
200            ));
201        }
202
203        serde_json::from_str(&body)
204            .map_err(|e| SpendPanelError::ParseError("copilot".into(), e.to_string()))
205    }
206
207    fn snapshot_from_response(
208        resp: &CopilotUsageResponse,
209    ) -> Result<UsageSnapshot, SpendPanelError> {
210        let premium = resp
211            .quota_snapshots
212            .premium_interactions
213            .as_ref()
214            .and_then(|s| s.to_rate_window("Premium"));
215        let chat = resp
216            .quota_snapshots
217            .chat
218            .as_ref()
219            .and_then(|s| s.to_rate_window("Chat"));
220
221        let (primary, secondary) = match (premium, chat) {
222            (Some(p), c) => (Some(p), c),
223            (None, Some(c)) => (None, Some(c)),
224            (None, None) => {
225                if resp.token_based_billing {
226                    (None, None)
227                } else {
228                    return Err(SpendPanelError::ProviderError(
229                        "copilot".into(),
230                        "no usable quota in response".into(),
231                    ));
232                }
233            }
234        };
235
236        let mut snapshot = UsageSnapshot::new("copilot");
237        snapshot.primary_rate_window = primary;
238        snapshot.secondary_rate_window = secondary;
239        if let Some(plan) = resp
240            .copilot_plan
241            .as_deref()
242            .filter(|p| !p.is_empty() && *p != "unknown")
243        {
244            snapshot.plan = Some(PlanInfo {
245                name: capitalize(plan),
246                tier: None,
247                features: Vec::new(),
248                price: None,
249                currency: None,
250                billing_period: None,
251            });
252        }
253        Ok(snapshot)
254    }
255}
256
257fn capitalize(s: &str) -> String {
258    let mut chars = s.chars();
259    match chars.next() {
260        Some(first) => first.to_uppercase().collect::<String>() + chars.as_str(),
261        None => String::new(),
262    }
263}
264
265impl Default for CopilotProvider {
266    fn default() -> Self {
267        Self::new()
268    }
269}
270
271#[async_trait]
272impl UsageProvider for CopilotProvider {
273    fn metadata(&self) -> &ProviderMetadata {
274        &self.metadata
275    }
276
277    fn detect_credentials(&self) -> bool {
278        ["COPILOT_API_TOKEN", "GITHUB_TOKEN", "GH_TOKEN"]
279            .iter()
280            .any(|env| {
281                std::env::var(env)
282                    .map(|v| !v.trim().is_empty())
283                    .unwrap_or(false)
284            })
285    }
286
287    async fn fetch_usage(&self, ctx: &ProviderContext) -> Result<UsageSnapshot, SpendPanelError> {
288        let token = Self::resolve_token(ctx)?;
289        let client = Self::build_client(ctx)?;
290        let response = Self::fetch_usage_response(self.api_base(), &client, &token).await?;
291        Self::snapshot_from_response(&response)
292    }
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298    use pretty_assertions::assert_eq;
299    use wiremock::matchers::{header, method, path};
300    use wiremock::{Mock, MockServer, ResponseTemplate};
301
302    const SAMPLE: &str = r#"{
303      "quota_snapshots": {
304        "premium_interactions": {"entitlement": 300, "remaining": 90, "percent_remaining": 30, "unlimited": false},
305        "chat": {"entitlement": 0, "remaining": 0, "unlimited": true}
306      },
307      "copilot_plan": "individual",
308      "token_based_billing": false
309    }"#;
310
311    fn parse(body: &str) -> CopilotUsageResponse {
312        serde_json::from_str(body).unwrap()
313    }
314
315    #[test]
316    fn test_metadata() {
317        let p = CopilotProvider::new();
318        assert_eq!(p.metadata().id, "copilot");
319        assert!(p.metadata().auth_methods.contains(&"token"));
320    }
321
322    #[test]
323    fn test_resolve_token_missing_is_error() {
324        let err = CopilotProvider::resolve_token(&ProviderContext::new()).unwrap_err();
325        assert!(matches!(err, SpendPanelError::AuthFailed(_, _)));
326    }
327
328    #[test]
329    fn test_premium_window_used_from_entitlement() {
330        let resp = parse(SAMPLE);
331        let window = resp
332            .quota_snapshots
333            .premium_interactions
334            .as_ref()
335            .unwrap()
336            .to_rate_window("Premium")
337            .unwrap();
338        assert_eq!(window.used, Some(210));
339        assert_eq!(window.limit, Some(300));
340    }
341
342    #[test]
343    fn test_unlimited_chat_full_remaining() {
344        let resp = parse(SAMPLE);
345        let window = resp
346            .quota_snapshots
347            .chat
348            .as_ref()
349            .unwrap()
350            .to_rate_window("Chat")
351            .unwrap();
352        // unlimited → 100% remaining → 0% used
353        assert_eq!(window.used, Some(0));
354    }
355
356    #[test]
357    fn test_placeholder_dropped() {
358        let snap = CopilotQuotaSnapshot {
359            entitlement: 0.0,
360            remaining: 0.0,
361            percent_remaining: Some(100.0),
362            unlimited: false,
363        };
364        assert!(snap.is_placeholder());
365        assert!(snap.to_rate_window("Premium").is_none());
366    }
367
368    #[test]
369    fn test_snapshot_plan_capitalized() {
370        let snapshot = CopilotProvider::snapshot_from_response(&parse(SAMPLE)).unwrap();
371        assert_eq!(snapshot.plan.unwrap().name, "Individual");
372        assert!(snapshot.primary_rate_window.is_some());
373    }
374
375    #[test]
376    fn test_token_based_billing_no_quota_ok() {
377        let body = r#"{
378          "quota_snapshots": {
379            "premium_interactions": {"entitlement": 0, "remaining": 0},
380            "chat": {"entitlement": 0, "remaining": 0}
381          },
382          "copilot_plan": "business",
383          "token_based_billing": true
384        }"#;
385        let snapshot = CopilotProvider::snapshot_from_response(&parse(body)).unwrap();
386        assert!(snapshot.primary_rate_window.is_none());
387        assert!(snapshot.secondary_rate_window.is_none());
388        assert_eq!(snapshot.plan.unwrap().name, "Business");
389    }
390
391    #[test]
392    fn test_chat_only_leaves_primary_empty() {
393        // No premium quota → primary stays empty, chat lands in secondary.
394        let body = r#"{
395          "quota_snapshots": {
396            "chat": {"entitlement": 50, "remaining": 20, "percent_remaining": 40}
397          },
398          "copilot_plan": "free"
399        }"#;
400        let snapshot = CopilotProvider::snapshot_from_response(&parse(body)).unwrap();
401        assert!(snapshot.primary_rate_window.is_none());
402        let chat = snapshot.secondary_rate_window.unwrap();
403        assert_eq!(chat.used, Some(30));
404        assert_eq!(chat.label, "Chat");
405    }
406
407    #[test]
408    fn test_no_usable_quota_without_token_billing_errors() {
409        let body = r#"{"quota_snapshots":{},"copilot_plan":"free","token_based_billing":false}"#;
410        assert!(matches!(
411            CopilotProvider::snapshot_from_response(&parse(body)).unwrap_err(),
412            SpendPanelError::ProviderError(_, _)
413        ));
414    }
415
416    #[tokio::test]
417    async fn test_fetch_usage_success() {
418        let server = MockServer::start().await;
419        Mock::given(method("GET"))
420            .and(path("/copilot_internal/user"))
421            .and(header("authorization", "token gho_test"))
422            .respond_with(ResponseTemplate::new(200).set_body_raw(SAMPLE, "application/json"))
423            .mount(&server)
424            .await;
425
426        let provider = CopilotProvider::with_base_url(&server.uri());
427        let mut ctx = ProviderContext::new();
428        ctx.config.insert("token".into(), "gho_test".into());
429        let snapshot = provider.fetch_usage(&ctx).await.unwrap();
430        assert_eq!(snapshot.primary_rate_window.unwrap().limit, Some(300));
431    }
432
433    #[tokio::test]
434    async fn test_fetch_usage_401_is_auth_failed() {
435        let server = MockServer::start().await;
436        Mock::given(method("GET"))
437            .and(path("/copilot_internal/user"))
438            .respond_with(ResponseTemplate::new(401))
439            .mount(&server)
440            .await;
441
442        let provider = CopilotProvider::with_base_url(&server.uri());
443        let mut ctx = ProviderContext::new();
444        ctx.config.insert("token".into(), "bad".into());
445        let err = provider.fetch_usage(&ctx).await.unwrap_err();
446        assert!(matches!(err, SpendPanelError::AuthFailed(_, _)));
447    }
448}
usage_monitor_cli/provider/copilot.rs

usage_monitor_cli/provider/
copilot.rs
