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

ai_usagebar/
vendor.rs
