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

ai_usagebar/
vendor.rs
