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    /// Compact three-letter code for the bar. This is the single source for
184    /// `{vendor_short}` in every renderer, the `usage --json` `short_name`
185    /// field, and any frontend that wants a Waybar-style provider tag; a
186    /// second copy in a placeholder map or a QML file is how the table forks.
187    pub const fn short_name(self) -> &'static str {
188        match self {
189            VendorId::Anthropic => "cld",
190            VendorId::AnthropicApi => "aac",
191            VendorId::Openai => "gpt",
192            VendorId::Zai => "zai",
193            VendorId::Openrouter => "opr",
194            VendorId::Deepseek => "dsk",
195            VendorId::Kimi => "kmi",
196            VendorId::Kilo => "klo",
197            VendorId::Novita => "nvt",
198            VendorId::Moonshot => "msh",
199            VendorId::Grok => "grk",
200            VendorId::Supergrok => "sgk",
201            VendorId::Antigravity => "agy",
202            VendorId::Cursor => "cur",
203            VendorId::Minimax => "mmx",
204            VendorId::Kiro => "kir",
205            VendorId::NousResearch => "nrs",
206            VendorId::OpenCodeGo => "ocg",
207        }
208    }
209
210    pub fn all() -> &'static [VendorId] {
211        &[
212            VendorId::Anthropic,
213            VendorId::AnthropicApi,
214            VendorId::Openai,
215            VendorId::Zai,
216            VendorId::Openrouter,
217            VendorId::Deepseek,
218            VendorId::Kimi,
219            VendorId::Kilo,
220            VendorId::Novita,
221            VendorId::Moonshot,
222            VendorId::Grok,
223            VendorId::Supergrok,
224            VendorId::Antigravity,
225            VendorId::Cursor,
226            VendorId::Minimax,
227            VendorId::Kiro,
228            VendorId::NousResearch,
229            VendorId::OpenCodeGo,
230        ]
231    }
232}
233
234/// What a vendor returns from a successful fetch — snapshot + meta. Mirrors
235/// `anthropic::fetch::FetchOutcome` but vendor-agnostic.
236#[derive(Debug, Clone)]
237pub struct VendorOutcome {
238    pub snapshot: VendorSnapshot,
239    pub stale: bool,
240    pub last_error: Option<(u16, String)>,
241    pub cache_age: Option<std::time::Duration>,
242}
243
244/// Options forwarded to renderers from the CLI.
245#[derive(Debug, Clone)]
246pub struct RenderOpts {
247    pub format: Option<String>,
248    pub tooltip_format: Option<String>,
249    pub icon: Option<String>,
250    pub pace_tolerance: u32,
251    pub format_pace_color: bool,
252    pub tooltip_pace_pts: bool,
253}
254
255impl RenderOpts {
256    pub fn from_cli(cli: &Cli) -> Self {
257        Self {
258            format: cli.format.clone(),
259            tooltip_format: cli.tooltip_format.clone(),
260            icon: cli.icon.clone(),
261            pace_tolerance: cli.pace_tolerance,
262            format_pace_color: cli.format_pace_color,
263            tooltip_pace_pts: cli.tooltip_pace_pts,
264        }
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271
272    #[test]
273    fn every_vendor_has_stable_machine_and_display_names() {
274        for vendor in VendorId::all() {
275            assert!(!vendor.slug().is_empty());
276            assert!(!vendor.display_name().is_empty());
277        }
278        assert_eq!(VendorId::Anthropic.slug(), "anthropic");
279        assert_eq!(VendorId::Anthropic.display_name(), "Claude");
280        assert_eq!(VendorId::Openai.display_name(), "Codex");
281        assert_eq!(VendorId::Zai.display_name(), "Z.AI");
282    }
283
284    /// `{vendor_short}` is a documented format placeholder and now also rides
285    /// the `usage --json` report, so a duplicate or a re-typed code would make
286    /// two providers indistinguishable in a bar that shows nothing else.
287    #[test]
288    fn every_vendor_short_name_is_a_unique_three_letter_code() {
289        let mut seen = std::collections::BTreeSet::new();
290        for vendor in VendorId::all() {
291            let short = vendor.short_name();
292            assert_eq!(short.len(), 3, "{} is not three letters", vendor.slug());
293            assert!(
294                short.chars().all(|c| c.is_ascii_lowercase()),
295                "{} is not lowercase ascii",
296                vendor.slug()
297            );
298            assert!(seen.insert(short), "{short} is used by two vendors");
299        }
300        assert_eq!(VendorId::Anthropic.short_name(), "cld");
301        assert_eq!(VendorId::Openai.short_name(), "gpt");
302        assert_eq!(VendorId::Zai.short_name(), "zai");
303        assert_eq!(VendorId::Antigravity.short_name(), "agy");
304    }
305
306    #[test]
307    fn new_vendor_contracts_keep_public_names_and_slugs() {
308        assert_eq!(VendorId::NousResearch.slug(), "nous");
309        assert_eq!(VendorId::NousResearch.display_name(), "Nous Research");
310        assert_eq!(VendorId::OpenCodeGo.slug(), "opencode-go");
311        assert_eq!(VendorId::OpenCodeGo.display_name(), "OpenCode Go");
312        assert_eq!(
313            serde_json::to_value(VendorId::OpenCodeGo).unwrap(),
314            serde_json::json!("opencode-go")
315        );
316    }
317
318    #[test]
319    fn vendor_secret_env_vars_cover_config_defaults() {
320        let configured_defaults = [
321            "ZAI_API_KEY",
322            "OPENROUTER_API_KEY",
323            "DEEPSEEK_API_KEY",
324            "KIMI_API_KEY",
325            "KILO_API_KEY",
326            "NOVITA_API_KEY",
327            "MINIMAX_API_KEY",
328            "MOONSHOT_API_KEY",
329            "XAI_MANAGEMENT_KEY",
330            "ANTHROPIC_ADMIN_KEY",
331        ];
332        for name in configured_defaults {
333            assert!(VENDOR_SECRET_ENV_VARS.contains(&name), "missing {name}");
334        }
335    }
336
337    #[test]
338    fn vars_to_remove_preserves_only_requested_grok_credentials() {
339        let removed = vendor_secret_env_vars_to_remove(&["XAI_API_KEY", "GROK_API_KEY"]);
340        assert!(!removed.contains(&"XAI_API_KEY"));
341        assert!(!removed.contains(&"GROK_API_KEY"));
342        assert!(removed.contains(&"ANTHROPIC_ADMIN_KEY"));
343        assert!(removed.contains(&"OPENROUTER_API_KEY"));
344        assert_eq!(removed.len(), VENDOR_SECRET_ENV_VARS.len() - 2);
345    }
346
347    #[tokio::test]
348    async fn body_over_the_cap_is_refused_and_under_it_round_trips() {
349        let mut server = mockito::Server::new_async().await;
350        server
351            .mock("GET", "/big")
352            .with_status(200)
353            .with_body("x".repeat(4096))
354            .create_async()
355            .await;
356        server
357            .mock("GET", "/small")
358            .with_status(200)
359            .with_body("hello")
360            .create_async()
361            .await;
362
363        let client = reqwest::Client::new();
364
365        // Over the cap: refused rather than buffered.
366        let resp = client
367            .get(format!("{}/big", server.url()))
368            .send()
369            .await
370            .unwrap();
371        let err = read_body_capped(resp, 1024).await.unwrap_err();
372        assert!(
373            err.to_string().contains("exceeds"),
374            "unexpected error: {err}"
375        );
376
377        // Under the cap: identical to the previous `resp.bytes()` behaviour.
378        let resp = client
379            .get(format!("{}/small", server.url()))
380            .send()
381            .await
382            .unwrap();
383        assert_eq!(read_body_capped(resp, 1024).await.unwrap(), b"hello");
384    }
385
386    #[tokio::test]
387    async fn chunked_body_without_content_length_still_hits_the_cap() {
388        let mut server = mockito::Server::new_async().await;
389        server
390            .mock("GET", "/chunked")
391            .with_status(200)
392            .with_chunked_body(|writer| writer.write_all(&[b'x'; 4096]))
393            .create_async()
394            .await;
395
396        let response = reqwest::Client::new()
397            .get(format!("{}/chunked", server.url()))
398            .send()
399            .await
400            .unwrap();
401        assert!(response.content_length().is_none());
402        let error = read_body_capped(response, 1024).await.unwrap_err();
403        assert!(error.to_string().contains("exceeds"), "{error}");
404    }
405
406    #[tokio::test]
407    async fn same_origin_redirects_still_work_with_vendor_headers() {
408        let mut server = mockito::Server::new_async().await;
409        let redirect = server
410            .mock("GET", "/start")
411            .match_header("x-api-key", "secret")
412            .with_status(302)
413            .with_header("location", "/finish")
414            .create_async()
415            .await;
416        let finish = server
417            .mock("GET", "/finish")
418            .match_header("x-api-key", "secret")
419            .with_status(200)
420            .create_async()
421            .await;
422        let client = reqwest::Client::builder()
423            .redirect(same_origin_redirect_policy())
424            .build()
425            .unwrap();
426
427        let response = client
428            .get(format!("{}/start", server.url()))
429            .header("x-api-key", "secret")
430            .send()
431            .await
432            .unwrap();
433
434        assert_eq!(response.status(), reqwest::StatusCode::OK);
435        redirect.assert_async().await;
436        finish.assert_async().await;
437    }
438
439    #[tokio::test]
440    async fn cross_origin_redirects_are_not_followed_with_vendor_headers() {
441        let mut origin = mockito::Server::new_async().await;
442        let mut target = mockito::Server::new_async().await;
443        let target_url = format!("{}/capture", target.url());
444        let redirect = origin
445            .mock("GET", "/start")
446            .match_header("x-api-key", "secret")
447            .with_status(302)
448            .with_header("location", &target_url)
449            .create_async()
450            .await;
451        let capture = target
452            .mock("GET", "/capture")
453            .expect(0)
454            .create_async()
455            .await;
456        let client = reqwest::Client::builder()
457            .redirect(same_origin_redirect_policy())
458            .build()
459            .unwrap();
460
461        let response = client
462            .get(format!("{}/start", origin.url()))
463            .header("x-api-key", "secret")
464            .send()
465            .await
466            .unwrap();
467
468        assert_eq!(response.status(), reqwest::StatusCode::FOUND);
469        redirect.assert_async().await;
470        capture.assert_async().await;
471    }
472
473    #[test]
474    fn vendor_id_slug_round_trip() {
475        for id in VendorId::all() {
476            assert_eq!(
477                id.slug(),
478                serde_json::to_value(id).unwrap().as_str().unwrap()
479            );
480        }
481    }
482}
ai_usagebar/vendor.rs

ai_usagebar/
vendor.rs
