ai_usagebar/
vendor.rs

1//! Shared vendor IDs and renderer/fetcher structs used by the widget and TUI.
2//!
3//! Snapshots remain a discriminated `VendorSnapshot` enum because the vendors
4//! have genuinely different shapes — see `usage.rs`.
5
6use std::time::Duration;
7
8use clap::ValueEnum;
9
10use crate::usage::VendorSnapshot;
11use crate::widget::cli::Cli;
12
13/// Outer reqwest client timeout shared by widget and TUI entry points.
14/// Vendor fetchers still apply their own tighter per-request timeouts.
15pub const HTTP_CLIENT_TIMEOUT: Duration = Duration::from_secs(30);
16
17/// Upper bound on a vendor response body. Every one of these endpoints returns
18/// a small JSON document — the largest observed is a few kilobytes — so this is
19/// generous by three orders of magnitude while still bounding the damage from a
20/// misbehaving proxy or a hijacked endpoint.
21pub const MAX_BODY_BYTES: usize = 2 * 1024 * 1024;
22
23/// Credential-bearing environment variables owned by ai-usagebar vendors.
24/// Subprocesses receive only the entries that belong to their own provider.
25pub(crate) const VENDOR_SECRET_ENV_VARS: &[&str] = &[
26    "ZAI_API_KEY",
27    "OPENROUTER_API_KEY",
28    "DEEPSEEK_API_KEY",
29    "KIMI_API_KEY",
30    "KILO_API_KEY",
31    "NOVITA_API_KEY",
32    "MINIMAX_API_KEY",
33    "MOONSHOT_API_KEY",
34    "XAI_MANAGEMENT_KEY",
35    "ANTHROPIC_ADMIN_KEY",
36    "XAI_API_KEY",
37    "GROK_API_KEY",
38    "OPENCODE_GO_API_KEY",
39];
40
41pub(crate) fn vendor_secret_env_vars_to_remove(keep: &[&str]) -> Vec<&'static str> {
42    VENDOR_SECRET_ENV_VARS
43        .iter()
44        .copied()
45        .filter(|var| !keep.contains(var))
46        .collect()
47}
48
49/// Follow ordinary vendor redirects without forwarding non-standard API-key
50/// headers to a different origin. Reqwest strips `Authorization` on sensitive
51/// redirects, but vendors also use headers such as `x-api-key`, which are not
52/// covered by that built-in list.
53pub fn same_origin_redirect_policy() -> reqwest::redirect::Policy {
54    reqwest::redirect::Policy::custom(|attempt| {
55        if attempt.previous().len() >= 10 {
56            return attempt.error("too many redirects");
57        }
58        let Some(origin) = attempt.previous().first() else {
59            return attempt.stop();
60        };
61        let target = attempt.url();
62        if target.scheme() == origin.scheme()
63            && target.host_str() == origin.host_str()
64            && target.port_or_known_default() == origin.port_or_known_default()
65        {
66            attempt.follow()
67        } else {
68            attempt.stop()
69        }
70    })
71}
72
73/// Read a response body with an upper bound.
74///
75/// Every vendor buffered the whole body with `resp.bytes()` *before* anything
76/// validated it. The widget is re-executed by Waybar every 60s, so an endpoint
77/// answering with an unbounded stream had a free hand at the machine's memory.
78/// `Content-Length` is checked first when present, then the body is read in
79/// chunks so a lying or absent length cannot get past the cap either.
80pub async fn read_body_capped(
81    mut resp: reqwest::Response,
82    max: usize,
83) -> crate::error::Result<Vec<u8>> {
84    let too_big = |n: u64| {
85        crate::error::AppError::Schema(format!(
86            "response body exceeds the {max}-byte limit ({n} bytes); refusing to buffer it"
87        ))
88    };
89    if let Some(len) = resp.content_length()
90        && len > max as u64
91    {
92        return Err(too_big(len));
93    }
94    let mut buf: Vec<u8> = Vec::new();
95    while let Some(chunk) = resp.chunk().await? {
96        if chunk.len() > max.saturating_sub(buf.len()) {
97            return Err(too_big(buf.len().saturating_add(chunk.len()) as u64));
98        }
99        buf.extend_from_slice(&chunk);
100    }
101    Ok(buf)
102}
103
104/// Stable enum used by `--vendor` and in config files.
105#[derive(
106    Debug, Clone, Copy, ValueEnum, PartialEq, Eq, Hash, serde::Deserialize, serde::Serialize,
107)]
108#[serde(rename_all = "lowercase")]
109pub enum VendorId {
110    Anthropic,
111    #[serde(rename = "anthropic_api")]
112    AnthropicApi,
113    Openai,
114    Zai,
115    Openrouter,
116    Deepseek,
117    Kimi,
118    Kilo,
119    Novita,
120    Moonshot,
121    Grok,
122    Supergrok,
123    Antigravity,
124    Cursor,
125    Minimax,
126    Kiro,
127    #[serde(rename = "nous")]
128    NousResearch,
129    #[serde(rename = "opencode-go")]
130    OpenCodeGo,
131}
132
133impl VendorId {
134    pub fn slug(self) -> &'static str {
135        match self {
136            VendorId::Anthropic => "anthropic",
137            VendorId::AnthropicApi => "anthropic_api",
138            VendorId::Openai => "openai",
139            VendorId::Zai => "zai",
140            VendorId::Openrouter => "openrouter",
141            VendorId::Deepseek => "deepseek",
142            VendorId::Kimi => "kimi",
143            VendorId::Kilo => "kilo",
144            VendorId::Novita => "novita",
145            VendorId::Moonshot => "moonshot",
146            VendorId::Grok => "grok",
147            VendorId::Supergrok => "supergrok",
148            VendorId::Antigravity => "antigravity",
149            VendorId::Cursor => "cursor",
150            VendorId::Minimax => "minimax",
151            VendorId::Kiro => "kiro",
152            VendorId::NousResearch => "nous",
153            VendorId::OpenCodeGo => "opencode-go",
154        }
155    }
156
157    /// Canonical human-readable name for shared reports and compact UI labels.
158    /// Platform frontends may add context (for example, "GLM (Z.AI)" in a
159    /// wide TUI tab), but should not carry their own full vendor-name table.
160    pub fn display_name(self) -> &'static str {
161        match self {
162            VendorId::Anthropic => "Claude",
163            VendorId::AnthropicApi => "Anthropic API",
164            VendorId::Openai => "Codex",
165            VendorId::Zai => "Z.AI",
166            VendorId::Openrouter => "OpenRouter",
167            VendorId::Deepseek => "DeepSeek",
168            VendorId::Kimi => "Kimi",
169            VendorId::Kilo => "Kilo",
170            VendorId::Novita => "Novita",
171            VendorId::Moonshot => "Moonshot",
172            VendorId::Grok => "Grok",
173            VendorId::Supergrok => "SuperGrok",
174            VendorId::Antigravity => "Antigravity",
175            VendorId::Cursor => "Cursor",
176            VendorId::Minimax => "MiniMax",
177            VendorId::Kiro => "Kiro",
178            VendorId::NousResearch => "Nous Research",
179            VendorId::OpenCodeGo => "OpenCode Go",
180        }
181    }
182
183    pub fn all() -> &'static [VendorId] {
184        &[
185            VendorId::Anthropic,
186            VendorId::AnthropicApi,
187            VendorId::Openai,
188            VendorId::Zai,
189            VendorId::Openrouter,
190            VendorId::Deepseek,
191            VendorId::Kimi,
192            VendorId::Kilo,
193            VendorId::Novita,
194            VendorId::Moonshot,
195            VendorId::Grok,
196            VendorId::Supergrok,
197            VendorId::Antigravity,
198            VendorId::Cursor,
199            VendorId::Minimax,
200            VendorId::Kiro,
201            VendorId::NousResearch,
202            VendorId::OpenCodeGo,
203        ]
204    }
205}
206
207/// What a vendor returns from a successful fetch — snapshot + meta. Mirrors
208/// `anthropic::fetch::FetchOutcome` but vendor-agnostic.
209#[derive(Debug, Clone)]
210pub struct VendorOutcome {
211    pub snapshot: VendorSnapshot,
212    pub stale: bool,
213    pub last_error: Option<(u16, String)>,
214    pub cache_age: Option<std::time::Duration>,
215}
216
217/// Options forwarded to renderers from the CLI.
218#[derive(Debug, Clone)]
219pub struct RenderOpts {
220    pub format: Option<String>,
221    pub tooltip_format: Option<String>,
222    pub icon: Option<String>,
223    pub pace_tolerance: u32,
224    pub format_pace_color: bool,
225    pub tooltip_pace_pts: bool,
226}
227
228impl RenderOpts {
229    pub fn from_cli(cli: &Cli) -> Self {
230        Self {
231            format: cli.format.clone(),
232            tooltip_format: cli.tooltip_format.clone(),
233            icon: cli.icon.clone(),
234            pace_tolerance: cli.pace_tolerance,
235            format_pace_color: cli.format_pace_color,
236            tooltip_pace_pts: cli.tooltip_pace_pts,
237        }
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244
245    #[test]
246    fn every_vendor_has_stable_machine_and_display_names() {
247        for vendor in VendorId::all() {
248            assert!(!vendor.slug().is_empty());
249            assert!(!vendor.display_name().is_empty());
250        }
251        assert_eq!(VendorId::Anthropic.slug(), "anthropic");
252        assert_eq!(VendorId::Anthropic.display_name(), "Claude");
253        assert_eq!(VendorId::Openai.display_name(), "Codex");
254        assert_eq!(VendorId::Zai.display_name(), "Z.AI");
255    }
256
257    #[test]
258    fn new_vendor_contracts_keep_public_names_and_slugs() {
259        assert_eq!(VendorId::NousResearch.slug(), "nous");
260        assert_eq!(VendorId::NousResearch.display_name(), "Nous Research");
261        assert_eq!(VendorId::OpenCodeGo.slug(), "opencode-go");
262        assert_eq!(VendorId::OpenCodeGo.display_name(), "OpenCode Go");
263        assert_eq!(
264            serde_json::to_value(VendorId::OpenCodeGo).unwrap(),
265            serde_json::json!("opencode-go")
266        );
267    }
268
269    #[test]
270    fn vendor_secret_env_vars_cover_config_defaults() {
271        let configured_defaults = [
272            "ZAI_API_KEY",
273            "OPENROUTER_API_KEY",
274            "DEEPSEEK_API_KEY",
275            "KIMI_API_KEY",
276            "KILO_API_KEY",
277            "NOVITA_API_KEY",
278            "MINIMAX_API_KEY",
279            "MOONSHOT_API_KEY",
280            "XAI_MANAGEMENT_KEY",
281            "ANTHROPIC_ADMIN_KEY",
282        ];
283        for name in configured_defaults {
284            assert!(VENDOR_SECRET_ENV_VARS.contains(&name), "missing {name}");
285        }
286    }
287
288    #[test]
289    fn vars_to_remove_preserves_only_requested_grok_credentials() {
290        let removed = vendor_secret_env_vars_to_remove(&["XAI_API_KEY", "GROK_API_KEY"]);
291        assert!(!removed.contains(&"XAI_API_KEY"));
292        assert!(!removed.contains(&"GROK_API_KEY"));
293        assert!(removed.contains(&"ANTHROPIC_ADMIN_KEY"));
294        assert!(removed.contains(&"OPENROUTER_API_KEY"));
295        assert_eq!(removed.len(), VENDOR_SECRET_ENV_VARS.len() - 2);
296    }
297
298    #[tokio::test]
299    async fn body_over_the_cap_is_refused_and_under_it_round_trips() {
300        let mut server = mockito::Server::new_async().await;
301        server
302            .mock("GET", "/big")
303            .with_status(200)
304            .with_body("x".repeat(4096))
305            .create_async()
306            .await;
307        server
308            .mock("GET", "/small")
309            .with_status(200)
310            .with_body("hello")
311            .create_async()
312            .await;
313
314        let client = reqwest::Client::new();
315
316        // Over the cap: refused rather than buffered.
317        let resp = client
318            .get(format!("{}/big", server.url()))
319            .send()
320            .await
321            .unwrap();
322        let err = read_body_capped(resp, 1024).await.unwrap_err();
323        assert!(
324            err.to_string().contains("exceeds"),
325            "unexpected error: {err}"
326        );
327
328        // Under the cap: identical to the previous `resp.bytes()` behaviour.
329        let resp = client
330            .get(format!("{}/small", server.url()))
331            .send()
332            .await
333            .unwrap();
334        assert_eq!(read_body_capped(resp, 1024).await.unwrap(), b"hello");
335    }
336
337    #[tokio::test]
338    async fn chunked_body_without_content_length_still_hits_the_cap() {
339        let mut server = mockito::Server::new_async().await;
340        server
341            .mock("GET", "/chunked")
342            .with_status(200)
343            .with_chunked_body(|writer| writer.write_all(&[b'x'; 4096]))
344            .create_async()
345            .await;
346
347        let response = reqwest::Client::new()
348            .get(format!("{}/chunked", server.url()))
349            .send()
350            .await
351            .unwrap();
352        assert!(response.content_length().is_none());
353        let error = read_body_capped(response, 1024).await.unwrap_err();
354        assert!(error.to_string().contains("exceeds"), "{error}");
355    }
356
357    #[tokio::test]
358    async fn same_origin_redirects_still_work_with_vendor_headers() {
359        let mut server = mockito::Server::new_async().await;
360        let redirect = server
361            .mock("GET", "/start")
362            .match_header("x-api-key", "secret")
363            .with_status(302)
364            .with_header("location", "/finish")
365            .create_async()
366            .await;
367        let finish = server
368            .mock("GET", "/finish")
369            .match_header("x-api-key", "secret")
370            .with_status(200)
371            .create_async()
372            .await;
373        let client = reqwest::Client::builder()
374            .redirect(same_origin_redirect_policy())
375            .build()
376            .unwrap();
377
378        let response = client
379            .get(format!("{}/start", server.url()))
380            .header("x-api-key", "secret")
381            .send()
382            .await
383            .unwrap();
384
385        assert_eq!(response.status(), reqwest::StatusCode::OK);
386        redirect.assert_async().await;
387        finish.assert_async().await;
388    }
389
390    #[tokio::test]
391    async fn cross_origin_redirects_are_not_followed_with_vendor_headers() {
392        let mut origin = mockito::Server::new_async().await;
393        let mut target = mockito::Server::new_async().await;
394        let target_url = format!("{}/capture", target.url());
395        let redirect = origin
396            .mock("GET", "/start")
397            .match_header("x-api-key", "secret")
398            .with_status(302)
399            .with_header("location", &target_url)
400            .create_async()
401            .await;
402        let capture = target
403            .mock("GET", "/capture")
404            .expect(0)
405            .create_async()
406            .await;
407        let client = reqwest::Client::builder()
408            .redirect(same_origin_redirect_policy())
409            .build()
410            .unwrap();
411
412        let response = client
413            .get(format!("{}/start", origin.url()))
414            .header("x-api-key", "secret")
415            .send()
416            .await
417            .unwrap();
418
419        assert_eq!(response.status(), reqwest::StatusCode::FOUND);
420        redirect.assert_async().await;
421        capture.assert_async().await;
422    }
423
424    #[test]
425    fn vendor_id_slug_round_trip() {
426        for id in VendorId::all() {
427            assert_eq!(
428                id.slug(),
429                serde_json::to_value(id).unwrap().as_str().unwrap()
430            );
431        }
432    }
433}
ai_usagebar/vendor.rs

ai_usagebar/
vendor.rs
