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 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
565/// What a vendor returns from a successful fetch — the same
566/// [`Outcome`](crate::outcome::Outcome) every vendor produces, once its own
567/// snapshot type has been widened to [`VendorSnapshot`]. Each vendor gets
568/// there with a single `outcome.map(VendorSnapshot::Whichever)`.
569pub type VendorOutcome = crate::outcome::Outcome<VendorSnapshot>;
570
571/// Options forwarded to renderers from the CLI.
572#[derive(Debug, Clone)]
573pub struct RenderOpts {
574    pub format: Option<String>,
575    pub tooltip_format: Option<String>,
576    pub icon: Option<String>,
577    pub pace_tolerance: u32,
578    pub format_pace_color: bool,
579    pub tooltip_pace_pts: bool,
580}
581
582impl RenderOpts {
583    pub fn from_cli(cli: &Cli) -> Self {
584        Self {
585            format: cli.format.clone(),
586            tooltip_format: cli.tooltip_format.clone(),
587            icon: cli.icon.clone(),
588            pace_tolerance: cli.pace_tolerance,
589            format_pace_color: cli.format_pace_color,
590            tooltip_pace_pts: cli.tooltip_pace_pts,
591        }
592    }
593}
594
595#[cfg(test)]
596mod tests {
597    use super::*;
598
599    #[test]
600    fn every_vendor_has_stable_machine_and_display_names() {
601        for vendor in VendorId::all() {
602            assert!(!vendor.slug().is_empty());
603            assert!(!vendor.display_name().is_empty());
604        }
605        assert_eq!(VendorId::Anthropic.slug(), "anthropic");
606        assert_eq!(VendorId::Anthropic.display_name(), "Claude");
607        assert_eq!(VendorId::Openai.display_name(), "Codex");
608        assert_eq!(VendorId::Zai.display_name(), "Z.AI");
609    }
610
611    /// `{vendor_short}` is a documented format placeholder and now also rides
612    /// the `usage --json` report, so a duplicate or a re-typed code would make
613    /// two providers indistinguishable in a bar that shows nothing else.
614    #[test]
615    fn every_vendor_short_name_is_a_unique_three_letter_code() {
616        let mut seen = std::collections::BTreeSet::new();
617        for vendor in VendorId::all() {
618            let short = vendor.short_name();
619            assert_eq!(short.len(), 3, "{} is not three letters", vendor.slug());
620            assert!(
621                short.chars().all(|c| c.is_ascii_lowercase()),
622                "{} is not lowercase ascii",
623                vendor.slug()
624            );
625            assert!(seen.insert(short), "{short} is used by two vendors");
626        }
627        assert_eq!(VendorId::Anthropic.short_name(), "cld");
628        assert_eq!(VendorId::Openai.short_name(), "gpt");
629        assert_eq!(VendorId::Zai.short_name(), "zai");
630        assert_eq!(VendorId::Antigravity.short_name(), "agy");
631    }
632
633    /// The bar can show every provider at once, so a glyph two providers share
634    /// tells the user nothing about which row is which. Grok and SuperGrok are
635    /// the one sanctioned pair — same brand, two products. Providers without a
636    /// distinct Nerd Font mark use their `short_name`, which is unique by
637    /// construction and cannot render as tofu.
638    #[test]
639    fn every_vendor_has_a_bar_icon_and_no_two_share_one() {
640        use std::collections::BTreeMap;
641
642        let mut by_icon: BTreeMap<&str, Vec<&str>> = BTreeMap::new();
643        for vendor in VendorId::all() {
644            assert!(!vendor.bar_icon().is_empty(), "{}", vendor.slug());
645            by_icon
646                .entry(vendor.bar_icon())
647                .or_default()
648                .push(vendor.slug());
649        }
650
651        let shared: Vec<_> = by_icon
652            .iter()
653            .filter(|(_, vendors)| vendors.len() > 1)
654            .filter(|(_, vendors)| vendors.as_slice() != ["grok", "supergrok"])
655            .collect();
656        assert!(
657            shared.is_empty(),
658            "these providers are indistinguishable in a bar that shows them \
659             side by side: {shared:#?}"
660        );
661        assert_eq!(VendorId::Anthropic.bar_icon(), "󰚩");
662        assert_eq!(VendorId::Openai.bar_icon(), "󱢆");
663        assert_eq!(VendorId::Supergrok.bar_icon(), VendorId::Grok.bar_icon());
664        assert_eq!(VendorId::CommandCode.bar_icon(), "cmc");
665    }
666
667    #[test]
668    fn new_vendor_contracts_keep_public_names_and_slugs() {
669        assert_eq!(VendorId::NousResearch.slug(), "nous");
670        assert_eq!(VendorId::NousResearch.display_name(), "Nous Research");
671        assert_eq!(VendorId::OpenCodeGo.slug(), "opencode-go");
672        assert_eq!(VendorId::OpenCodeGo.display_name(), "OpenCode Go");
673        assert_eq!(
674            serde_json::to_value(VendorId::OpenCodeGo).unwrap(),
675            serde_json::json!("opencode-go")
676        );
677    }
678
679    #[test]
680    fn vendor_secret_env_vars_cover_config_defaults() {
681        let configured_defaults = [
682            "ZAI_API_KEY",
683            "OPENROUTER_API_KEY",
684            "DEEPSEEK_API_KEY",
685            "KIMI_API_KEY",
686            "KILO_API_KEY",
687            "NOVITA_API_KEY",
688            "MINIMAX_API_KEY",
689            "MOONSHOT_API_KEY",
690            "XAI_MANAGEMENT_KEY",
691            "ANTHROPIC_ADMIN_KEY",
692            "GITHUB_COPILOT_TOKEN",
693            "ORCAROUTER_API_KEY",
694        ];
695        for name in configured_defaults {
696            assert!(VENDOR_SECRET_ENV_VARS.contains(&name), "missing {name}");
697        }
698    }
699
700    #[test]
701    fn vars_to_remove_preserves_only_requested_grok_credentials() {
702        let removed = vendor_secret_env_vars_to_remove(&["XAI_API_KEY", "GROK_API_KEY"]);
703        assert!(!removed.contains(&"XAI_API_KEY"));
704        assert!(!removed.contains(&"GROK_API_KEY"));
705        assert!(removed.contains(&"ANTHROPIC_ADMIN_KEY"));
706        assert!(removed.contains(&"OPENROUTER_API_KEY"));
707        // Counted against the static list: another test in this process may
708        // have registered a custom provider's env var, which belongs here too.
709        let builtins = removed
710            .iter()
711            .filter(|var| VENDOR_SECRET_ENV_VARS.contains(var))
712            .count();
713        assert_eq!(builtins, VENDOR_SECRET_ENV_VARS.len() - 2);
714    }
715
716    #[test]
717    fn a_registered_custom_env_var_is_scrubbed_like_a_builtin_one() {
718        let name = "AI_USAGEBAR_TEST_CUSTOM_TOKEN_7F3A";
719        assert!(!vendor_secret_env_vars_to_remove(&[]).contains(&name));
720
721        register_secret_env_vars(&[name.to_string(), "not a name!".to_string()]);
722        register_secret_env_vars(&[name.to_string()]);
723
724        let removed = vendor_secret_env_vars_to_remove(&[]);
725        assert_eq!(
726            removed.iter().filter(|var| **var == name).count(),
727            1,
728            "registering twice must not list it twice: {removed:?}"
729        );
730        assert!(!removed.contains(&"not a name!"), "{removed:?}");
731        assert!(
732            !vendor_secret_env_vars_to_remove(&[name]).contains(&name),
733            "`keep` applies to registered names too"
734        );
735    }
736
737    #[test]
738    fn copilot_token_is_removed_before_unrelated_subprocesses_launch() {
739        let removed = vendor_secret_env_vars_to_remove(&[]);
740        assert!(removed.contains(&"GITHUB_COPILOT_TOKEN"));
741    }
742
743    #[tokio::test]
744    async fn body_over_the_cap_is_refused_and_under_it_round_trips() {
745        let mut server = mockito::Server::new_async().await;
746        server
747            .mock("GET", "/big")
748            .with_status(200)
749            .with_body("x".repeat(4096))
750            .create_async()
751            .await;
752        server
753            .mock("GET", "/small")
754            .with_status(200)
755            .with_body("hello")
756            .create_async()
757            .await;
758
759        let client = reqwest::Client::new();
760
761        // Over the cap: refused rather than buffered.
762        let resp = client
763            .get(format!("{}/big", server.url()))
764            .send()
765            .await
766            .unwrap();
767        let err = read_body_capped(resp, 1024).await.unwrap_err();
768        assert!(
769            err.to_string().contains("exceeds"),
770            "unexpected error: {err}"
771        );
772
773        // Under the cap: identical to the previous `resp.bytes()` behaviour.
774        let resp = client
775            .get(format!("{}/small", server.url()))
776            .send()
777            .await
778            .unwrap();
779        assert_eq!(read_body_capped(resp, 1024).await.unwrap(), b"hello");
780    }
781
782    #[tokio::test]
783    async fn chunked_body_without_content_length_still_hits_the_cap() {
784        let mut server = mockito::Server::new_async().await;
785        server
786            .mock("GET", "/chunked")
787            .with_status(200)
788            .with_chunked_body(|writer| writer.write_all(&[b'x'; 4096]))
789            .create_async()
790            .await;
791
792        let response = reqwest::Client::new()
793            .get(format!("{}/chunked", server.url()))
794            .send()
795            .await
796            .unwrap();
797        assert!(response.content_length().is_none());
798        let error = read_body_capped(response, 1024).await.unwrap_err();
799        assert!(error.to_string().contains("exceeds"), "{error}");
800    }
801
802    #[tokio::test]
803    async fn same_origin_redirects_still_work_with_vendor_headers() {
804        let mut server = mockito::Server::new_async().await;
805        let redirect = server
806            .mock("GET", "/start")
807            .match_header("x-api-key", "secret")
808            .with_status(302)
809            .with_header("location", "/finish")
810            .create_async()
811            .await;
812        let finish = server
813            .mock("GET", "/finish")
814            .match_header("x-api-key", "secret")
815            .with_status(200)
816            .create_async()
817            .await;
818        let client = reqwest::Client::builder()
819            .redirect(same_origin_redirect_policy())
820            .build()
821            .unwrap();
822
823        let response = client
824            .get(format!("{}/start", server.url()))
825            .header("x-api-key", "secret")
826            .send()
827            .await
828            .unwrap();
829
830        assert_eq!(response.status(), reqwest::StatusCode::OK);
831        redirect.assert_async().await;
832        finish.assert_async().await;
833    }
834
835    #[tokio::test]
836    async fn cross_origin_redirects_are_not_followed_with_vendor_headers() {
837        let mut origin = mockito::Server::new_async().await;
838        let mut target = mockito::Server::new_async().await;
839        let target_url = format!("{}/capture", target.url());
840        let redirect = origin
841            .mock("GET", "/start")
842            .match_header("x-api-key", "secret")
843            .with_status(302)
844            .with_header("location", &target_url)
845            .create_async()
846            .await;
847        let capture = target
848            .mock("GET", "/capture")
849            .expect(0)
850            .create_async()
851            .await;
852        let client = reqwest::Client::builder()
853            .redirect(same_origin_redirect_policy())
854            .build()
855            .unwrap();
856
857        let response = client
858            .get(format!("{}/start", origin.url()))
859            .header("x-api-key", "secret")
860            .send()
861            .await
862            .unwrap();
863
864        assert_eq!(response.status(), reqwest::StatusCode::FOUND);
865        redirect.assert_async().await;
866        capture.assert_async().await;
867    }
868
869    #[test]
870    fn vendor_id_slug_round_trip() {
871        for id in VendorId::all() {
872            assert_eq!(
873                id.slug(),
874                serde_json::to_value(id).unwrap().as_str().unwrap()
875            );
876        }
877    }
878}
ai_usagebar/vendor.rs

ai_usagebar/
vendor.rs
