Skip to main content

ai_usagebar/
vendor.rs

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