use std::panic::{AssertUnwindSafe, catch_unwind};
use std::path::{Path, PathBuf};
use serde::{Deserialize, Serialize};
use crate::config::Config;
use crate::error::{AppError, Result};
use crate::vendor::VendorId;
pub fn has_local_credentials(vendor: VendorId, config: &Config) -> bool {
match vendor {
VendorId::Anthropic => anthropic_present(config),
VendorId::AnthropicApi => key_present(config, vendor),
VendorId::Openai => config
.openai
.resolve_auth_path(None)
.is_ok_and(|path| crate::openai::creds::read_from(&path).is_ok()),
VendorId::Copilot => copilot_present(),
VendorId::Zai => key_present(config, vendor),
VendorId::Openrouter => config.openrouter.resolve_api_key(None).is_ok(),
VendorId::Deepseek => key_present(config, vendor),
VendorId::Kimi => crate::kimi::resolve_auth(&config.kimi).is_ok(),
VendorId::Kilo => key_present(config, vendor),
VendorId::Novita => key_present(config, vendor),
VendorId::Moonshot => key_present(config, vendor),
VendorId::Grok => key_present(config, vendor),
VendorId::Supergrok => crate::supergrok::scope::ScopePaths::with_overrides(
config.supergrok.auth_path.as_deref(),
config.supergrok.config_path.as_deref(),
)
.is_ok_and(|paths| crate::supergrok::direct::read_billing_key(&paths.auth).is_ok()),
VendorId::Antigravity => antigravity_present(),
VendorId::Cursor => cursor_present(config),
VendorId::Minimax => key_present(config, vendor),
VendorId::Kiro => {
let path = match config.kiro.db_path.clone() {
Some(path) => path,
None => match crate::kiro::db::default_db_path() {
Ok(path) => path,
Err(_) => return false,
},
};
crate::kiro::db::read_credentials(&path).is_ok()
}
VendorId::NousResearch => {
let store = crate::nous::credentials::CredentialStore::at(
crate::nous::credentials::default_credentials_path(),
);
matches!(store.read_unlocked(), Ok(Some(_)))
}
VendorId::OpenCodeGo => key_present(config, vendor),
VendorId::CommandCode => {
crate::commandcode::creds::resolve(config.commandcode.auth_paths.as_deref()).is_ok()
}
VendorId::Ollama => key_present(config, vendor),
}
}
fn key_present(config: &Config, vendor: VendorId) -> bool {
crate::config::optional_api_key(
config.api_key_env_for(vendor),
config.inline_api_key(vendor),
)
.is_some()
}
fn anthropic_present(config: &Config) -> bool {
use crate::anthropic::creds::{CredsTarget, default_path, resolve};
let target = match config.anthropic.credentials_path.clone() {
Some(path) => CredsTarget::Explicit(path),
None => match default_path() {
Ok(path) => CredsTarget::Default(path),
Err(_) => return false,
},
};
resolve(&target).is_ok()
}
fn copilot_present() -> bool {
if std::env::var_os("GITHUB_COPILOT_TOKEN").is_some_and(|value| !value.is_empty()) {
return true;
}
crate::copilot::credentials::default_hosts_path()
.is_ok_and(|path| copilot_hosts_present_at(&path))
}
pub(crate) fn copilot_hosts_present_at(path: &Path) -> bool {
std::fs::metadata(path).is_ok_and(|meta| meta.is_file() && meta.len() > 0)
}
fn antigravity_present() -> bool {
if std::env::var_os("ANTIGRAVITY_LS_ADDRESS").is_some_and(|value| !value.is_empty()) {
return true;
}
if !crate::antigravity::fetch::discover_ls_ports().is_empty() {
return true;
}
crate::cache::Cache::for_vendor(crate::vendor::VendorId::Antigravity.slug()).is_ok_and(
|cache| {
crate::antigravity::cloud::has_persisted_session(
&crate::antigravity::cloud::oauth_cache_path(&cache),
)
},
)
}
fn cursor_present(config: &Config) -> bool {
let db_path = match config.cursor.db_path.clone() {
Some(path) => path,
None => match crate::cursor::db::default_db_path() {
Ok(path) => path,
Err(_) => return false,
},
};
let agent_auth_path = match config.cursor.agent_auth_path.clone() {
Some(path) => path,
None => match crate::cursor::db::default_agent_auth_path() {
Ok(path) => path,
Err(_) => return false,
},
};
crate::cursor::db::resolve_access_token(&db_path, &agent_auth_path).is_ok()
}
#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DetectState {
#[serde(default)]
pub known: Vec<VendorId>,
}
impl DetectState {
pub fn load_at(path: &Path) -> DetectState {
std::fs::read(path)
.ok()
.and_then(|bytes| serde_json::from_slice(&bytes).ok())
.unwrap_or_default()
}
pub fn save_at(&self, path: &Path) -> Result<()> {
let bytes = serde_json::to_vec_pretty(self)?;
crate::cache::atomic_write(path, &bytes)
}
}
pub fn default_state_path() -> Result<PathBuf> {
Ok(crate::cache::xdg_cache_dir()?
.join("ai-usagebar")
.join("detect.json"))
}
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct DetectPlan {
pub enable: Vec<VendorId>,
pub known: Vec<VendorId>,
pub probed: usize,
}
#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize)]
pub struct DetectReport {
pub enabled: Vec<VendorId>,
pub known: Vec<VendorId>,
pub probed: usize,
}
pub fn plan(
config: &Config,
state: &DetectState,
all: &[VendorId],
force: bool,
probe: impl Fn(VendorId) -> bool,
) -> DetectPlan {
let candidates: Vec<VendorId> = all
.iter()
.copied()
.filter(|vendor| force || !state.known.contains(vendor))
.collect();
let probed = candidates.len();
let enable = candidates
.into_iter()
.filter(|vendor| !config.is_enabled(*vendor))
.filter(|vendor| probe(*vendor))
.collect();
let known = VendorId::all()
.iter()
.copied()
.filter(|vendor| state.known.contains(vendor) || all.contains(vendor))
.collect();
DetectPlan {
enable,
known,
probed,
}
}
pub fn run_once(
config_path: Option<&Path>,
state_path: &Path,
force: bool,
) -> Result<Vec<VendorId>> {
run_once_report(config_path, state_path, force).map(|report| report.enabled)
}
pub fn run_once_report(
config_path: Option<&Path>,
state_path: &Path,
force: bool,
) -> Result<DetectReport> {
run_once_with(config_path, state_path, force, |vendor, config| {
catch_unwind(AssertUnwindSafe(|| has_local_credentials(vendor, config))).unwrap_or(false)
})
}
pub fn run_cli(all: bool, json: bool) -> i32 {
let report =
default_state_path().and_then(|state_path| run_once_report(None, &state_path, all));
match report {
Ok(report) if json => match serde_json::to_string(&report) {
Ok(text) => {
println!("{text}");
0
}
Err(error) => {
eprintln!("ai-usagebar detect: {error}");
1
}
},
Ok(report) => {
println!(
"{}",
format_report(&report, &crate::config::config_path_hint())
);
0
}
Err(error) => {
eprintln!("ai-usagebar detect: {}", error.user_message());
1
}
}
}
pub fn format_report(report: &DetectReport, config_hint: &str) -> String {
if report.enabled.is_empty() {
let noun = if report.probed == 1 {
"vendor"
} else {
"vendors"
};
return format!("Nothing new detected ({} {noun} checked)", report.probed);
}
let names: Vec<&str> = report
.enabled
.iter()
.map(|vendor| vendor.display_name())
.collect();
format!("Enabled: {}\nWritten to {config_hint}", names.join(", "))
}
pub fn run_once_with(
config_path: Option<&Path>,
state_path: &Path,
force: bool,
probe: impl Fn(VendorId, &Config) -> bool,
) -> Result<DetectReport> {
let resolved = match config_path {
Some(path) => Some(path.to_path_buf()),
None => crate::config::resolved_path(),
};
let config = match &resolved {
Some(path) => Config::load_from(path)?,
None => Config::default(),
};
let state = DetectState::load_at(state_path);
let plan = plan(&config, &state, VendorId::all(), force, |vendor| {
probe(vendor, &config)
});
let enabled = if plan.enable.is_empty() {
Vec::new()
} else {
let path = resolved.ok_or_else(|| {
AppError::Other("could not resolve the config.toml path to enable vendors in".into())
})?;
crate::config::enable_vendors_in(&path, &plan.enable)?
};
DetectState {
known: plan.known.clone(),
}
.save_at(state_path)?;
Ok(DetectReport {
enabled,
known: plan.known,
probed: plan.probed,
})
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
fn probe_in(present: &[VendorId]) -> impl Fn(VendorId) -> bool + '_ {
move |vendor| present.contains(&vendor)
}
#[test]
fn state_round_trips_through_json_by_slug() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("nested").join("detect.json");
let state = DetectState {
known: vec![
VendorId::Cursor,
VendorId::OpenCodeGo,
VendorId::NousResearch,
],
};
state.save_at(&path).unwrap();
let text = std::fs::read_to_string(&path).unwrap();
assert!(text.contains("\"opencode-go\""), "{text}");
assert!(text.contains("\"nous\""), "{text}");
assert_eq!(DetectState::load_at(&path), state);
}
#[test]
fn missing_or_corrupt_state_is_the_default() {
let dir = TempDir::new().unwrap();
assert_eq!(
DetectState::load_at(&dir.path().join("absent.json")),
DetectState::default()
);
let corrupt = dir.path().join("corrupt.json");
std::fs::write(&corrupt, "{\"known\": [\"not-a-vendor\"").unwrap();
assert_eq!(DetectState::load_at(&corrupt), DetectState::default());
let unknown_slug = dir.path().join("unknown.json");
std::fs::write(&unknown_slug, "{\"known\": [\"not-a-vendor\"]}").unwrap();
assert_eq!(DetectState::load_at(&unknown_slug), DetectState::default());
}
#[test]
fn plan_enables_only_unknown_probed_vendors_that_are_off() {
let config = Config::default(); let state = DetectState {
known: vec![VendorId::Grok],
};
let all = [
VendorId::Anthropic, VendorId::Grok, VendorId::Cursor, VendorId::Kiro, ];
let present = [VendorId::Anthropic, VendorId::Grok, VendorId::Cursor];
let plan = plan(&config, &state, &all, false, probe_in(&present));
assert_eq!(plan.enable, vec![VendorId::Cursor]);
}
#[test]
fn force_reconsiders_known_vendors_but_never_enabled_ones() {
let config = Config::default();
let state = DetectState {
known: vec![VendorId::Grok, VendorId::Zai],
};
let all = [VendorId::Zai, VendorId::Grok];
let present = [VendorId::Zai, VendorId::Grok];
let plan = plan(&config, &state, &all, true, probe_in(&present));
assert_eq!(plan.enable, vec![VendorId::Grok]);
}
#[test]
fn plan_orders_enable_by_the_candidate_list() {
let config = Config::default();
let all = [VendorId::Kiro, VendorId::Cursor, VendorId::Grok];
let present = [VendorId::Grok, VendorId::Cursor, VendorId::Kiro];
let plan = plan(
&config,
&DetectState::default(),
&all,
false,
probe_in(&present),
);
assert_eq!(
plan.enable,
vec![VendorId::Kiro, VendorId::Cursor, VendorId::Grok]
);
}
#[test]
fn known_becomes_the_union_in_canonical_order_without_duplicates() {
let config = Config::default();
let state = DetectState {
known: vec![VendorId::Grok, VendorId::Cursor],
};
let all = [VendorId::Cursor, VendorId::Anthropic, VendorId::Cursor];
let plan = plan(&config, &state, &all, false, |_| false);
assert_eq!(
plan.known,
vec![VendorId::Anthropic, VendorId::Grok, VendorId::Cursor]
);
assert!(plan.enable.is_empty());
}
#[test]
fn probe_is_not_consulted_for_skipped_vendors() {
let config = Config::default();
let state = DetectState {
known: vec![VendorId::Grok],
};
let all = [VendorId::Grok, VendorId::Anthropic];
let plan = plan(&config, &state, &all, false, |vendor| {
panic!("probe called for {}", vendor.slug())
});
assert!(plan.enable.is_empty());
}
#[test]
fn copilot_hosts_file_must_be_a_non_empty_regular_file() {
let dir = TempDir::new().unwrap();
let hosts = dir.path().join("hosts.yml");
assert!(!copilot_hosts_present_at(&hosts));
std::fs::write(&hosts, "").unwrap();
assert!(!copilot_hosts_present_at(&hosts));
std::fs::write(&hosts, "github.com:\n user: octocat\n").unwrap();
assert!(copilot_hosts_present_at(&hosts));
assert!(!copilot_hosts_present_at(dir.path()));
}
#[test]
fn run_once_writes_enables_into_the_config_and_marks_everything_known() {
let dir = TempDir::new().unwrap();
let config_path = dir.path().join("config.toml");
let state_path = dir.path().join("detect.json");
std::fs::write(
&config_path,
"# mine
[zai]
enabled = false
",
)
.unwrap();
let present = [VendorId::Zai, VendorId::Cursor, VendorId::Anthropic];
let probe = |vendor: VendorId, _: &Config| present.contains(&vendor);
let report = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
assert_eq!(report.enabled, vec![VendorId::Cursor]);
assert_eq!(report.known, VendorId::all());
assert_eq!(report.probed, VendorId::all().len());
let after = Config::load_from(&config_path).unwrap();
assert!(!after.is_enabled(VendorId::Zai), "an opt-out must survive");
assert!(after.is_enabled(VendorId::Cursor));
let text = std::fs::read_to_string(&config_path).unwrap();
assert!(
text.starts_with(
"# mine
"
),
"{text}"
);
assert_eq!(DetectState::load_at(&state_path).known, VendorId::all());
let again = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
assert!(again.enabled.is_empty(), "{again:?}");
assert_eq!(again.probed, 0, "everything is known: nothing to check");
std::fs::write(
&config_path,
"[cursor]
enabled = false
",
)
.unwrap();
let third = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
assert!(third.enabled.is_empty(), "{third:?}");
assert!(
!Config::load_from(&config_path)
.unwrap()
.is_enabled(VendorId::Cursor)
);
let forced = run_once_with(Some(&config_path), &state_path, true, probe).unwrap();
assert!(forced.enabled.is_empty(), "{forced:?}");
assert_eq!(forced.probed, VendorId::all().len());
assert!(
!Config::load_from(&config_path)
.unwrap()
.is_enabled(VendorId::Cursor)
);
}
#[test]
fn plan_counts_candidates_not_enables() {
let config = Config::default();
let state = DetectState {
known: vec![VendorId::Grok],
};
let all = [VendorId::Anthropic, VendorId::Grok, VendorId::Cursor];
let unforced = plan(&config, &state, &all, false, |_| false);
assert_eq!(unforced.probed, 2, "Grok is known and skipped");
let forced = plan(&config, &state, &all, true, |_| false);
assert_eq!(forced.probed, 3);
}
#[test]
fn format_report_lists_display_names_and_where_they_were_written() {
let report = DetectReport {
enabled: vec![VendorId::Cursor, VendorId::Kiro],
known: VendorId::all().to_vec(),
probed: 3,
};
let text = format_report(&report, "/home/u/.config/ai-usagebar/config.toml");
assert_eq!(
text,
"Enabled: Cursor, Kiro\nWritten to /home/u/.config/ai-usagebar/config.toml"
);
}
#[test]
fn format_report_says_how_many_were_checked_when_nothing_changed() {
let none = DetectReport {
enabled: vec![],
known: VendorId::all().to_vec(),
probed: 3,
};
assert_eq!(
format_report(&none, "unused"),
"Nothing new detected (3 vendors checked)"
);
let one = DetectReport {
probed: 1,
..none.clone()
};
assert_eq!(
format_report(&one, "unused"),
"Nothing new detected (1 vendor checked)"
);
}
#[test]
fn report_serializes_slugs_and_the_probed_count() {
let dir = TempDir::new().unwrap();
let config_path = dir.path().join("config.toml");
let state_path = dir.path().join("detect.json");
std::fs::write(&config_path, "").unwrap();
DetectState {
known: vec![VendorId::Anthropic, VendorId::Grok],
}
.save_at(&state_path)
.unwrap();
let probe = |vendor: VendorId, _: &Config| vendor == VendorId::Cursor;
let report = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
let json: serde_json::Value = serde_json::to_value(&report).unwrap();
assert_eq!(json["enabled"], serde_json::json!(["cursor"]));
assert_eq!(
json["probed"],
serde_json::json!(VendorId::all().len() - 2),
"the two known vendors were not candidates"
);
let known = json["known"].as_array().unwrap();
assert_eq!(known.len(), VendorId::all().len());
assert_eq!(known[0], serde_json::json!("anthropic"));
assert_eq!(json.as_object().unwrap().len(), 3, "{json}");
}
#[test]
fn a_panicking_probe_counts_as_absent_and_still_saves_state() {
let dir = TempDir::new().unwrap();
let config_path = dir.path().join("config.toml");
let state_path = dir.path().join("detect.json");
let probe = |vendor: VendorId, _: &Config| {
catch_unwind(AssertUnwindSafe(|| {
if vendor == VendorId::Kiro {
panic!("boom");
}
vendor == VendorId::Grok
}))
.unwrap_or(false)
};
let enabled = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
assert_eq!(enabled.enabled, vec![VendorId::Grok]);
assert_eq!(DetectState::load_at(&state_path).known, VendorId::all());
}
}