Skip to main content

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::collections::BTreeSet;
7use std::sync::{Mutex, OnceLock};
8use std::time::Duration;
9
10use clap::ValueEnum;
11
12use crate::usage::VendorSnapshot;
13use crate::widget::cli::Cli;
14
15/// Outer reqwest client timeout shared by widget and TUI entry points.
16/// Vendor fetchers still apply their own tighter per-request timeouts.
17pub const HTTP_CLIENT_TIMEOUT: Duration = Duration::from_secs(30);
18
19/// Upper bound on a vendor response body. Every one of these endpoints returns
20/// a small JSON document — the largest observed is a few kilobytes — so this is
21/// generous by three orders of magnitude while still bounding the damage from a
22/// misbehaving proxy or a hijacked endpoint.
23pub const MAX_BODY_BYTES: usize = 2 * 1024 * 1024;
24
25/// Credential-bearing environment variables owned by ai-usagebar vendors.
26/// Subprocesses receive only the entries that belong to their own provider.
27pub(crate) const VENDOR_SECRET_ENV_VARS: &[&str] = &[
28    "ZAI_API_KEY",
29    "OPENROUTER_API_KEY",
30    "DEEPSEEK_API_KEY",
31    "KIMI_API_KEY",
32    "KILO_API_KEY",
33    "NOVITA_API_KEY",
34    "MINIMAX_API_KEY",
35    "MOONSHOT_API_KEY",
36    "XAI_MANAGEMENT_KEY",
37    "ANTHROPIC_ADMIN_KEY",
38    "XAI_API_KEY",
39    "GROK_API_KEY",
40    "OPENCODE_GO_API_KEY",
41    "COMMANDCODE_API_KEY",
42    "GITHUB_COPILOT_TOKEN",
43    "GH_TOKEN",
44    "GITHUB_TOKEN",
45];
46
47/// Env var names a `[[custom]]` provider reads its token from. They are not
48/// known until the config is parsed, so they cannot sit in the static list
49/// above, but they are exactly as secret as `DEEPSEEK_API_KEY` and must be
50/// scrubbed from every subprocess the same way.
51fn registered_secret_env_vars() -> &'static Mutex<BTreeSet<&'static str>> {
52    static REGISTERED: OnceLock<Mutex<BTreeSet<&'static str>>> = OnceLock::new();
53    REGISTERED.get_or_init(|| Mutex::new(BTreeSet::new()))
54}
55
56/// Extra env var names (custom providers' `api_key_env`) that must be
57/// scrubbed from every child process. Additive and idempotent; names that are
58/// not valid env var names, or already in [`VENDOR_SECRET_ENV_VARS`], are
59/// ignored.
60///
61/// A name is interned once, on first registration, so the removal list keeps
62/// its `&'static str` element type and the three call sites and their tests
63/// stay untouched. The set is bounded by the user's config, and re-loading
64/// the same config registers nothing new, so the leak is a handful of short
65/// strings for the life of the process.
66pub fn register_secret_env_vars(names: &[String]) {
67    let mut registered = registered_secret_env_vars()
68        .lock()
69        .unwrap_or_else(|poisoned| poisoned.into_inner());
70    for name in names {
71        if !crate::config::is_valid_env_var_name(name)
72            || VENDOR_SECRET_ENV_VARS.contains(&name.as_str())
73            || registered.contains(name.as_str())
74        {
75            continue;
76        }
77        registered.insert(Box::leak(name.clone().into_boxed_str()));
78    }
79}
80
81pub(crate) fn vendor_secret_env_vars_to_remove(keep: &[&str]) -> Vec<&'static str> {
82    let registered = registered_secret_env_vars()
83        .lock()
84        .unwrap_or_else(|poisoned| poisoned.into_inner());
85    VENDOR_SECRET_ENV_VARS
86        .iter()
87        .copied()
88        .chain(registered.iter().copied())
89        .filter(|var| !keep.contains(var))
90        .collect()
91}
92
93/// Follow ordinary vendor redirects without forwarding non-standard API-key
94/// headers to a different origin. Reqwest strips `Authorization` on sensitive
95/// redirects, but vendors also use headers such as `x-api-key`, which are not
96/// covered by that built-in list.
97pub fn same_origin_redirect_policy() -> reqwest::redirect::Policy {
98    reqwest::redirect::Policy::custom(|attempt| {
99        if attempt.previous().len() >= 10 {
100            return attempt.error("too many redirects");
101        }
102        let Some(origin) = attempt.previous().first() else {
103            return attempt.stop();
104        };
105        let target = attempt.url();
106        if target.scheme() == origin.scheme()
107            && target.host_str() == origin.host_str()
108            && target.port_or_known_default() == origin.port_or_known_default()
109        {
110            attempt.follow()
111        } else {
112            attempt.stop()
113        }
114    })
115}
116
117/// Read a response body with an upper bound.
118///
119/// Every vendor buffered the whole body with `resp.bytes()` *before* anything
120/// validated it. The widget is re-executed by Waybar every 60s, so an endpoint
121/// answering with an unbounded stream had a free hand at the machine's memory.
122/// `Content-Length` is checked first when present, then the body is read in
123/// chunks so a lying or absent length cannot get past the cap either.
124pub async fn read_body_capped(
125    mut resp: reqwest::Response,
126    max: usize,
127) -> crate::error::Result<Vec<u8>> {
128    let too_big = |n: u64| {
129        crate::error::AppError::Schema(format!(
130            "response body exceeds the {max}-byte limit ({n} bytes); refusing to buffer it"
131        ))
132    };
133    if let Some(len) = resp.content_length()
134        && len > max as u64
135    {
136        return Err(too_big(len));
137    }
138    let mut buf: Vec<u8> = Vec::new();
139    while let Some(chunk) = resp.chunk().await? {
140        if chunk.len() > max.saturating_sub(buf.len()) {
141            return Err(too_big(buf.len().saturating_add(chunk.len()) as u64));
142        }
143        buf.extend_from_slice(&chunk);
144    }
145    Ok(buf)
146}
147
148/// Stable enum used by `--vendor` and in config files.
149#[derive(
150    Debug, Clone, Copy, ValueEnum, PartialEq, Eq, Hash, serde::Deserialize, serde::Serialize,
151)]
152#[serde(rename_all = "lowercase")]
153pub enum VendorId {
154    Anthropic,
155    #[serde(rename = "anthropic_api")]
156    AnthropicApi,
157    Openai,
158    Copilot,
159    Zai,
160    Openrouter,
161    Deepseek,
162    Kimi,
163    Kilo,
164    Novita,
165    Moonshot,
166    Grok,
167    Supergrok,
168    Antigravity,
169    Cursor,
170    Minimax,
171    Kiro,
172    #[serde(rename = "nous")]
173    NousResearch,
174    #[serde(rename = "opencode-go")]
175    OpenCodeGo,
176    #[serde(rename = "commandcode")]
177    CommandCode,
178}
179
180/// How a provider authenticates. Drives what a frontend offers a provider that
181/// is not usable yet: a command to run, a variable to set, or an app to sign
182/// in to.
183#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
184#[serde(rename_all = "lowercase")]
185pub enum AuthKind {
186    /// An interactive login writes a credential file. `login_command` runs it.
187    Oauth,
188    /// An API key, from the environment or an inline `api_key` in config.
189    ApiKey,
190    /// No credential of its own — a local product's session or state file is
191    /// the login, and there is nothing for the user to paste.
192    Local,
193}
194
195impl AuthKind {
196    pub const fn as_str(self) -> &'static str {
197        match self {
198            AuthKind::Oauth => "oauth",
199            AuthKind::ApiKey => "apikey",
200            AuthKind::Local => "local",
201        }
202    }
203}
204
205impl VendorId {
206    pub fn slug(self) -> &'static str {
207        match self {
208            VendorId::Anthropic => "anthropic",
209            VendorId::AnthropicApi => "anthropic_api",
210            VendorId::Openai => "openai",
211            VendorId::Copilot => "copilot",
212            VendorId::Zai => "zai",
213            VendorId::Openrouter => "openrouter",
214            VendorId::Deepseek => "deepseek",
215            VendorId::Kimi => "kimi",
216            VendorId::Kilo => "kilo",
217            VendorId::Novita => "novita",
218            VendorId::Moonshot => "moonshot",
219            VendorId::Grok => "grok",
220            VendorId::Supergrok => "supergrok",
221            VendorId::Antigravity => "antigravity",
222            VendorId::Cursor => "cursor",
223            VendorId::Minimax => "minimax",
224            VendorId::Kiro => "kiro",
225            VendorId::NousResearch => "nous",
226            VendorId::OpenCodeGo => "opencode-go",
227            VendorId::CommandCode => "commandcode",
228        }
229    }
230
231    /// Canonical human-readable name for shared reports and compact UI labels.
232    /// Platform frontends may add context (for example, "GLM (Z.AI)" in a
233    /// wide TUI tab), but should not carry their own full vendor-name table.
234    pub fn display_name(self) -> &'static str {
235        match self {
236            VendorId::Anthropic => "Claude",
237            VendorId::AnthropicApi => "Anthropic API",
238            VendorId::Openai => "Codex",
239            VendorId::Copilot => "GitHub Copilot",
240            VendorId::Zai => "Z.AI",
241            VendorId::Openrouter => "OpenRouter",
242            VendorId::Deepseek => "DeepSeek",
243            VendorId::Kimi => "Kimi",
244            VendorId::Kilo => "Kilo",
245            VendorId::Novita => "Novita",
246            VendorId::Moonshot => "Moonshot",
247            VendorId::Grok => "Grok",
248            VendorId::Supergrok => "SuperGrok",
249            VendorId::Antigravity => "Antigravity",
250            VendorId::Cursor => "Cursor",
251            VendorId::Minimax => "MiniMax",
252            VendorId::Kiro => "Kiro",
253            VendorId::NousResearch => "Nous Research",
254            VendorId::OpenCodeGo => "OpenCode Go",
255            VendorId::CommandCode => "Command Code",
256        }
257    }
258
259    /// Glyph for a compact bar chip. Same role as [`Self::short_name`]: the
260    /// Omarchy top bar (and any other frontend) takes it from `usage --json`
261    /// rather than keeping its own provider-icon table.
262    pub const fn bar_icon(self) -> &'static str {
263        match self {
264            VendorId::Anthropic => "󰚩",
265            VendorId::AnthropicApi => "󰢗",
266            VendorId::Openai => "󱢆",
267            VendorId::Copilot => "󰊤",
268            VendorId::Zai => VendorId::Zai.short_name(),
269            VendorId::Openrouter => "󱙺",
270            VendorId::Deepseek => "󰧑",
271            VendorId::Kimi => VendorId::Kimi.short_name(),
272            VendorId::Kilo => "󰭟",
273            VendorId::Novita => "󰄔",
274            VendorId::Moonshot => VendorId::Moonshot.short_name(),
275            VendorId::Grok | VendorId::Supergrok => "󰇷",
276            VendorId::Antigravity => VendorId::Antigravity.short_name(),
277            VendorId::Cursor => "❯",
278            VendorId::Minimax => VendorId::Minimax.short_name(),
279            VendorId::Kiro => "◆",
280            VendorId::NousResearch => VendorId::NousResearch.short_name(),
281            VendorId::OpenCodeGo => VendorId::OpenCodeGo.short_name(),
282            VendorId::CommandCode => VendorId::CommandCode.short_name(),
283        }
284    }
285
286    /// Compact three-letter code for the bar. This is the single source for
287    /// `{vendor_short}` in every renderer, the `usage --json` `short_name`
288    /// field, and any frontend that wants a Waybar-style provider tag; a
289    /// second copy in a placeholder map or a QML file is how the table forks.
290    pub const fn short_name(self) -> &'static str {
291        match self {
292            VendorId::Anthropic => "cld",
293            VendorId::AnthropicApi => "aac",
294            VendorId::Openai => "gpt",
295            VendorId::Copilot => "ghc",
296            VendorId::Zai => "zai",
297            VendorId::Openrouter => "opr",
298            VendorId::Deepseek => "dsk",
299            VendorId::Kimi => "kmi",
300            VendorId::Kilo => "klo",
301            VendorId::Novita => "nvt",
302            VendorId::Moonshot => "msh",
303            VendorId::Grok => "grk",
304            VendorId::Supergrok => "sgk",
305            VendorId::Antigravity => "agy",
306            VendorId::Cursor => "cur",
307            VendorId::Minimax => "mmx",
308            VendorId::Kiro => "kir",
309            VendorId::NousResearch => "nrs",
310            VendorId::OpenCodeGo => "ocg",
311            VendorId::CommandCode => "cmc",
312        }
313    }
314
315    /// The `config.toml` table this vendor's settings live in — the `Config`
316    /// field name, or its `#[serde(rename)]` where one applies. This is the
317    /// single source for every writer that edits a vendor section by name
318    /// (the Settings overlay's `KEY_VENDORS`, `config::enable_vendors_in`), so
319    /// a section can't be spelled one way by the parser and another by a
320    /// writer. A guard test in `config` parses `[<section>] enabled = true` for
321    /// every vendor and checks `is_enabled`.
322    pub const fn config_section(self) -> &'static str {
323        match self {
324            VendorId::Anthropic => "anthropic",
325            VendorId::AnthropicApi => "anthropic_api",
326            VendorId::Openai => "openai",
327            VendorId::Copilot => "copilot",
328            VendorId::Zai => "zai",
329            VendorId::Openrouter => "openrouter",
330            VendorId::Deepseek => "deepseek",
331            VendorId::Kimi => "kimi",
332            VendorId::Kilo => "kilo",
333            VendorId::Novita => "novita",
334            VendorId::Moonshot => "moonshot",
335            VendorId::Grok => "grok",
336            VendorId::Supergrok => "supergrok",
337            VendorId::Antigravity => "antigravity",
338            VendorId::Cursor => "cursor",
339            VendorId::Minimax => "minimax",
340            VendorId::Kiro => "kiro",
341            VendorId::NousResearch => "nous",
342            VendorId::OpenCodeGo => "opencode-go",
343            VendorId::CommandCode => "commandcode",
344        }
345    }
346
347    /// How a provider proves who you are. This is the fact a frontend needs to
348    /// say what an unconfigured provider is still missing, and it is the one
349    /// thing neither `usage --json` nor the config file carries: the report
350    /// lists only *enabled* providers, so the switched-off and the
351    /// never-credentialed are exactly the rows it cannot describe.
352    pub const fn auth_kind(self) -> AuthKind {
353        match self {
354            VendorId::Anthropic
355            | VendorId::Openai
356            | VendorId::Copilot
357            | VendorId::NousResearch
358            | VendorId::CommandCode => AuthKind::Oauth,
359            VendorId::AnthropicApi
360            | VendorId::Zai
361            | VendorId::Openrouter
362            | VendorId::Deepseek
363            | VendorId::Kimi
364            | VendorId::Kilo
365            | VendorId::Novita
366            | VendorId::Moonshot
367            | VendorId::Grok
368            | VendorId::Minimax
369            | VendorId::OpenCodeGo => AuthKind::ApiKey,
370            // No credential of their own: another local product's session is
371            // the login. Antigravity has no credential file at all (the binary
372            // probes whichever local server answers), Cursor and Kiro read the
373            // IDE's and kiro-cli's own state, and SuperGrok uses the Grok Build
374            // CLI's login.
375            VendorId::Supergrok | VendorId::Antigravity | VendorId::Cursor | VendorId::Kiro => {
376                AuthKind::Local
377            }
378        }
379    }
380
381    /// Default environment variable holding this provider's key, or `""` for a
382    /// provider that has none. This is only the *default*: most key vendors
383    /// accept an `api_key_env` override in config, so a frontend showing the
384    /// variable a user must set wants [`Config::api_key_env_for`], not this.
385    pub const fn api_key_env(self) -> &'static str {
386        match self {
387            VendorId::AnthropicApi => "ANTHROPIC_ADMIN_KEY",
388            VendorId::Zai => "ZAI_API_KEY",
389            VendorId::Openrouter => "OPENROUTER_API_KEY",
390            VendorId::Deepseek => "DEEPSEEK_API_KEY",
391            VendorId::Kimi => "KIMI_API_KEY",
392            VendorId::Kilo => "KILO_API_KEY",
393            VendorId::Novita => "NOVITA_API_KEY",
394            VendorId::Moonshot => "MOONSHOT_API_KEY",
395            VendorId::Grok => "XAI_MANAGEMENT_KEY",
396            VendorId::Minimax => "MINIMAX_API_KEY",
397            VendorId::OpenCodeGo => "OPENCODE_GO_API_KEY",
398            // OAuth-first, with an environment override for CI and headless
399            // use. Neither name is configurable, so neither has an
400            // `api_key_env` field in its config section.
401            VendorId::Copilot => "GITHUB_COPILOT_TOKEN",
402            VendorId::CommandCode => "COMMANDCODE_API_KEY",
403            VendorId::Anthropic
404            | VendorId::Openai
405            | VendorId::Supergrok
406            | VendorId::Antigravity
407            | VendorId::Cursor
408            | VendorId::Kiro
409            | VendorId::NousResearch => "",
410        }
411    }
412
413    /// Command that signs this provider in, or `""` when signing in happens
414    /// somewhere this cannot name — a desktop app's own window. The strings
415    /// are the ones the vendor modules' own credential errors already print,
416    /// so a status row and a failed fetch tell the user to run the same thing.
417    /// One sentence telling the user how to sign this provider in, for a UI
418    /// that has an error card and no room for a manual.
419    ///
420    /// This is product knowledge, so it lives beside [`Self::login_command`]
421    /// rather than in a frontend table. The Windows popover grew its own copy
422    /// first and it disagreed with this one for five of eight providers before
423    /// it had shipped — the match here is exhaustive, so a new provider cannot
424    /// be added without saying how a person signs into it.
425    pub const fn sign_in_hint(self) -> &'static str {
426        match self {
427            VendorId::Anthropic => "Run `claude` in a terminal, then Refresh.",
428            VendorId::Openai => "Run `codex login` in a terminal, then Refresh.",
429            VendorId::Copilot => "Run `gh auth login` in a terminal, then Refresh.",
430            VendorId::Kiro => "Run `kiro-cli login` in a terminal, then Refresh.",
431            VendorId::Kimi => "Run `kimi` in a terminal, or set an API key.",
432            VendorId::CommandCode => "Run `commandcode` in a terminal, then Refresh.",
433            VendorId::NousResearch => {
434                "Run `ai-usagebar auth nous login` in a terminal, then Refresh."
435            }
436            VendorId::Cursor => "Sign in to the Cursor app, then Refresh.",
437            VendorId::Antigravity => "Open Antigravity or run `agy`, then Refresh.",
438            VendorId::Grok | VendorId::Supergrok => "Sign in with `grok`, then Refresh.",
439            // Key-only providers: there is nothing to log into, only a key to
440            // put in the config.
441            VendorId::AnthropicApi
442            | VendorId::Zai
443            | VendorId::Openrouter
444            | VendorId::Deepseek
445            | VendorId::Kilo
446            | VendorId::Novita
447            | VendorId::Moonshot
448            | VendorId::Minimax
449            | VendorId::OpenCodeGo => "Add an API key in Settings, then Refresh.",
450        }
451    }
452
453    pub const fn login_command(self) -> &'static str {
454        match self {
455            VendorId::Anthropic => "claude",
456            VendorId::Openai => "codex login",
457            VendorId::Copilot => "gh auth login",
458            VendorId::CommandCode => "commandcode",
459            VendorId::NousResearch => "ai-usagebar auth nous login",
460            VendorId::Kiro => "kiro-cli login",
461            // Kimi takes a key *or* the Kimi Code CLI's own OAuth login, which
462            // is what a subscriber already has locally.
463            VendorId::Kimi => "kimi",
464            VendorId::AnthropicApi
465            | VendorId::Zai
466            | VendorId::Openrouter
467            | VendorId::Deepseek
468            | VendorId::Kilo
469            | VendorId::Novita
470            | VendorId::Moonshot
471            | VendorId::Grok
472            | VendorId::Supergrok
473            | VendorId::Antigravity
474            | VendorId::Cursor
475            | VendorId::Minimax
476            | VendorId::OpenCodeGo => "",
477        }
478    }
479
480    pub fn all() -> &'static [VendorId] {
481        &[
482            VendorId::Anthropic,
483            VendorId::AnthropicApi,
484            VendorId::Openai,
485            VendorId::Copilot,
486            VendorId::Zai,
487            VendorId::Openrouter,
488            VendorId::Deepseek,
489            VendorId::Kimi,
490            VendorId::Kilo,
491            VendorId::Novita,
492            VendorId::Moonshot,
493            VendorId::Grok,
494            VendorId::Supergrok,
495            VendorId::Antigravity,
496            VendorId::Cursor,
497            VendorId::Minimax,
498            VendorId::Kiro,
499            VendorId::NousResearch,
500            VendorId::OpenCodeGo,
501            VendorId::CommandCode,
502        ]
503    }
504}
505
506/// What a vendor returns from a successful fetch — the same
507/// [`Outcome`](crate::outcome::Outcome) every vendor produces, once its own
508/// snapshot type has been widened to [`VendorSnapshot`]. Each vendor gets
509/// there with a single `outcome.map(VendorSnapshot::Whichever)`.
510pub type VendorOutcome = crate::outcome::Outcome<VendorSnapshot>;
511
512/// Options forwarded to renderers from the CLI.
513#[derive(Debug, Clone)]
514pub struct RenderOpts {
515    pub format: Option<String>,
516    pub tooltip_format: Option<String>,
517    pub icon: Option<String>,
518    pub pace_tolerance: u32,
519    pub format_pace_color: bool,
520    pub tooltip_pace_pts: bool,
521}
522
523impl RenderOpts {
524    pub fn from_cli(cli: &Cli) -> Self {
525        Self {
526            format: cli.format.clone(),
527            tooltip_format: cli.tooltip_format.clone(),
528            icon: cli.icon.clone(),
529            pace_tolerance: cli.pace_tolerance,
530            format_pace_color: cli.format_pace_color,
531            tooltip_pace_pts: cli.tooltip_pace_pts,
532        }
533    }
534}
535
536#[cfg(test)]
537mod tests {
538    use super::*;
539
540    #[test]
541    fn every_vendor_has_stable_machine_and_display_names() {
542        for vendor in VendorId::all() {
543            assert!(!vendor.slug().is_empty());
544            assert!(!vendor.display_name().is_empty());
545        }
546        assert_eq!(VendorId::Anthropic.slug(), "anthropic");
547        assert_eq!(VendorId::Anthropic.display_name(), "Claude");
548        assert_eq!(VendorId::Openai.display_name(), "Codex");
549        assert_eq!(VendorId::Zai.display_name(), "Z.AI");
550    }
551
552    /// `{vendor_short}` is a documented format placeholder and now also rides
553    /// the `usage --json` report, so a duplicate or a re-typed code would make
554    /// two providers indistinguishable in a bar that shows nothing else.
555    #[test]
556    fn every_vendor_short_name_is_a_unique_three_letter_code() {
557        let mut seen = std::collections::BTreeSet::new();
558        for vendor in VendorId::all() {
559            let short = vendor.short_name();
560            assert_eq!(short.len(), 3, "{} is not three letters", vendor.slug());
561            assert!(
562                short.chars().all(|c| c.is_ascii_lowercase()),
563                "{} is not lowercase ascii",
564                vendor.slug()
565            );
566            assert!(seen.insert(short), "{short} is used by two vendors");
567        }
568        assert_eq!(VendorId::Anthropic.short_name(), "cld");
569        assert_eq!(VendorId::Openai.short_name(), "gpt");
570        assert_eq!(VendorId::Zai.short_name(), "zai");
571        assert_eq!(VendorId::Antigravity.short_name(), "agy");
572    }
573
574    /// The bar can show every provider at once, so a glyph two providers share
575    /// tells the user nothing about which row is which. Grok and SuperGrok are
576    /// the one sanctioned pair — same brand, two products. Providers without a
577    /// distinct Nerd Font mark use their `short_name`, which is unique by
578    /// construction and cannot render as tofu.
579    #[test]
580    fn every_vendor_has_a_bar_icon_and_no_two_share_one() {
581        use std::collections::BTreeMap;
582
583        let mut by_icon: BTreeMap<&str, Vec<&str>> = BTreeMap::new();
584        for vendor in VendorId::all() {
585            assert!(!vendor.bar_icon().is_empty(), "{}", vendor.slug());
586            by_icon
587                .entry(vendor.bar_icon())
588                .or_default()
589                .push(vendor.slug());
590        }
591
592        let shared: Vec<_> = by_icon
593            .iter()
594            .filter(|(_, vendors)| vendors.len() > 1)
595            .filter(|(_, vendors)| vendors.as_slice() != ["grok", "supergrok"])
596            .collect();
597        assert!(
598            shared.is_empty(),
599            "these providers are indistinguishable in a bar that shows them \
600             side by side: {shared:#?}"
601        );
602        assert_eq!(VendorId::Anthropic.bar_icon(), "󰚩");
603        assert_eq!(VendorId::Openai.bar_icon(), "󱢆");
604        assert_eq!(VendorId::Supergrok.bar_icon(), VendorId::Grok.bar_icon());
605        assert_eq!(VendorId::CommandCode.bar_icon(), "cmc");
606    }
607
608    #[test]
609    fn new_vendor_contracts_keep_public_names_and_slugs() {
610        assert_eq!(VendorId::NousResearch.slug(), "nous");
611        assert_eq!(VendorId::NousResearch.display_name(), "Nous Research");
612        assert_eq!(VendorId::OpenCodeGo.slug(), "opencode-go");
613        assert_eq!(VendorId::OpenCodeGo.display_name(), "OpenCode Go");
614        assert_eq!(
615            serde_json::to_value(VendorId::OpenCodeGo).unwrap(),
616            serde_json::json!("opencode-go")
617        );
618    }
619
620    #[test]
621    fn vendor_secret_env_vars_cover_config_defaults() {
622        let configured_defaults = [
623            "ZAI_API_KEY",
624            "OPENROUTER_API_KEY",
625            "DEEPSEEK_API_KEY",
626            "KIMI_API_KEY",
627            "KILO_API_KEY",
628            "NOVITA_API_KEY",
629            "MINIMAX_API_KEY",
630            "MOONSHOT_API_KEY",
631            "XAI_MANAGEMENT_KEY",
632            "ANTHROPIC_ADMIN_KEY",
633            "GITHUB_COPILOT_TOKEN",
634        ];
635        for name in configured_defaults {
636            assert!(VENDOR_SECRET_ENV_VARS.contains(&name), "missing {name}");
637        }
638    }
639
640    #[test]
641    fn vars_to_remove_preserves_only_requested_grok_credentials() {
642        let removed = vendor_secret_env_vars_to_remove(&["XAI_API_KEY", "GROK_API_KEY"]);
643        assert!(!removed.contains(&"XAI_API_KEY"));
644        assert!(!removed.contains(&"GROK_API_KEY"));
645        assert!(removed.contains(&"ANTHROPIC_ADMIN_KEY"));
646        assert!(removed.contains(&"OPENROUTER_API_KEY"));
647        // Counted against the static list: another test in this process may
648        // have registered a custom provider's env var, which belongs here too.
649        let builtins = removed
650            .iter()
651            .filter(|var| VENDOR_SECRET_ENV_VARS.contains(var))
652            .count();
653        assert_eq!(builtins, VENDOR_SECRET_ENV_VARS.len() - 2);
654    }
655
656    #[test]
657    fn a_registered_custom_env_var_is_scrubbed_like_a_builtin_one() {
658        let name = "AI_USAGEBAR_TEST_CUSTOM_TOKEN_7F3A";
659        assert!(!vendor_secret_env_vars_to_remove(&[]).contains(&name));
660
661        register_secret_env_vars(&[name.to_string(), "not a name!".to_string()]);
662        register_secret_env_vars(&[name.to_string()]);
663
664        let removed = vendor_secret_env_vars_to_remove(&[]);
665        assert_eq!(
666            removed.iter().filter(|var| **var == name).count(),
667            1,
668            "registering twice must not list it twice: {removed:?}"
669        );
670        assert!(!removed.contains(&"not a name!"), "{removed:?}");
671        assert!(
672            !vendor_secret_env_vars_to_remove(&[name]).contains(&name),
673            "`keep` applies to registered names too"
674        );
675    }
676
677    #[test]
678    fn copilot_token_is_removed_before_unrelated_subprocesses_launch() {
679        let removed = vendor_secret_env_vars_to_remove(&[]);
680        assert!(removed.contains(&"GITHUB_COPILOT_TOKEN"));
681    }
682
683    #[tokio::test]
684    async fn body_over_the_cap_is_refused_and_under_it_round_trips() {
685        let mut server = mockito::Server::new_async().await;
686        server
687            .mock("GET", "/big")
688            .with_status(200)
689            .with_body("x".repeat(4096))
690            .create_async()
691            .await;
692        server
693            .mock("GET", "/small")
694            .with_status(200)
695            .with_body("hello")
696            .create_async()
697            .await;
698
699        let client = reqwest::Client::new();
700
701        // Over the cap: refused rather than buffered.
702        let resp = client
703            .get(format!("{}/big", server.url()))
704            .send()
705            .await
706            .unwrap();
707        let err = read_body_capped(resp, 1024).await.unwrap_err();
708        assert!(
709            err.to_string().contains("exceeds"),
710            "unexpected error: {err}"
711        );
712
713        // Under the cap: identical to the previous `resp.bytes()` behaviour.
714        let resp = client
715            .get(format!("{}/small", server.url()))
716            .send()
717            .await
718            .unwrap();
719        assert_eq!(read_body_capped(resp, 1024).await.unwrap(), b"hello");
720    }
721
722    #[tokio::test]
723    async fn chunked_body_without_content_length_still_hits_the_cap() {
724        let mut server = mockito::Server::new_async().await;
725        server
726            .mock("GET", "/chunked")
727            .with_status(200)
728            .with_chunked_body(|writer| writer.write_all(&[b'x'; 4096]))
729            .create_async()
730            .await;
731
732        let response = reqwest::Client::new()
733            .get(format!("{}/chunked", server.url()))
734            .send()
735            .await
736            .unwrap();
737        assert!(response.content_length().is_none());
738        let error = read_body_capped(response, 1024).await.unwrap_err();
739        assert!(error.to_string().contains("exceeds"), "{error}");
740    }
741
742    #[tokio::test]
743    async fn same_origin_redirects_still_work_with_vendor_headers() {
744        let mut server = mockito::Server::new_async().await;
745        let redirect = server
746            .mock("GET", "/start")
747            .match_header("x-api-key", "secret")
748            .with_status(302)
749            .with_header("location", "/finish")
750            .create_async()
751            .await;
752        let finish = server
753            .mock("GET", "/finish")
754            .match_header("x-api-key", "secret")
755            .with_status(200)
756            .create_async()
757            .await;
758        let client = reqwest::Client::builder()
759            .redirect(same_origin_redirect_policy())
760            .build()
761            .unwrap();
762
763        let response = client
764            .get(format!("{}/start", server.url()))
765            .header("x-api-key", "secret")
766            .send()
767            .await
768            .unwrap();
769
770        assert_eq!(response.status(), reqwest::StatusCode::OK);
771        redirect.assert_async().await;
772        finish.assert_async().await;
773    }
774
775    #[tokio::test]
776    async fn cross_origin_redirects_are_not_followed_with_vendor_headers() {
777        let mut origin = mockito::Server::new_async().await;
778        let mut target = mockito::Server::new_async().await;
779        let target_url = format!("{}/capture", target.url());
780        let redirect = origin
781            .mock("GET", "/start")
782            .match_header("x-api-key", "secret")
783            .with_status(302)
784            .with_header("location", &target_url)
785            .create_async()
786            .await;
787        let capture = target
788            .mock("GET", "/capture")
789            .expect(0)
790            .create_async()
791            .await;
792        let client = reqwest::Client::builder()
793            .redirect(same_origin_redirect_policy())
794            .build()
795            .unwrap();
796
797        let response = client
798            .get(format!("{}/start", origin.url()))
799            .header("x-api-key", "secret")
800            .send()
801            .await
802            .unwrap();
803
804        assert_eq!(response.status(), reqwest::StatusCode::FOUND);
805        redirect.assert_async().await;
806        capture.assert_async().await;
807    }
808
809    #[test]
810    fn vendor_id_slug_round_trip() {
811        for id in VendorId::all() {
812            assert_eq!(
813                id.slug(),
814                serde_json::to_value(id).unwrap().as_str().unwrap()
815            );
816        }
817    }
818}