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    "COMMANDCODE_API_KEY",
40    "GITHUB_COPILOT_TOKEN",
41    "GH_TOKEN",
42    "GITHUB_TOKEN",
43];
44
45pub(crate) fn vendor_secret_env_vars_to_remove(keep: &[&str]) -> Vec<&'static str> {
46    VENDOR_SECRET_ENV_VARS
47        .iter()
48        .copied()
49        .filter(|var| !keep.contains(var))
50        .collect()
51}
52
53/// Follow ordinary vendor redirects without forwarding non-standard API-key
54/// headers to a different origin. Reqwest strips `Authorization` on sensitive
55/// redirects, but vendors also use headers such as `x-api-key`, which are not
56/// covered by that built-in list.
57pub fn same_origin_redirect_policy() -> reqwest::redirect::Policy {
58    reqwest::redirect::Policy::custom(|attempt| {
59        if attempt.previous().len() >= 10 {
60            return attempt.error("too many redirects");
61        }
62        let Some(origin) = attempt.previous().first() else {
63            return attempt.stop();
64        };
65        let target = attempt.url();
66        if target.scheme() == origin.scheme()
67            && target.host_str() == origin.host_str()
68            && target.port_or_known_default() == origin.port_or_known_default()
69        {
70            attempt.follow()
71        } else {
72            attempt.stop()
73        }
74    })
75}
76
77/// Read a response body with an upper bound.
78///
79/// Every vendor buffered the whole body with `resp.bytes()` *before* anything
80/// validated it. The widget is re-executed by Waybar every 60s, so an endpoint
81/// answering with an unbounded stream had a free hand at the machine's memory.
82/// `Content-Length` is checked first when present, then the body is read in
83/// chunks so a lying or absent length cannot get past the cap either.
84pub async fn read_body_capped(
85    mut resp: reqwest::Response,
86    max: usize,
87) -> crate::error::Result<Vec<u8>> {
88    let too_big = |n: u64| {
89        crate::error::AppError::Schema(format!(
90            "response body exceeds the {max}-byte limit ({n} bytes); refusing to buffer it"
91        ))
92    };
93    if let Some(len) = resp.content_length()
94        && len > max as u64
95    {
96        return Err(too_big(len));
97    }
98    let mut buf: Vec<u8> = Vec::new();
99    while let Some(chunk) = resp.chunk().await? {
100        if chunk.len() > max.saturating_sub(buf.len()) {
101            return Err(too_big(buf.len().saturating_add(chunk.len()) as u64));
102        }
103        buf.extend_from_slice(&chunk);
104    }
105    Ok(buf)
106}
107
108/// Stable enum used by `--vendor` and in config files.
109#[derive(
110    Debug, Clone, Copy, ValueEnum, PartialEq, Eq, Hash, serde::Deserialize, serde::Serialize,
111)]
112#[serde(rename_all = "lowercase")]
113pub enum VendorId {
114    Anthropic,
115    #[serde(rename = "anthropic_api")]
116    AnthropicApi,
117    Openai,
118    Copilot,
119    Zai,
120    Openrouter,
121    Deepseek,
122    Kimi,
123    Kilo,
124    Novita,
125    Moonshot,
126    Grok,
127    Supergrok,
128    Antigravity,
129    Cursor,
130    Minimax,
131    Kiro,
132    #[serde(rename = "nous")]
133    NousResearch,
134    #[serde(rename = "opencode-go")]
135    OpenCodeGo,
136    #[serde(rename = "commandcode")]
137    CommandCode,
138}
139
140impl VendorId {
141    pub fn slug(self) -> &'static str {
142        match self {
143            VendorId::Anthropic => "anthropic",
144            VendorId::AnthropicApi => "anthropic_api",
145            VendorId::Openai => "openai",
146            VendorId::Copilot => "copilot",
147            VendorId::Zai => "zai",
148            VendorId::Openrouter => "openrouter",
149            VendorId::Deepseek => "deepseek",
150            VendorId::Kimi => "kimi",
151            VendorId::Kilo => "kilo",
152            VendorId::Novita => "novita",
153            VendorId::Moonshot => "moonshot",
154            VendorId::Grok => "grok",
155            VendorId::Supergrok => "supergrok",
156            VendorId::Antigravity => "antigravity",
157            VendorId::Cursor => "cursor",
158            VendorId::Minimax => "minimax",
159            VendorId::Kiro => "kiro",
160            VendorId::NousResearch => "nous",
161            VendorId::OpenCodeGo => "opencode-go",
162            VendorId::CommandCode => "commandcode",
163        }
164    }
165
166    /// Canonical human-readable name for shared reports and compact UI labels.
167    /// Platform frontends may add context (for example, "GLM (Z.AI)" in a
168    /// wide TUI tab), but should not carry their own full vendor-name table.
169    pub fn display_name(self) -> &'static str {
170        match self {
171            VendorId::Anthropic => "Claude",
172            VendorId::AnthropicApi => "Anthropic API",
173            VendorId::Openai => "Codex",
174            VendorId::Copilot => "GitHub Copilot",
175            VendorId::Zai => "Z.AI",
176            VendorId::Openrouter => "OpenRouter",
177            VendorId::Deepseek => "DeepSeek",
178            VendorId::Kimi => "Kimi",
179            VendorId::Kilo => "Kilo",
180            VendorId::Novita => "Novita",
181            VendorId::Moonshot => "Moonshot",
182            VendorId::Grok => "Grok",
183            VendorId::Supergrok => "SuperGrok",
184            VendorId::Antigravity => "Antigravity",
185            VendorId::Cursor => "Cursor",
186            VendorId::Minimax => "MiniMax",
187            VendorId::Kiro => "Kiro",
188            VendorId::NousResearch => "Nous Research",
189            VendorId::OpenCodeGo => "OpenCode Go",
190            VendorId::CommandCode => "Command Code",
191        }
192    }
193
194    /// Compact three-letter code for the bar. This is the single source for
195    /// `{vendor_short}` in every renderer, the `usage --json` `short_name`
196    /// field, and any frontend that wants a Waybar-style provider tag; a
197    /// second copy in a placeholder map or a QML file is how the table forks.
198    pub const fn short_name(self) -> &'static str {
199        match self {
200            VendorId::Anthropic => "cld",
201            VendorId::AnthropicApi => "aac",
202            VendorId::Openai => "gpt",
203            VendorId::Copilot => "ghc",
204            VendorId::Zai => "zai",
205            VendorId::Openrouter => "opr",
206            VendorId::Deepseek => "dsk",
207            VendorId::Kimi => "kmi",
208            VendorId::Kilo => "klo",
209            VendorId::Novita => "nvt",
210            VendorId::Moonshot => "msh",
211            VendorId::Grok => "grk",
212            VendorId::Supergrok => "sgk",
213            VendorId::Antigravity => "agy",
214            VendorId::Cursor => "cur",
215            VendorId::Minimax => "mmx",
216            VendorId::Kiro => "kir",
217            VendorId::NousResearch => "nrs",
218            VendorId::OpenCodeGo => "ocg",
219            VendorId::CommandCode => "cmc",
220        }
221    }
222
223    pub fn all() -> &'static [VendorId] {
224        &[
225            VendorId::Anthropic,
226            VendorId::AnthropicApi,
227            VendorId::Openai,
228            VendorId::Copilot,
229            VendorId::Zai,
230            VendorId::Openrouter,
231            VendorId::Deepseek,
232            VendorId::Kimi,
233            VendorId::Kilo,
234            VendorId::Novita,
235            VendorId::Moonshot,
236            VendorId::Grok,
237            VendorId::Supergrok,
238            VendorId::Antigravity,
239            VendorId::Cursor,
240            VendorId::Minimax,
241            VendorId::Kiro,
242            VendorId::NousResearch,
243            VendorId::OpenCodeGo,
244            VendorId::CommandCode,
245        ]
246    }
247}
248
249/// What a vendor returns from a successful fetch — the same
250/// [`Outcome`](crate::outcome::Outcome) every vendor produces, once its own
251/// snapshot type has been widened to [`VendorSnapshot`]. Each vendor gets
252/// there with a single `outcome.map(VendorSnapshot::Whichever)`.
253pub type VendorOutcome = crate::outcome::Outcome<VendorSnapshot>;
254
255/// Options forwarded to renderers from the CLI.
256#[derive(Debug, Clone)]
257pub struct RenderOpts {
258    pub format: Option<String>,
259    pub tooltip_format: Option<String>,
260    pub icon: Option<String>,
261    pub pace_tolerance: u32,
262    pub format_pace_color: bool,
263    pub tooltip_pace_pts: bool,
264}
265
266impl RenderOpts {
267    pub fn from_cli(cli: &Cli) -> Self {
268        Self {
269            format: cli.format.clone(),
270            tooltip_format: cli.tooltip_format.clone(),
271            icon: cli.icon.clone(),
272            pace_tolerance: cli.pace_tolerance,
273            format_pace_color: cli.format_pace_color,
274            tooltip_pace_pts: cli.tooltip_pace_pts,
275        }
276    }
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282
283    #[test]
284    fn every_vendor_has_stable_machine_and_display_names() {
285        for vendor in VendorId::all() {
286            assert!(!vendor.slug().is_empty());
287            assert!(!vendor.display_name().is_empty());
288        }
289        assert_eq!(VendorId::Anthropic.slug(), "anthropic");
290        assert_eq!(VendorId::Anthropic.display_name(), "Claude");
291        assert_eq!(VendorId::Openai.display_name(), "Codex");
292        assert_eq!(VendorId::Zai.display_name(), "Z.AI");
293    }
294
295    /// `{vendor_short}` is a documented format placeholder and now also rides
296    /// the `usage --json` report, so a duplicate or a re-typed code would make
297    /// two providers indistinguishable in a bar that shows nothing else.
298    #[test]
299    fn every_vendor_short_name_is_a_unique_three_letter_code() {
300        let mut seen = std::collections::BTreeSet::new();
301        for vendor in VendorId::all() {
302            let short = vendor.short_name();
303            assert_eq!(short.len(), 3, "{} is not three letters", vendor.slug());
304            assert!(
305                short.chars().all(|c| c.is_ascii_lowercase()),
306                "{} is not lowercase ascii",
307                vendor.slug()
308            );
309            assert!(seen.insert(short), "{short} is used by two vendors");
310        }
311        assert_eq!(VendorId::Anthropic.short_name(), "cld");
312        assert_eq!(VendorId::Openai.short_name(), "gpt");
313        assert_eq!(VendorId::Zai.short_name(), "zai");
314        assert_eq!(VendorId::Antigravity.short_name(), "agy");
315    }
316
317    #[test]
318    fn new_vendor_contracts_keep_public_names_and_slugs() {
319        assert_eq!(VendorId::NousResearch.slug(), "nous");
320        assert_eq!(VendorId::NousResearch.display_name(), "Nous Research");
321        assert_eq!(VendorId::OpenCodeGo.slug(), "opencode-go");
322        assert_eq!(VendorId::OpenCodeGo.display_name(), "OpenCode Go");
323        assert_eq!(
324            serde_json::to_value(VendorId::OpenCodeGo).unwrap(),
325            serde_json::json!("opencode-go")
326        );
327    }
328
329    #[test]
330    fn vendor_secret_env_vars_cover_config_defaults() {
331        let configured_defaults = [
332            "ZAI_API_KEY",
333            "OPENROUTER_API_KEY",
334            "DEEPSEEK_API_KEY",
335            "KIMI_API_KEY",
336            "KILO_API_KEY",
337            "NOVITA_API_KEY",
338            "MINIMAX_API_KEY",
339            "MOONSHOT_API_KEY",
340            "XAI_MANAGEMENT_KEY",
341            "ANTHROPIC_ADMIN_KEY",
342            "GITHUB_COPILOT_TOKEN",
343        ];
344        for name in configured_defaults {
345            assert!(VENDOR_SECRET_ENV_VARS.contains(&name), "missing {name}");
346        }
347    }
348
349    #[test]
350    fn vars_to_remove_preserves_only_requested_grok_credentials() {
351        let removed = vendor_secret_env_vars_to_remove(&["XAI_API_KEY", "GROK_API_KEY"]);
352        assert!(!removed.contains(&"XAI_API_KEY"));
353        assert!(!removed.contains(&"GROK_API_KEY"));
354        assert!(removed.contains(&"ANTHROPIC_ADMIN_KEY"));
355        assert!(removed.contains(&"OPENROUTER_API_KEY"));
356        assert_eq!(removed.len(), VENDOR_SECRET_ENV_VARS.len() - 2);
357    }
358
359    #[test]
360    fn copilot_token_is_removed_before_unrelated_subprocesses_launch() {
361        let removed = vendor_secret_env_vars_to_remove(&[]);
362        assert!(removed.contains(&"GITHUB_COPILOT_TOKEN"));
363    }
364
365    #[tokio::test]
366    async fn body_over_the_cap_is_refused_and_under_it_round_trips() {
367        let mut server = mockito::Server::new_async().await;
368        server
369            .mock("GET", "/big")
370            .with_status(200)
371            .with_body("x".repeat(4096))
372            .create_async()
373            .await;
374        server
375            .mock("GET", "/small")
376            .with_status(200)
377            .with_body("hello")
378            .create_async()
379            .await;
380
381        let client = reqwest::Client::new();
382
383        // Over the cap: refused rather than buffered.
384        let resp = client
385            .get(format!("{}/big", server.url()))
386            .send()
387            .await
388            .unwrap();
389        let err = read_body_capped(resp, 1024).await.unwrap_err();
390        assert!(
391            err.to_string().contains("exceeds"),
392            "unexpected error: {err}"
393        );
394
395        // Under the cap: identical to the previous `resp.bytes()` behaviour.
396        let resp = client
397            .get(format!("{}/small", server.url()))
398            .send()
399            .await
400            .unwrap();
401        assert_eq!(read_body_capped(resp, 1024).await.unwrap(), b"hello");
402    }
403
404    #[tokio::test]
405    async fn chunked_body_without_content_length_still_hits_the_cap() {
406        let mut server = mockito::Server::new_async().await;
407        server
408            .mock("GET", "/chunked")
409            .with_status(200)
410            .with_chunked_body(|writer| writer.write_all(&[b'x'; 4096]))
411            .create_async()
412            .await;
413
414        let response = reqwest::Client::new()
415            .get(format!("{}/chunked", server.url()))
416            .send()
417            .await
418            .unwrap();
419        assert!(response.content_length().is_none());
420        let error = read_body_capped(response, 1024).await.unwrap_err();
421        assert!(error.to_string().contains("exceeds"), "{error}");
422    }
423
424    #[tokio::test]
425    async fn same_origin_redirects_still_work_with_vendor_headers() {
426        let mut server = mockito::Server::new_async().await;
427        let redirect = server
428            .mock("GET", "/start")
429            .match_header("x-api-key", "secret")
430            .with_status(302)
431            .with_header("location", "/finish")
432            .create_async()
433            .await;
434        let finish = server
435            .mock("GET", "/finish")
436            .match_header("x-api-key", "secret")
437            .with_status(200)
438            .create_async()
439            .await;
440        let client = reqwest::Client::builder()
441            .redirect(same_origin_redirect_policy())
442            .build()
443            .unwrap();
444
445        let response = client
446            .get(format!("{}/start", server.url()))
447            .header("x-api-key", "secret")
448            .send()
449            .await
450            .unwrap();
451
452        assert_eq!(response.status(), reqwest::StatusCode::OK);
453        redirect.assert_async().await;
454        finish.assert_async().await;
455    }
456
457    #[tokio::test]
458    async fn cross_origin_redirects_are_not_followed_with_vendor_headers() {
459        let mut origin = mockito::Server::new_async().await;
460        let mut target = mockito::Server::new_async().await;
461        let target_url = format!("{}/capture", target.url());
462        let redirect = origin
463            .mock("GET", "/start")
464            .match_header("x-api-key", "secret")
465            .with_status(302)
466            .with_header("location", &target_url)
467            .create_async()
468            .await;
469        let capture = target
470            .mock("GET", "/capture")
471            .expect(0)
472            .create_async()
473            .await;
474        let client = reqwest::Client::builder()
475            .redirect(same_origin_redirect_policy())
476            .build()
477            .unwrap();
478
479        let response = client
480            .get(format!("{}/start", origin.url()))
481            .header("x-api-key", "secret")
482            .send()
483            .await
484            .unwrap();
485
486        assert_eq!(response.status(), reqwest::StatusCode::FOUND);
487        redirect.assert_async().await;
488        capture.assert_async().await;
489    }
490
491    #[test]
492    fn vendor_id_slug_round_trip() {
493        for id in VendorId::all() {
494            assert_eq!(
495                id.slug(),
496                serde_json::to_value(id).unwrap().as_str().unwrap()
497            );
498        }
499    }
500}
ai_usagebar/vendor.rs

ai_usagebar/
vendor.rs
