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

ai_usagebar/
vendor.rs
