use std::collections::BTreeSet;
use std::sync::{Mutex, OnceLock};
use std::time::Duration;
use clap::ValueEnum;
use crate::usage::VendorSnapshot;
use crate::widget::cli::Cli;
pub const HTTP_CLIENT_TIMEOUT: Duration = Duration::from_secs(30);
pub const MAX_BODY_BYTES: usize = 2 * 1024 * 1024;
pub(crate) const VENDOR_SECRET_ENV_VARS: &[&str] = &[
"ZAI_API_KEY",
"OPENROUTER_API_KEY",
"DEEPSEEK_API_KEY",
"KIMI_API_KEY",
"KILO_API_KEY",
"NOVITA_API_KEY",
"MINIMAX_API_KEY",
"MOONSHOT_API_KEY",
"XAI_MANAGEMENT_KEY",
"ANTHROPIC_ADMIN_KEY",
"XAI_API_KEY",
"GROK_API_KEY",
"OPENCODE_GO_API_KEY",
"COMMANDCODE_API_KEY",
"GITHUB_COPILOT_TOKEN",
"GH_TOKEN",
"GITHUB_TOKEN",
"OLLAMA_API_KEY",
];
fn registered_secret_env_vars() -> &'static Mutex<BTreeSet<&'static str>> {
static REGISTERED: OnceLock<Mutex<BTreeSet<&'static str>>> = OnceLock::new();
REGISTERED.get_or_init(|| Mutex::new(BTreeSet::new()))
}
pub fn register_secret_env_vars(names: &[String]) {
let mut registered = registered_secret_env_vars()
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
for name in names {
if !crate::config::is_valid_env_var_name(name)
|| VENDOR_SECRET_ENV_VARS.contains(&name.as_str())
|| registered.contains(name.as_str())
{
continue;
}
registered.insert(Box::leak(name.clone().into_boxed_str()));
}
}
pub(crate) fn vendor_secret_env_vars_to_remove(keep: &[&str]) -> Vec<&'static str> {
let registered = registered_secret_env_vars()
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
VENDOR_SECRET_ENV_VARS
.iter()
.copied()
.chain(registered.iter().copied())
.filter(|var| !keep.contains(var))
.collect()
}
pub fn same_origin_redirect_policy() -> reqwest::redirect::Policy {
reqwest::redirect::Policy::custom(|attempt| {
if attempt.previous().len() >= 10 {
return attempt.error("too many redirects");
}
let Some(origin) = attempt.previous().first() else {
return attempt.stop();
};
let target = attempt.url();
if target.scheme() == origin.scheme()
&& target.host_str() == origin.host_str()
&& target.port_or_known_default() == origin.port_or_known_default()
{
attempt.follow()
} else {
attempt.stop()
}
})
}
pub async fn read_body_capped(
mut resp: reqwest::Response,
max: usize,
) -> crate::error::Result<Vec<u8>> {
let too_big = |n: u64| {
crate::error::AppError::Schema(format!(
"response body exceeds the {max}-byte limit ({n} bytes); refusing to buffer it"
))
};
if let Some(len) = resp.content_length()
&& len > max as u64
{
return Err(too_big(len));
}
let mut buf: Vec<u8> = Vec::new();
while let Some(chunk) = resp.chunk().await? {
if chunk.len() > max.saturating_sub(buf.len()) {
return Err(too_big(buf.len().saturating_add(chunk.len()) as u64));
}
buf.extend_from_slice(&chunk);
}
Ok(buf)
}
#[derive(
Debug, Clone, Copy, ValueEnum, PartialEq, Eq, Hash, serde::Deserialize, serde::Serialize,
)]
#[serde(rename_all = "lowercase")]
pub enum VendorId {
Anthropic,
#[serde(rename = "anthropic_api")]
AnthropicApi,
Openai,
Copilot,
Zai,
Openrouter,
Deepseek,
Kimi,
Kilo,
Novita,
Moonshot,
Grok,
Supergrok,
Antigravity,
Cursor,
Minimax,
Kiro,
#[serde(rename = "nous")]
NousResearch,
#[serde(rename = "opencode-go")]
OpenCodeGo,
#[serde(rename = "commandcode")]
CommandCode,
Ollama,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum AuthKind {
Oauth,
ApiKey,
Local,
}
impl AuthKind {
pub const fn as_str(self) -> &'static str {
match self {
AuthKind::Oauth => "oauth",
AuthKind::ApiKey => "apikey",
AuthKind::Local => "local",
}
}
}
impl VendorId {
pub fn slug(self) -> &'static str {
match self {
VendorId::Anthropic => "anthropic",
VendorId::AnthropicApi => "anthropic_api",
VendorId::Openai => "openai",
VendorId::Copilot => "copilot",
VendorId::Zai => "zai",
VendorId::Openrouter => "openrouter",
VendorId::Deepseek => "deepseek",
VendorId::Kimi => "kimi",
VendorId::Kilo => "kilo",
VendorId::Novita => "novita",
VendorId::Moonshot => "moonshot",
VendorId::Grok => "grok",
VendorId::Supergrok => "supergrok",
VendorId::Antigravity => "antigravity",
VendorId::Cursor => "cursor",
VendorId::Minimax => "minimax",
VendorId::Kiro => "kiro",
VendorId::NousResearch => "nous",
VendorId::OpenCodeGo => "opencode-go",
VendorId::CommandCode => "commandcode",
VendorId::Ollama => "ollama",
}
}
pub fn display_name(self) -> &'static str {
match self {
VendorId::Anthropic => "Claude",
VendorId::AnthropicApi => "Anthropic API",
VendorId::Openai => "Codex",
VendorId::Copilot => "GitHub Copilot",
VendorId::Zai => "Z.AI",
VendorId::Openrouter => "OpenRouter",
VendorId::Deepseek => "DeepSeek",
VendorId::Kimi => "Kimi",
VendorId::Kilo => "Kilo",
VendorId::Novita => "Novita",
VendorId::Moonshot => "Moonshot",
VendorId::Grok => "Grok",
VendorId::Supergrok => "SuperGrok",
VendorId::Antigravity => "Antigravity",
VendorId::Cursor => "Cursor",
VendorId::Minimax => "MiniMax",
VendorId::Kiro => "Kiro",
VendorId::NousResearch => "Nous Research",
VendorId::OpenCodeGo => "OpenCode Go",
VendorId::CommandCode => "Command Code",
VendorId::Ollama => "Ollama Cloud",
}
}
pub const fn bar_icon(self) -> &'static str {
match self {
VendorId::Anthropic => "",
VendorId::AnthropicApi => "",
VendorId::Openai => "",
VendorId::Copilot => "",
VendorId::Zai => VendorId::Zai.short_name(),
VendorId::Openrouter => "",
VendorId::Deepseek => "",
VendorId::Kimi => VendorId::Kimi.short_name(),
VendorId::Kilo => "",
VendorId::Novita => "",
VendorId::Moonshot => VendorId::Moonshot.short_name(),
VendorId::Grok | VendorId::Supergrok => "",
VendorId::Antigravity => VendorId::Antigravity.short_name(),
VendorId::Cursor => "❯",
VendorId::Minimax => VendorId::Minimax.short_name(),
VendorId::Kiro => "◆",
VendorId::NousResearch => VendorId::NousResearch.short_name(),
VendorId::OpenCodeGo => VendorId::OpenCodeGo.short_name(),
VendorId::CommandCode => VendorId::CommandCode.short_name(),
VendorId::Ollama => VendorId::Ollama.short_name(),
}
}
pub const fn short_name(self) -> &'static str {
match self {
VendorId::Anthropic => "cld",
VendorId::AnthropicApi => "aac",
VendorId::Openai => "gpt",
VendorId::Copilot => "ghc",
VendorId::Zai => "zai",
VendorId::Openrouter => "opr",
VendorId::Deepseek => "dsk",
VendorId::Kimi => "kmi",
VendorId::Kilo => "klo",
VendorId::Novita => "nvt",
VendorId::Moonshot => "msh",
VendorId::Grok => "grk",
VendorId::Supergrok => "sgk",
VendorId::Antigravity => "agy",
VendorId::Cursor => "cur",
VendorId::Minimax => "mmx",
VendorId::Kiro => "kir",
VendorId::NousResearch => "nrs",
VendorId::OpenCodeGo => "ocg",
VendorId::CommandCode => "cmc",
VendorId::Ollama => "oll",
}
}
pub const fn config_section(self) -> &'static str {
match self {
VendorId::Anthropic => "anthropic",
VendorId::AnthropicApi => "anthropic_api",
VendorId::Openai => "openai",
VendorId::Copilot => "copilot",
VendorId::Zai => "zai",
VendorId::Openrouter => "openrouter",
VendorId::Deepseek => "deepseek",
VendorId::Kimi => "kimi",
VendorId::Kilo => "kilo",
VendorId::Novita => "novita",
VendorId::Moonshot => "moonshot",
VendorId::Grok => "grok",
VendorId::Supergrok => "supergrok",
VendorId::Antigravity => "antigravity",
VendorId::Cursor => "cursor",
VendorId::Minimax => "minimax",
VendorId::Kiro => "kiro",
VendorId::NousResearch => "nous",
VendorId::OpenCodeGo => "opencode-go",
VendorId::CommandCode => "commandcode",
VendorId::Ollama => "ollama",
}
}
pub const fn auth_kind(self) -> AuthKind {
match self {
VendorId::Anthropic
| VendorId::Openai
| VendorId::Copilot
| VendorId::NousResearch
| VendorId::CommandCode => AuthKind::Oauth,
VendorId::AnthropicApi
| VendorId::Zai
| VendorId::Openrouter
| VendorId::Deepseek
| VendorId::Kimi
| VendorId::Kilo
| VendorId::Novita
| VendorId::Moonshot
| VendorId::Grok
| VendorId::Minimax
| VendorId::OpenCodeGo
| VendorId::Ollama => AuthKind::ApiKey,
VendorId::Supergrok | VendorId::Antigravity | VendorId::Cursor | VendorId::Kiro => {
AuthKind::Local
}
}
}
pub const fn api_key_env(self) -> &'static str {
match self {
VendorId::AnthropicApi => "ANTHROPIC_ADMIN_KEY",
VendorId::Zai => "ZAI_API_KEY",
VendorId::Openrouter => "OPENROUTER_API_KEY",
VendorId::Deepseek => "DEEPSEEK_API_KEY",
VendorId::Kimi => "KIMI_API_KEY",
VendorId::Kilo => "KILO_API_KEY",
VendorId::Novita => "NOVITA_API_KEY",
VendorId::Moonshot => "MOONSHOT_API_KEY",
VendorId::Grok => "XAI_MANAGEMENT_KEY",
VendorId::Minimax => "MINIMAX_API_KEY",
VendorId::OpenCodeGo => "OPENCODE_GO_API_KEY",
VendorId::Ollama => "OLLAMA_API_KEY",
VendorId::Copilot => "GITHUB_COPILOT_TOKEN",
VendorId::CommandCode => "COMMANDCODE_API_KEY",
VendorId::Anthropic
| VendorId::Openai
| VendorId::Supergrok
| VendorId::Antigravity
| VendorId::Cursor
| VendorId::Kiro
| VendorId::NousResearch => "",
}
}
pub const fn sign_in_hint(self) -> &'static str {
match self {
VendorId::Anthropic => "Run `claude` in a terminal, then Refresh.",
VendorId::Openai => "Run `codex login` in a terminal, then Refresh.",
VendorId::Copilot => "Run `gh auth login` in a terminal, then Refresh.",
VendorId::Kiro => "Run `kiro-cli login` in a terminal, then Refresh.",
VendorId::Kimi => "Run `kimi` in a terminal, or set an API key.",
VendorId::CommandCode => "Run `commandcode` in a terminal, then Refresh.",
VendorId::NousResearch => {
"Run `ai-usagebar auth nous login` in a terminal, then Refresh."
}
VendorId::Cursor => "Sign in to the Cursor app, then Refresh.",
VendorId::Antigravity => "Open Antigravity or run `agy`, then Refresh.",
VendorId::Grok | VendorId::Supergrok => "Sign in with `grok`, then Refresh.",
VendorId::AnthropicApi
| VendorId::Zai
| VendorId::Openrouter
| VendorId::Deepseek
| VendorId::Kilo
| VendorId::Novita
| VendorId::Moonshot
| VendorId::Minimax
| VendorId::OpenCodeGo
| VendorId::Ollama => "Add an API key in Settings, then Refresh.",
}
}
pub const fn login_command(self) -> &'static str {
match self {
VendorId::Anthropic => "claude",
VendorId::Openai => "codex login",
VendorId::Copilot => "gh auth login",
VendorId::CommandCode => "commandcode",
VendorId::NousResearch => "ai-usagebar auth nous login",
VendorId::Kiro => "kiro-cli login",
VendorId::Kimi => "kimi",
VendorId::AnthropicApi
| VendorId::Zai
| VendorId::Openrouter
| VendorId::Deepseek
| VendorId::Kilo
| VendorId::Novita
| VendorId::Moonshot
| VendorId::Grok
| VendorId::Supergrok
| VendorId::Antigravity
| VendorId::Cursor
| VendorId::Minimax
| VendorId::OpenCodeGo
| VendorId::Ollama => "",
}
}
pub fn all() -> &'static [VendorId] {
&[
VendorId::Anthropic,
VendorId::AnthropicApi,
VendorId::Openai,
VendorId::Copilot,
VendorId::Zai,
VendorId::Openrouter,
VendorId::Deepseek,
VendorId::Kimi,
VendorId::Kilo,
VendorId::Novita,
VendorId::Moonshot,
VendorId::Grok,
VendorId::Supergrok,
VendorId::Antigravity,
VendorId::Cursor,
VendorId::Minimax,
VendorId::Kiro,
VendorId::NousResearch,
VendorId::OpenCodeGo,
VendorId::CommandCode,
VendorId::Ollama,
]
}
}
pub type VendorOutcome = crate::outcome::Outcome<VendorSnapshot>;
#[derive(Debug, Clone)]
pub struct RenderOpts {
pub format: Option<String>,
pub tooltip_format: Option<String>,
pub icon: Option<String>,
pub pace_tolerance: u32,
pub format_pace_color: bool,
pub tooltip_pace_pts: bool,
}
impl RenderOpts {
pub fn from_cli(cli: &Cli) -> Self {
Self {
format: cli.format.clone(),
tooltip_format: cli.tooltip_format.clone(),
icon: cli.icon.clone(),
pace_tolerance: cli.pace_tolerance,
format_pace_color: cli.format_pace_color,
tooltip_pace_pts: cli.tooltip_pace_pts,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn every_vendor_has_stable_machine_and_display_names() {
for vendor in VendorId::all() {
assert!(!vendor.slug().is_empty());
assert!(!vendor.display_name().is_empty());
}
assert_eq!(VendorId::Anthropic.slug(), "anthropic");
assert_eq!(VendorId::Anthropic.display_name(), "Claude");
assert_eq!(VendorId::Openai.display_name(), "Codex");
assert_eq!(VendorId::Zai.display_name(), "Z.AI");
}
#[test]
fn every_vendor_short_name_is_a_unique_three_letter_code() {
let mut seen = std::collections::BTreeSet::new();
for vendor in VendorId::all() {
let short = vendor.short_name();
assert_eq!(short.len(), 3, "{} is not three letters", vendor.slug());
assert!(
short.chars().all(|c| c.is_ascii_lowercase()),
"{} is not lowercase ascii",
vendor.slug()
);
assert!(seen.insert(short), "{short} is used by two vendors");
}
assert_eq!(VendorId::Anthropic.short_name(), "cld");
assert_eq!(VendorId::Openai.short_name(), "gpt");
assert_eq!(VendorId::Zai.short_name(), "zai");
assert_eq!(VendorId::Antigravity.short_name(), "agy");
}
#[test]
fn every_vendor_has_a_bar_icon_and_no_two_share_one() {
use std::collections::BTreeMap;
let mut by_icon: BTreeMap<&str, Vec<&str>> = BTreeMap::new();
for vendor in VendorId::all() {
assert!(!vendor.bar_icon().is_empty(), "{}", vendor.slug());
by_icon
.entry(vendor.bar_icon())
.or_default()
.push(vendor.slug());
}
let shared: Vec<_> = by_icon
.iter()
.filter(|(_, vendors)| vendors.len() > 1)
.filter(|(_, vendors)| vendors.as_slice() != ["grok", "supergrok"])
.collect();
assert!(
shared.is_empty(),
"these providers are indistinguishable in a bar that shows them \
side by side: {shared:#?}"
);
assert_eq!(VendorId::Anthropic.bar_icon(), "");
assert_eq!(VendorId::Openai.bar_icon(), "");
assert_eq!(VendorId::Supergrok.bar_icon(), VendorId::Grok.bar_icon());
assert_eq!(VendorId::CommandCode.bar_icon(), "cmc");
}
#[test]
fn new_vendor_contracts_keep_public_names_and_slugs() {
assert_eq!(VendorId::NousResearch.slug(), "nous");
assert_eq!(VendorId::NousResearch.display_name(), "Nous Research");
assert_eq!(VendorId::OpenCodeGo.slug(), "opencode-go");
assert_eq!(VendorId::OpenCodeGo.display_name(), "OpenCode Go");
assert_eq!(
serde_json::to_value(VendorId::OpenCodeGo).unwrap(),
serde_json::json!("opencode-go")
);
}
#[test]
fn vendor_secret_env_vars_cover_config_defaults() {
let configured_defaults = [
"ZAI_API_KEY",
"OPENROUTER_API_KEY",
"DEEPSEEK_API_KEY",
"KIMI_API_KEY",
"KILO_API_KEY",
"NOVITA_API_KEY",
"MINIMAX_API_KEY",
"MOONSHOT_API_KEY",
"XAI_MANAGEMENT_KEY",
"ANTHROPIC_ADMIN_KEY",
"GITHUB_COPILOT_TOKEN",
];
for name in configured_defaults {
assert!(VENDOR_SECRET_ENV_VARS.contains(&name), "missing {name}");
}
}
#[test]
fn vars_to_remove_preserves_only_requested_grok_credentials() {
let removed = vendor_secret_env_vars_to_remove(&["XAI_API_KEY", "GROK_API_KEY"]);
assert!(!removed.contains(&"XAI_API_KEY"));
assert!(!removed.contains(&"GROK_API_KEY"));
assert!(removed.contains(&"ANTHROPIC_ADMIN_KEY"));
assert!(removed.contains(&"OPENROUTER_API_KEY"));
let builtins = removed
.iter()
.filter(|var| VENDOR_SECRET_ENV_VARS.contains(var))
.count();
assert_eq!(builtins, VENDOR_SECRET_ENV_VARS.len() - 2);
}
#[test]
fn a_registered_custom_env_var_is_scrubbed_like_a_builtin_one() {
let name = "AI_USAGEBAR_TEST_CUSTOM_TOKEN_7F3A";
assert!(!vendor_secret_env_vars_to_remove(&[]).contains(&name));
register_secret_env_vars(&[name.to_string(), "not a name!".to_string()]);
register_secret_env_vars(&[name.to_string()]);
let removed = vendor_secret_env_vars_to_remove(&[]);
assert_eq!(
removed.iter().filter(|var| **var == name).count(),
1,
"registering twice must not list it twice: {removed:?}"
);
assert!(!removed.contains(&"not a name!"), "{removed:?}");
assert!(
!vendor_secret_env_vars_to_remove(&[name]).contains(&name),
"`keep` applies to registered names too"
);
}
#[test]
fn copilot_token_is_removed_before_unrelated_subprocesses_launch() {
let removed = vendor_secret_env_vars_to_remove(&[]);
assert!(removed.contains(&"GITHUB_COPILOT_TOKEN"));
}
#[tokio::test]
async fn body_over_the_cap_is_refused_and_under_it_round_trips() {
let mut server = mockito::Server::new_async().await;
server
.mock("GET", "/big")
.with_status(200)
.with_body("x".repeat(4096))
.create_async()
.await;
server
.mock("GET", "/small")
.with_status(200)
.with_body("hello")
.create_async()
.await;
let client = reqwest::Client::new();
let resp = client
.get(format!("{}/big", server.url()))
.send()
.await
.unwrap();
let err = read_body_capped(resp, 1024).await.unwrap_err();
assert!(
err.to_string().contains("exceeds"),
"unexpected error: {err}"
);
let resp = client
.get(format!("{}/small", server.url()))
.send()
.await
.unwrap();
assert_eq!(read_body_capped(resp, 1024).await.unwrap(), b"hello");
}
#[tokio::test]
async fn chunked_body_without_content_length_still_hits_the_cap() {
let mut server = mockito::Server::new_async().await;
server
.mock("GET", "/chunked")
.with_status(200)
.with_chunked_body(|writer| writer.write_all(&[b'x'; 4096]))
.create_async()
.await;
let response = reqwest::Client::new()
.get(format!("{}/chunked", server.url()))
.send()
.await
.unwrap();
assert!(response.content_length().is_none());
let error = read_body_capped(response, 1024).await.unwrap_err();
assert!(error.to_string().contains("exceeds"), "{error}");
}
#[tokio::test]
async fn same_origin_redirects_still_work_with_vendor_headers() {
let mut server = mockito::Server::new_async().await;
let redirect = server
.mock("GET", "/start")
.match_header("x-api-key", "secret")
.with_status(302)
.with_header("location", "/finish")
.create_async()
.await;
let finish = server
.mock("GET", "/finish")
.match_header("x-api-key", "secret")
.with_status(200)
.create_async()
.await;
let client = reqwest::Client::builder()
.redirect(same_origin_redirect_policy())
.build()
.unwrap();
let response = client
.get(format!("{}/start", server.url()))
.header("x-api-key", "secret")
.send()
.await
.unwrap();
assert_eq!(response.status(), reqwest::StatusCode::OK);
redirect.assert_async().await;
finish.assert_async().await;
}
#[tokio::test]
async fn cross_origin_redirects_are_not_followed_with_vendor_headers() {
let mut origin = mockito::Server::new_async().await;
let mut target = mockito::Server::new_async().await;
let target_url = format!("{}/capture", target.url());
let redirect = origin
.mock("GET", "/start")
.match_header("x-api-key", "secret")
.with_status(302)
.with_header("location", &target_url)
.create_async()
.await;
let capture = target
.mock("GET", "/capture")
.expect(0)
.create_async()
.await;
let client = reqwest::Client::builder()
.redirect(same_origin_redirect_policy())
.build()
.unwrap();
let response = client
.get(format!("{}/start", origin.url()))
.header("x-api-key", "secret")
.send()
.await
.unwrap();
assert_eq!(response.status(), reqwest::StatusCode::FOUND);
redirect.assert_async().await;
capture.assert_async().await;
}
#[test]
fn vendor_id_slug_round_trip() {
for id in VendorId::all() {
assert_eq!(
id.slug(),
serde_json::to_value(id).unwrap().as_str().unwrap()
);
}
}
}