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

ai_usagebar/
vendor.rs
