ai_usagebar/opencode_go/
fetch.rs

1use std::fmt::Write as _;
2use std::time::Duration;
3
4use sha2::{Digest, Sha256};
5
6use crate::cache::{Cache, MAX_STALE, acquire_lock_async};
7use crate::error::{AUTH_FAILURE_MESSAGE, AppError, Result};
8use crate::vendor::{MAX_BODY_BYTES, read_body_capped};
9
10use super::types::{Usage, parse_usage};
11
12pub const BASE_URL: &str = "https://opencode.ai/zen/go/v1/usage";
13
14const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
15const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
16const SCHEMA_ERROR: &str = "OpenCode Go usage response schema mismatch";
17
18#[derive(Debug, Clone)]
19pub struct Endpoints {
20    pub usage: String,
21}
22
23impl Default for Endpoints {
24    fn default() -> Self {
25        Self {
26            usage: BASE_URL.to_string(),
27        }
28    }
29}
30
31#[derive(Debug, Clone)]
32pub struct FetchOutcome {
33    pub snapshot: Usage,
34    pub stale: bool,
35    pub last_error: Option<(u16, String)>,
36    pub cache_age: Option<Duration>,
37}
38
39pub async fn fetch_snapshot(
40    client: &reqwest::Client,
41    api_key: &str,
42    cache: &Cache,
43    endpoints: &Endpoints,
44    ttl: Duration,
45) -> Result<FetchOutcome> {
46    cache.ensure_dir()?;
47    let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
48    let target = target_key(endpoints, api_key);
49
50    if let Some(bytes) = cache.fresh_payload(ttl)?
51        && let Ok(snapshot) = parse_cache(&bytes, &target)
52    {
53        return Ok(FetchOutcome {
54            snapshot,
55            stale: false,
56            last_error: cache.read_last_error(),
57            cache_age: cache.payload_age(),
58        });
59    }
60
61    match fetch_live(client, &endpoints.usage, api_key).await {
62        Ok(snapshot) => {
63            let body = serde_json::to_vec(&serde_json::json!({
64                "target": target,
65                "response": usage_repr(&snapshot),
66            }))?;
67            cache.write_payload(&body)?;
68            Ok(FetchOutcome {
69                snapshot,
70                stale: false,
71                last_error: None,
72                cache_age: Some(Duration::ZERO),
73            })
74        }
75        Err(error @ AppError::Transport(_)) => fallback_or_error(cache, None, &target, error),
76        Err(AppError::Http { status, .. }) => {
77            let message = status_message(status).to_string();
78            cache.mark_stale();
79            cache.write_last_error(status, &message);
80            fallback_or_error(
81                cache,
82                Some((status, message.clone())),
83                &target,
84                AppError::Http {
85                    status,
86                    body: message,
87                },
88            )
89        }
90        Err(AppError::Schema(_)) => {
91            let message = SCHEMA_ERROR.to_string();
92            cache.mark_stale();
93            cache.write_last_error(0, &message);
94            fallback_or_error(
95                cache,
96                Some((0, message.clone())),
97                &target,
98                AppError::Schema(message),
99            )
100        }
101        Err(error) => fallback_or_error(cache, None, &target, error),
102    }
103}
104
105async fn fetch_live(client: &reqwest::Client, url: &str, api_key: &str) -> Result<Usage> {
106    let response = tokio::time::timeout(
107        HTTP_TIMEOUT,
108        client
109            .get(url)
110            .bearer_auth(api_key)
111            .header(reqwest::header::ACCEPT, "application/json")
112            .send(),
113    )
114    .await
115    .map_err(|_| AppError::Transport("OpenCode Go request timed out".to_string()))??;
116
117    let status = response.status();
118    let body = read_body_capped(response, MAX_BODY_BYTES).await?;
119    if !status.is_success() {
120        return Err(AppError::Http {
121            status: status.as_u16(),
122            body: status_message(status.as_u16()).to_string(),
123        });
124    }
125
126    let snapshot = parse_payload(&body)?;
127    Ok(snapshot)
128}
129
130/// Stable, non-secret identity for the endpoint and account selected by the
131/// API key. The usage endpoint resolves the key to a user/workspace, so cache
132/// reuse must fail closed when either input changes.
133fn target_key(endpoints: &Endpoints, api_key: &str) -> String {
134    let digest = Sha256::digest(api_key.as_bytes());
135    let mut fingerprint = String::with_capacity(digest.len() * 2);
136    for byte in digest {
137        let _ = write!(fingerprint, "{byte:02x}");
138    }
139    format!("{}|key:{fingerprint}", endpoints.usage)
140}
141
142fn usage_repr(usage: &Usage) -> serde_json::Value {
143    let window = |window: &super::types::Window| {
144        serde_json::json!({
145            "status": window.status,
146            "percent": window.percent,
147            "resetsAt": window.resets_at.to_rfc3339(),
148        })
149    };
150    let mut windows = serde_json::Map::new();
151    for (name, value) in [
152        ("rolling", usage.rolling.as_ref()),
153        ("weekly", usage.weekly.as_ref()),
154        ("monthly", usage.monthly.as_ref()),
155    ] {
156        if let Some(value) = value {
157            windows.insert(name.into(), window(value));
158        }
159    }
160    serde_json::json!({ "usage": windows })
161}
162
163fn parse_payload(body: &[u8]) -> Result<Usage> {
164    let value: serde_json::Value =
165        serde_json::from_slice(body).map_err(|_| AppError::Schema(SCHEMA_ERROR.to_string()))?;
166    parse_usage(&value).map_err(|_| AppError::Schema(SCHEMA_ERROR.to_string()))
167}
168
169fn status_message(status: u16) -> &'static str {
170    match status {
171        401 | 403 => AUTH_FAILURE_MESSAGE,
172        429 => "OpenCode Go request was rate limited",
173        500..=599 => "OpenCode Go service is temporarily unavailable",
174        _ => "OpenCode Go request failed",
175    }
176}
177
178fn fallback_or_error(
179    cache: &Cache,
180    last_error: Option<(u16, String)>,
181    target: &str,
182    error: AppError,
183) -> Result<FetchOutcome> {
184    if let Some(snapshot) = cached_outcome(cache, last_error, target)? {
185        return Ok(snapshot);
186    }
187    Err(error)
188}
189
190fn cached_outcome(
191    cache: &Cache,
192    last_error: Option<(u16, String)>,
193    target: &str,
194) -> Result<Option<FetchOutcome>> {
195    let Some(body) = cache.fallback_payload(MAX_STALE)? else {
196        return Ok(None);
197    };
198    let Ok(snapshot) = parse_cache(&body, target) else {
199        return Ok(None);
200    };
201    Ok(Some(FetchOutcome {
202        snapshot,
203        stale: true,
204        last_error: last_error.or_else(|| cache.read_last_error()),
205        cache_age: cache.payload_age(),
206    }))
207}
208
209fn parse_cache(body: &[u8], target: &str) -> Result<Usage> {
210    let value: serde_json::Value = serde_json::from_slice(body)
211        .map_err(|_| AppError::Schema("OpenCode Go cache is invalid".into()))?;
212    if value.get("target").and_then(serde_json::Value::as_str) != Some(target) {
213        return Err(AppError::Schema(
214            "OpenCode Go cache belongs to a different account".into(),
215        ));
216    }
217    let response = value
218        .get("response")
219        .ok_or_else(|| AppError::Schema("OpenCode Go cache is missing its response".into()))?;
220    parse_usage(response).map_err(|_| AppError::Schema("OpenCode Go cache is invalid".into()))
221}
222
223#[cfg(test)]
224mod tests {
225    use std::time::Duration;
226
227    use tempfile::TempDir;
228
229    use super::*;
230    use crate::cache::Cache;
231
232    const GOOD_BODY: &str = r#"{
233        "usage": {
234            "rolling": {"status":"ok","percent":12.3,"resetsAt":"2026-08-16T20:00:00Z"},
235            "weekly": {"status":"ok","percent":45.6,"resetsAt":"2026-08-20T00:00:00Z"},
236            "monthly": {"status":"ok","percent":78.9,"resetsAt":"2026-09-01T00:00:00Z"}
237        }
238    }"#;
239
240    fn cache_fixture() -> (TempDir, Cache) {
241        let dir = TempDir::new().expect("temporary cache directory");
242        let cache = Cache::at(dir.path().join("opencode-go"));
243        cache.ensure_dir().expect("cache directory");
244        (dir, cache)
245    }
246
247    #[tokio::test]
248    async fn fetches_usage_with_bearer_and_accept_headers() {
249        let mut server = mockito::Server::new_async().await;
250        let mock = server
251            .mock("GET", "/zen/go/v1/usage")
252            .match_header("authorization", "Bearer test-key")
253            .match_header("accept", "application/json")
254            .with_status(200)
255            .with_body(GOOD_BODY)
256            .create_async()
257            .await;
258        let (_dir, cache) = cache_fixture();
259        let endpoints = Endpoints {
260            usage: format!("{}/zen/go/v1/usage", server.url()),
261        };
262
263        let output = fetch_snapshot(
264            &reqwest::Client::new(),
265            "test-key",
266            &cache,
267            &endpoints,
268            Duration::from_secs(60),
269        )
270        .await
271        .expect("successful usage response");
272
273        assert_eq!(output.snapshot.rolling.expect("rolling").percent, 12.3);
274        assert!(!output.stale);
275        assert!(output.last_error.is_none());
276        mock.assert_async().await;
277    }
278
279    #[tokio::test]
280    async fn unauthorized_without_cache_returns_redacted_error() {
281        let mut server = mockito::Server::new_async().await;
282        server
283            .mock("GET", "/usage")
284            .with_status(401)
285            .with_body("secret-token should never be returned")
286            .create_async()
287            .await;
288        let (_dir, cache) = cache_fixture();
289        let endpoints = Endpoints {
290            usage: format!("{}/usage", server.url()),
291        };
292
293        let error = fetch_snapshot(
294            &reqwest::Client::new(),
295            "test-key",
296            &cache,
297            &endpoints,
298            Duration::ZERO,
299        )
300        .await
301        .expect_err("401 without a cache must fail");
302
303        let rendered = error.to_string();
304        assert!(rendered.contains("401"));
305        assert!(!rendered.contains("secret-token"));
306        assert!(!rendered.contains("test-key"));
307    }
308
309    #[tokio::test]
310    async fn schema_error_does_not_include_response_body() {
311        let mut server = mockito::Server::new_async().await;
312        server
313            .mock("GET", "/usage")
314            .with_status(200)
315            .with_body(
316                r#"{"error":"schema-secret", "usage":{"rolling":{"status":"ok","percent":"schema-secret","resetsAt":"2026-08-16T20:00:00Z"}}}"#,
317            )
318            .create_async()
319            .await;
320        let (_dir, cache) = cache_fixture();
321        let endpoints = Endpoints {
322            usage: format!("{}/usage", server.url()),
323        };
324
325        let error = fetch_snapshot(
326            &reqwest::Client::new(),
327            "test-key",
328            &cache,
329            &endpoints,
330            Duration::ZERO,
331        )
332        .await
333        .expect_err("schema mismatch must fail");
334
335        let rendered = error.to_string();
336        assert!(rendered.to_ascii_lowercase().contains("schema"));
337        assert!(!rendered.contains("schema-secret"));
338        assert!(!rendered.contains("test-key"));
339    }
340
341    #[tokio::test]
342    async fn status_errors_are_redacted_for_403_429_and_server_errors() {
343        for status in [403, 429, 500, 503] {
344            let mut server = mockito::Server::new_async().await;
345            server
346                .mock("GET", "/usage")
347                .with_status(status)
348                .with_body("body-secret")
349                .create_async()
350                .await;
351            let (_dir, cache) = cache_fixture();
352            let endpoints = Endpoints {
353                usage: format!("{}/usage", server.url()),
354            };
355
356            let error = fetch_snapshot(
357                &reqwest::Client::new(),
358                "test-key",
359                &cache,
360                &endpoints,
361                Duration::ZERO,
362            )
363            .await
364            .expect_err("status without cache must fail");
365            let rendered = error.to_string();
366            assert!(rendered.contains(&status.to_string()));
367            assert!(!rendered.contains("body-secret"));
368            assert!(!rendered.contains("test-key"));
369        }
370    }
371
372    #[tokio::test]
373    async fn changing_api_keys_never_reuses_another_accounts_fresh_cache() {
374        let mut server = mockito::Server::new_async().await;
375        let first = server
376            .mock("GET", "/usage")
377            .match_header("authorization", "Bearer first-key")
378            .with_status(200)
379            .with_body(GOOD_BODY)
380            .expect(1)
381            .create_async()
382            .await;
383        let second_body = GOOD_BODY.replace("12.3", "91.2");
384        let second = server
385            .mock("GET", "/usage")
386            .match_header("authorization", "Bearer second-key")
387            .with_status(200)
388            .with_body(second_body)
389            .expect(1)
390            .create_async()
391            .await;
392        let (_dir, cache) = cache_fixture();
393        let endpoints = Endpoints {
394            usage: format!("{}/usage", server.url()),
395        };
396
397        fetch_snapshot(
398            &reqwest::Client::new(),
399            "first-key",
400            &cache,
401            &endpoints,
402            Duration::from_secs(60),
403        )
404        .await
405        .unwrap();
406        let second_output = fetch_snapshot(
407            &reqwest::Client::new(),
408            "second-key",
409            &cache,
410            &endpoints,
411            Duration::from_secs(60),
412        )
413        .await
414        .unwrap();
415
416        assert_eq!(second_output.snapshot.rolling.unwrap().percent, 91.2);
417        let persisted = String::from_utf8(cache.maybe_payload().unwrap().unwrap()).unwrap();
418        assert!(!persisted.contains("first-key"));
419        assert!(!persisted.contains("second-key"));
420        first.assert_async().await;
421        second.assert_async().await;
422    }
423
424    #[tokio::test]
425    async fn changing_api_keys_rejects_another_accounts_stale_fallback() {
426        let mut server = mockito::Server::new_async().await;
427        server
428            .mock("GET", "/usage")
429            .match_header("authorization", "Bearer first-key")
430            .with_status(200)
431            .with_body(GOOD_BODY)
432            .create_async()
433            .await;
434        server
435            .mock("GET", "/usage")
436            .match_header("authorization", "Bearer second-key")
437            .with_status(503)
438            .create_async()
439            .await;
440        let (_dir, cache) = cache_fixture();
441        let endpoints = Endpoints {
442            usage: format!("{}/usage", server.url()),
443        };
444        fetch_snapshot(
445            &reqwest::Client::new(),
446            "first-key",
447            &cache,
448            &endpoints,
449            Duration::ZERO,
450        )
451        .await
452        .unwrap();
453
454        let error = fetch_snapshot(
455            &reqwest::Client::new(),
456            "second-key",
457            &cache,
458            &endpoints,
459            Duration::ZERO,
460        )
461        .await
462        .expect_err("a cache for another endpoint/account must not be served");
463        assert!(matches!(error, AppError::Http { status: 503, .. }));
464    }
465}
ai_usagebar/opencode_go/fetch.rs

ai_usagebar/opencode_go/
fetch.rs
