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

ai_usagebar/
vendor.rs
