ai_usagebar/copilot/
fetch.rs

1//! GitHub Copilot quota fetch with cache isolation and no credential storage.
2
3use std::fmt::Write as _;
4use std::time::Duration;
5
6use sha2::{Digest, Sha256};
7
8use crate::cache::{Cache, acquire_lock_async};
9use crate::error::{AppError, Result};
10use crate::vendor::{MAX_BODY_BYTES, read_body_capped};
11
12use super::types::{Response, Snapshot, to_snapshot};
13
14pub const USER_URL: &str = "https://api.github.com/copilot_internal/user";
15const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
16const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
17const SCHEMA_ERROR: &str = "GitHub Copilot quota response schema mismatch";
18
19#[derive(Debug, Clone)]
20pub struct Endpoints {
21    pub user: String,
22}
23
24impl Default for Endpoints {
25    fn default() -> Self {
26        Self {
27            user: USER_URL.to_string(),
28        }
29    }
30}
31
32pub type FetchOutcome = crate::outcome::Outcome<Snapshot>;
33
34pub async fn fetch_snapshot(
35    client: &reqwest::Client,
36    token: &str,
37    cache: &Cache,
38    endpoints: &Endpoints,
39    ttl: Duration,
40) -> Result<FetchOutcome> {
41    cache.ensure_dir()?;
42    let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
43    let target = target_key(endpoints, token);
44    if let Some(bytes) = cache.fresh_payload(ttl)?
45        && let Ok(snapshot) = parse_cache(&bytes, &target)
46    {
47        return Ok(crate::outcome::Outcome::cached(snapshot, cache, false));
48    }
49    match fetch_live(client, token, endpoints).await {
50        Ok(snapshot) => {
51            let bytes = serde_json::to_vec(&serde_json::json!({
52                "target": target,
53                "snapshot": snapshot,
54            }))?;
55            cache.write_payload(&bytes)?;
56            Ok(crate::outcome::Outcome::fresh(snapshot))
57        }
58        Err(error @ AppError::Transport(_)) => fallback_or_error(cache, None, &target, error),
59        Err(AppError::Http { status, .. }) => {
60            let message = status_message(status).to_string();
61            cache.mark_stale();
62            let diagnostic = cache.write_last_error(status, &message);
63            fallback_or_error(
64                cache,
65                Some(diagnostic),
66                &target,
67                AppError::Http {
68                    status,
69                    body: message,
70                },
71            )
72        }
73        Err(AppError::Schema(_)) => {
74            cache.mark_stale();
75            let diagnostic = cache.write_last_error(0, SCHEMA_ERROR);
76            fallback_or_error(
77                cache,
78                Some(diagnostic),
79                &target,
80                AppError::Schema(SCHEMA_ERROR.to_string()),
81            )
82        }
83        Err(error) => fallback_or_error(cache, None, &target, error),
84    }
85}
86
87async fn fetch_live(
88    client: &reqwest::Client,
89    token: &str,
90    endpoints: &Endpoints,
91) -> Result<Snapshot> {
92    let response = tokio::time::timeout(
93        HTTP_TIMEOUT,
94        client
95            .get(&endpoints.user)
96            .header(reqwest::header::AUTHORIZATION, format!("token {token}"))
97            .header(reqwest::header::ACCEPT, "application/json")
98            .header("Editor-Version", "vscode/1.96.2")
99            .header("Editor-Plugin-Version", "copilot-chat/0.26.7")
100            .header(reqwest::header::USER_AGENT, "GitHubCopilotChat/0.26.7")
101            .header("X-GitHub-Api-Version", "2025-04-01")
102            .send(),
103    )
104    .await
105    .map_err(|_| AppError::Transport("GitHub Copilot request timed out".into()))??;
106    let status = response.status();
107    let bytes = read_body_capped(response, MAX_BODY_BYTES).await?;
108    if !status.is_success() {
109        return Err(AppError::Http {
110            status: status.as_u16(),
111            body: status_message(status.as_u16()).into(),
112        });
113    }
114    let response: Response =
115        serde_json::from_slice(&bytes).map_err(|_| AppError::Schema(SCHEMA_ERROR.to_string()))?;
116    to_snapshot(response)
117}
118
119/// Bind cache reuse to both endpoint and token without writing either raw
120/// credential or the full response to disk.
121fn target_key(endpoints: &Endpoints, token: &str) -> String {
122    let digest = Sha256::digest(token.as_bytes());
123    let mut fingerprint = String::with_capacity(digest.len() * 2);
124    for byte in digest {
125        let _ = write!(fingerprint, "{byte:02x}");
126    }
127    format!("{}|token:{fingerprint}", endpoints.user)
128}
129
130fn parse_cache(bytes: &[u8], target: &str) -> Result<Snapshot> {
131    let value: serde_json::Value = serde_json::from_slice(bytes)
132        .map_err(|_| AppError::Schema("GitHub Copilot cache is invalid".into()))?;
133    if value.get("target").and_then(serde_json::Value::as_str) != Some(target) {
134        return Err(AppError::Schema(
135            "GitHub Copilot cache belongs to a different account".into(),
136        ));
137    }
138    serde_json::from_value(
139        value.get("snapshot").cloned().ok_or_else(|| {
140            AppError::Schema("GitHub Copilot cache is missing its snapshot".into())
141        })?,
142    )
143    .map_err(|_| AppError::Schema("GitHub Copilot cache has an invalid snapshot".into()))
144}
145
146fn fallback_or_error(
147    cache: &Cache,
148    diagnostic: Option<(u16, String)>,
149    target: &str,
150    error: AppError,
151) -> Result<FetchOutcome> {
152    crate::outcome::fallback(cache, diagnostic, error, |bytes| parse_cache(bytes, target))
153}
154
155fn status_message(status: u16) -> &'static str {
156    match status {
157        401 | 403 => crate::error::AUTH_FAILURE_MESSAGE,
158        429 => "GitHub Copilot rate limited the quota request",
159        500..=599 => "GitHub Copilot quota endpoint is unavailable",
160        _ => "GitHub Copilot quota request failed",
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167    use crate::copilot::types::Quota;
168
169    fn cache_in(dir: &std::path::Path) -> Cache {
170        Cache::at(dir.join("copilot"))
171    }
172
173    #[tokio::test]
174    async fn requests_vs_code_endpoint_with_only_copilot_token_and_normalizes_quotas() {
175        let mut server = mockito::Server::new_async().await;
176        let request = server
177            .mock("GET", "/copilot_internal/user")
178            .expect(1)
179            .match_header("authorization", "token mock-oauth-token")
180            .match_header("accept", "application/json")
181            .match_header("editor-version", "vscode/1.96.2")
182            .match_header("editor-plugin-version", "copilot-chat/0.26.7")
183            .match_header("user-agent", "GitHubCopilotChat/0.26.7")
184            .match_header("x-github-api-version", "2025-04-01")
185            .with_status(200)
186            .with_body(
187                r#"{"copilot_plan":"pro","quota_reset_date":"2026-09-15","quota_snapshots":{"premium_interactions":{"entitlement":300,"remaining":45,"percent_remaining":15},"chat":{"entitlement":1000,"remaining":250},"completions":{"unlimited":true}}}"#,
188            )
189            .create_async()
190            .await;
191        let dir = tempfile::tempdir().unwrap();
192        let cache = cache_in(dir.path());
193        let endpoints = Endpoints {
194            user: format!("{}/copilot_internal/user", server.url()),
195        };
196        let outcome = fetch_snapshot(
197            &reqwest::Client::new(),
198            "mock-oauth-token",
199            &cache,
200            &endpoints,
201            Duration::from_secs(60),
202        )
203        .await
204        .unwrap();
205        let cached = fetch_snapshot(
206            &reqwest::Client::new(),
207            "mock-oauth-token",
208            &cache,
209            &endpoints,
210            Duration::from_secs(60),
211        )
212        .await
213        .unwrap();
214        request.assert_async().await;
215        assert_eq!(outcome.snapshot.premium.unwrap().used_pct(), 85);
216        assert!(!cached.stale);
217        assert!(cached.cache_age.is_some());
218        assert_eq!(
219            outcome.snapshot.chat.unwrap().used_and_entitlement(),
220            Some((750, 1000))
221        );
222        assert!(outcome.snapshot.completions.unwrap().unlimited);
223    }
224
225    #[tokio::test]
226    async fn rejected_refresh_uses_stale_cache_without_storing_token_or_response() {
227        let mut server = mockito::Server::new_async().await;
228        server
229            .mock("GET", "/copilot_internal/user")
230            .with_status(401)
231            .with_body(r#"{"message":"contains private account data"}"#)
232            .create_async()
233            .await;
234        let dir = tempfile::tempdir().unwrap();
235        let cache = cache_in(dir.path());
236        let endpoints = Endpoints {
237            user: format!("{}/copilot_internal/user", server.url()),
238        };
239        let target = target_key(&endpoints, "mock-oauth-token");
240        let snapshot = Snapshot {
241            plan: "pro".into(),
242            premium: Some(Quota {
243                percent_remaining: 80,
244                entitlement: Some(300),
245                remaining: Some(240),
246                unlimited: false,
247            }),
248            chat: None,
249            completions: None,
250            reset_at: None,
251        };
252        cache
253            .write_payload(
254                serde_json::json!({"target": target, "snapshot": snapshot})
255                    .to_string()
256                    .as_bytes(),
257            )
258            .unwrap();
259
260        let outcome = fetch_snapshot(
261            &reqwest::Client::new(),
262            "mock-oauth-token",
263            &cache,
264            &endpoints,
265            Duration::ZERO,
266        )
267        .await
268        .unwrap();
269        assert!(outcome.stale);
270        assert_eq!(outcome.last_error.unwrap().0, 401);
271        assert_eq!(outcome.snapshot.premium.unwrap().used_pct(), 20);
272        let cached = std::fs::read_to_string(cache.payload_path()).unwrap();
273        assert!(!cached.contains("mock-oauth-token"));
274        assert!(!cached.contains("private account data"));
275    }
276
277    #[tokio::test]
278    async fn invalid_success_body_is_schema_error_on_a_cold_cache() {
279        let mut server = mockito::Server::new_async().await;
280        server
281            .mock("GET", "/copilot_internal/user")
282            .with_status(200)
283            .with_body(r#"{"quota_snapshots":{}}"#)
284            .create_async()
285            .await;
286        let dir = tempfile::tempdir().unwrap();
287        let error = fetch_snapshot(
288            &reqwest::Client::new(),
289            "mock-oauth-token",
290            &cache_in(dir.path()),
291            &Endpoints {
292                user: format!("{}/copilot_internal/user", server.url()),
293            },
294            Duration::ZERO,
295        )
296        .await
297        .unwrap_err();
298        assert!(matches!(error, AppError::Schema(message) if message == SCHEMA_ERROR));
299    }
300}
ai_usagebar/copilot/fetch.rs

ai_usagebar/copilot/
fetch.rs
