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

ai_usagebar/opencode_go/
fetch.rs
