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

ai_usagebar/
vendor.rs
