1use std::panic::{AssertUnwindSafe, catch_unwind};
27use std::path::{Path, PathBuf};
28
29use serde::{Deserialize, Serialize};
30
31use crate::config::Config;
32use crate::error::{AppError, Result};
33use crate::vendor::VendorId;
34
35pub fn has_local_credentials(vendor: VendorId, config: &Config) -> bool {
42 match vendor {
43 VendorId::Anthropic => anthropic_present(config),
44 VendorId::AnthropicApi => key_present(config, vendor),
45 VendorId::Openai => config
46 .openai
47 .resolve_auth_path(None)
48 .is_ok_and(|path| crate::openai::creds::read_from(&path).is_ok()),
49 VendorId::Copilot => copilot_present(),
50 VendorId::Zai => key_present(config, vendor),
51 VendorId::Openrouter => config.openrouter.resolve_api_key(None).is_ok(),
52 VendorId::Deepseek => key_present(config, vendor),
53 VendorId::Kimi => crate::kimi::resolve_auth(&config.kimi).is_ok(),
54 VendorId::Kilo => key_present(config, vendor),
55 VendorId::Novita => key_present(config, vendor),
56 VendorId::Moonshot => key_present(config, vendor),
57 VendorId::Grok => key_present(config, vendor),
58 VendorId::Supergrok => crate::supergrok::scope::ScopePaths::with_overrides(
59 config.supergrok.auth_path.as_deref(),
60 config.supergrok.config_path.as_deref(),
61 )
62 .is_ok_and(|paths| crate::supergrok::direct::read_billing_key(&paths.auth).is_ok()),
63 VendorId::Grokbot => crate::grokbot::secrets_path(&config.grokbot)
67 .map(|path| crate::grokbot::creds::secrets_present_at(&path))
68 .unwrap_or(false),
69 VendorId::Antigravity => antigravity_present(),
70 VendorId::Cursor => cursor_present(config),
71 VendorId::Minimax => key_present(config, vendor),
72 VendorId::Kiro => {
73 let path = match config.kiro.db_path.clone() {
74 Some(path) => path,
75 None => match crate::kiro::db::default_db_path() {
76 Ok(path) => path,
77 Err(_) => return false,
78 },
79 };
80 crate::kiro::db::read_credentials(&path).is_ok()
81 }
82 VendorId::NousResearch => {
83 let store = crate::nous::credentials::CredentialStore::at(
87 crate::nous::credentials::default_credentials_path(),
88 );
89 matches!(store.read_unlocked(), Ok(Some(_)))
90 }
91 VendorId::OpenCodeGo => key_present(config, vendor),
92 VendorId::CommandCode => {
93 crate::commandcode::creds::resolve(config.commandcode.auth_paths.as_deref()).is_ok()
94 }
95 VendorId::Ollama => key_present(config, vendor),
96 VendorId::OrcaRouter => key_present(config, vendor),
97 VendorId::ModelStudio => crate::modelstudio::config_path(&config.modelstudio)
100 .map(|path| crate::modelstudio::creds::config_present_at(&path))
101 .unwrap_or(false),
102 }
103}
104
105fn key_present(config: &Config, vendor: VendorId) -> bool {
110 crate::config::optional_api_key(
111 config.api_key_env_for(vendor),
112 config.inline_api_key(vendor),
113 )
114 .is_some()
115}
116
117fn anthropic_present(config: &Config) -> bool {
121 use crate::anthropic::creds::{CredsTarget, default_path, resolve};
122 let target = match config.anthropic.credentials_path.clone() {
123 Some(path) => CredsTarget::Explicit(path),
124 None => match default_path() {
125 Ok(path) => CredsTarget::Default(path),
126 Err(_) => return false,
127 },
128 };
129 resolve(&target).is_ok()
130}
131
132fn copilot_present() -> bool {
142 if std::env::var_os("GITHUB_COPILOT_TOKEN").is_some_and(|value| !value.is_empty()) {
143 return true;
144 }
145 crate::copilot::credentials::default_hosts_path()
146 .is_ok_and(|path| copilot_hosts_present_at(&path))
147}
148
149pub(crate) fn copilot_hosts_present_at(path: &Path) -> bool {
152 std::fs::metadata(path).is_ok_and(|meta| meta.is_file() && meta.len() > 0)
153}
154
155fn antigravity_present() -> bool {
159 if std::env::var_os("ANTIGRAVITY_LS_ADDRESS").is_some_and(|value| !value.is_empty()) {
160 return true;
161 }
162 if !crate::antigravity::fetch::discover_ls_ports().is_empty() {
163 return true;
164 }
165 crate::cache::Cache::for_vendor(crate::vendor::VendorId::Antigravity.slug()).is_ok_and(
172 |cache| {
173 crate::antigravity::cloud::has_persisted_session(
174 &crate::antigravity::cloud::oauth_cache_path(&cache),
175 )
176 },
177 )
178}
179
180fn cursor_present(config: &Config) -> bool {
181 let db_path = match config.cursor.db_path.clone() {
182 Some(path) => path,
183 None => match crate::cursor::db::default_db_path() {
184 Ok(path) => path,
185 Err(_) => return false,
186 },
187 };
188 let agent_auth_path = match config.cursor.agent_auth_path.clone() {
189 Some(path) => path,
190 None => match crate::cursor::db::default_agent_auth_path() {
191 Ok(path) => path,
192 Err(_) => return false,
193 },
194 };
195 crate::cursor::db::resolve_access_token(&db_path, &agent_auth_path).is_ok()
196}
197
198#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)]
202pub struct DetectState {
203 #[serde(default)]
204 pub known: Vec<VendorId>,
205}
206
207impl DetectState {
208 pub fn load_at(path: &Path) -> DetectState {
212 std::fs::read(path)
213 .ok()
214 .and_then(|bytes| serde_json::from_slice(&bytes).ok())
215 .unwrap_or_default()
216 }
217
218 pub fn save_at(&self, path: &Path) -> Result<()> {
220 let bytes = serde_json::to_vec_pretty(self)?;
221 crate::cache::atomic_write(path, &bytes)
222 }
223}
224
225pub fn default_state_path() -> Result<PathBuf> {
228 Ok(crate::cache::xdg_cache_dir()?
229 .join("ai-usagebar")
230 .join("detect.json"))
231}
232
233#[derive(Debug, Default, Clone, PartialEq, Eq)]
235pub struct DetectPlan {
236 pub enable: Vec<VendorId>,
238 pub known: Vec<VendorId>,
240 pub probed: usize,
243}
244
245#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize)]
248pub struct DetectReport {
249 pub enabled: Vec<VendorId>,
251 pub known: Vec<VendorId>,
253 pub probed: usize,
255}
256
257pub fn plan(
263 config: &Config,
264 state: &DetectState,
265 all: &[VendorId],
266 force: bool,
267 probe: impl Fn(VendorId) -> bool,
268) -> DetectPlan {
269 let candidates: Vec<VendorId> = all
270 .iter()
271 .copied()
272 .filter(|vendor| force || !state.known.contains(vendor))
273 .collect();
274 let probed = candidates.len();
275 let enable = candidates
276 .into_iter()
277 .filter(|vendor| !config.is_enabled(*vendor))
278 .filter(|vendor| probe(*vendor))
279 .collect();
280 let known = VendorId::all()
281 .iter()
282 .copied()
283 .filter(|vendor| state.known.contains(vendor) || all.contains(vendor))
284 .collect();
285 DetectPlan {
286 enable,
287 known,
288 probed,
289 }
290}
291
292pub fn run_once(
301 config_path: Option<&Path>,
302 state_path: &Path,
303 force: bool,
304) -> Result<Vec<VendorId>> {
305 run_once_report(config_path, state_path, force).map(|report| report.enabled)
306}
307
308pub fn run_once_report(
311 config_path: Option<&Path>,
312 state_path: &Path,
313 force: bool,
314) -> Result<DetectReport> {
315 run_once_with(config_path, state_path, force, |vendor, config| {
316 catch_unwind(AssertUnwindSafe(|| has_local_credentials(vendor, config))).unwrap_or(false)
317 })
318}
319
320pub fn run_cli(all: bool, json: bool) -> i32 {
325 let report =
326 default_state_path().and_then(|state_path| run_once_report(None, &state_path, all));
327 match report {
328 Ok(report) if json => match serde_json::to_string(&report) {
329 Ok(text) => {
330 println!("{text}");
331 0
332 }
333 Err(error) => {
334 eprintln!("ai-usagebar detect: {error}");
335 1
336 }
337 },
338 Ok(report) => {
339 println!(
340 "{}",
341 format_report(&report, &crate::config::config_path_hint())
342 );
343 0
344 }
345 Err(error) => {
346 eprintln!("ai-usagebar detect: {}", error.user_message());
347 1
348 }
349 }
350}
351
352pub fn format_report(report: &DetectReport, config_hint: &str) -> String {
355 if report.enabled.is_empty() {
356 let noun = if report.probed == 1 {
357 "vendor"
358 } else {
359 "vendors"
360 };
361 return format!("Nothing new detected ({} {noun} checked)", report.probed);
362 }
363 let names: Vec<&str> = report
364 .enabled
365 .iter()
366 .map(|vendor| vendor.display_name())
367 .collect();
368 format!("Enabled: {}\nWritten to {config_hint}", names.join(", "))
369}
370
371pub fn run_once_with(
375 config_path: Option<&Path>,
376 state_path: &Path,
377 force: bool,
378 probe: impl Fn(VendorId, &Config) -> bool,
379) -> Result<DetectReport> {
380 let resolved = match config_path {
381 Some(path) => Some(path.to_path_buf()),
382 None => crate::config::resolved_path(),
383 };
384 let config = match &resolved {
385 Some(path) => Config::load_from(path)?,
386 None => Config::default(),
387 };
388 let state = DetectState::load_at(state_path);
389 let plan = plan(&config, &state, VendorId::all(), force, |vendor| {
390 probe(vendor, &config)
391 });
392 let enabled = if plan.enable.is_empty() {
396 Vec::new()
397 } else {
398 let path = resolved.ok_or_else(|| {
399 AppError::Other("could not resolve the config.toml path to enable vendors in".into())
400 })?;
401 crate::config::enable_vendors_in(&path, &plan.enable)?
402 };
403 DetectState {
404 known: plan.known.clone(),
405 }
406 .save_at(state_path)?;
407 Ok(DetectReport {
408 enabled,
409 known: plan.known,
410 probed: plan.probed,
411 })
412}
413
414#[cfg(test)]
415mod tests {
416 use super::*;
417 use tempfile::TempDir;
418
419 fn probe_in(present: &[VendorId]) -> impl Fn(VendorId) -> bool + '_ {
420 move |vendor| present.contains(&vendor)
421 }
422
423 #[test]
424 fn state_round_trips_through_json_by_slug() {
425 let dir = TempDir::new().unwrap();
426 let path = dir.path().join("nested").join("detect.json");
427 let state = DetectState {
428 known: vec![
429 VendorId::Cursor,
430 VendorId::OpenCodeGo,
431 VendorId::NousResearch,
432 ],
433 };
434
435 state.save_at(&path).unwrap();
436
437 let text = std::fs::read_to_string(&path).unwrap();
438 assert!(text.contains("\"opencode-go\""), "{text}");
439 assert!(text.contains("\"nous\""), "{text}");
440 assert_eq!(DetectState::load_at(&path), state);
441 }
442
443 #[test]
444 fn missing_or_corrupt_state_is_the_default() {
445 let dir = TempDir::new().unwrap();
446 assert_eq!(
447 DetectState::load_at(&dir.path().join("absent.json")),
448 DetectState::default()
449 );
450
451 let corrupt = dir.path().join("corrupt.json");
452 std::fs::write(&corrupt, "{\"known\": [\"not-a-vendor\"").unwrap();
453 assert_eq!(DetectState::load_at(&corrupt), DetectState::default());
454
455 let unknown_slug = dir.path().join("unknown.json");
456 std::fs::write(&unknown_slug, "{\"known\": [\"not-a-vendor\"]}").unwrap();
457 assert_eq!(DetectState::load_at(&unknown_slug), DetectState::default());
458 }
459
460 #[test]
461 fn plan_enables_only_unknown_probed_vendors_that_are_off() {
462 let config = Config::default(); let state = DetectState {
464 known: vec![VendorId::Grok],
465 };
466 let all = [
467 VendorId::Anthropic, VendorId::Grok, VendorId::Cursor, VendorId::Kiro, ];
472 let present = [VendorId::Anthropic, VendorId::Grok, VendorId::Cursor];
473
474 let plan = plan(&config, &state, &all, false, probe_in(&present));
475
476 assert_eq!(plan.enable, vec![VendorId::Cursor]);
477 }
478
479 #[test]
480 fn force_reconsiders_known_vendors_but_never_enabled_ones() {
481 let config = Config::default();
482 let state = DetectState {
483 known: vec![VendorId::Grok, VendorId::Zai],
484 };
485 let all = [VendorId::Zai, VendorId::Grok];
486 let present = [VendorId::Zai, VendorId::Grok];
487
488 let plan = plan(&config, &state, &all, true, probe_in(&present));
489
490 assert_eq!(plan.enable, vec![VendorId::Grok]);
491 }
492
493 #[test]
494 fn plan_orders_enable_by_the_candidate_list() {
495 let config = Config::default();
496 let all = [VendorId::Kiro, VendorId::Cursor, VendorId::Grok];
497 let present = [VendorId::Grok, VendorId::Cursor, VendorId::Kiro];
498
499 let plan = plan(
500 &config,
501 &DetectState::default(),
502 &all,
503 false,
504 probe_in(&present),
505 );
506
507 assert_eq!(
508 plan.enable,
509 vec![VendorId::Kiro, VendorId::Cursor, VendorId::Grok]
510 );
511 }
512
513 #[test]
514 fn known_becomes_the_union_in_canonical_order_without_duplicates() {
515 let config = Config::default();
516 let state = DetectState {
517 known: vec![VendorId::Grok, VendorId::Cursor],
518 };
519 let all = [VendorId::Cursor, VendorId::Anthropic, VendorId::Cursor];
520
521 let plan = plan(&config, &state, &all, false, |_| false);
522
523 assert_eq!(
524 plan.known,
525 vec![VendorId::Anthropic, VendorId::Grok, VendorId::Cursor]
526 );
527 assert!(plan.enable.is_empty());
528 }
529
530 #[test]
531 fn probe_is_not_consulted_for_skipped_vendors() {
532 let config = Config::default();
533 let state = DetectState {
534 known: vec![VendorId::Grok],
535 };
536 let all = [VendorId::Grok, VendorId::Anthropic];
537
538 let plan = plan(&config, &state, &all, false, |vendor| {
539 panic!("probe called for {}", vendor.slug())
540 });
541
542 assert!(plan.enable.is_empty());
543 }
544
545 #[test]
546 fn copilot_hosts_file_must_be_a_non_empty_regular_file() {
547 let dir = TempDir::new().unwrap();
548 let hosts = dir.path().join("hosts.yml");
549 assert!(!copilot_hosts_present_at(&hosts));
550
551 std::fs::write(&hosts, "").unwrap();
552 assert!(!copilot_hosts_present_at(&hosts));
553
554 std::fs::write(&hosts, "github.com:\n user: octocat\n").unwrap();
555 assert!(copilot_hosts_present_at(&hosts));
556
557 assert!(!copilot_hosts_present_at(dir.path()));
558 }
559
560 #[test]
564 fn run_once_writes_enables_into_the_config_and_marks_everything_known() {
565 let dir = TempDir::new().unwrap();
566 let config_path = dir.path().join("config.toml");
567 let state_path = dir.path().join("detect.json");
568 std::fs::write(
569 &config_path,
570 "# mine
571[zai]
572enabled = false
573",
574 )
575 .unwrap();
576 let present = [VendorId::Zai, VendorId::Cursor, VendorId::Anthropic];
577 let probe = |vendor: VendorId, _: &Config| present.contains(&vendor);
578
579 let report = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
580
581 assert_eq!(report.enabled, vec![VendorId::Cursor]);
585 assert_eq!(report.known, VendorId::all());
586 assert_eq!(report.probed, VendorId::all().len());
587 let after = Config::load_from(&config_path).unwrap();
588 assert!(!after.is_enabled(VendorId::Zai), "an opt-out must survive");
589 assert!(after.is_enabled(VendorId::Cursor));
590 let text = std::fs::read_to_string(&config_path).unwrap();
591 assert!(
592 text.starts_with(
593 "# mine
594"
595 ),
596 "{text}"
597 );
598 assert_eq!(DetectState::load_at(&state_path).known, VendorId::all());
599
600 let again = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
601 assert!(again.enabled.is_empty(), "{again:?}");
602 assert_eq!(again.probed, 0, "everything is known: nothing to check");
603
604 std::fs::write(
606 &config_path,
607 "[cursor]
608enabled = false
609",
610 )
611 .unwrap();
612 let third = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
613 assert!(third.enabled.is_empty(), "{third:?}");
614 assert!(
615 !Config::load_from(&config_path)
616 .unwrap()
617 .is_enabled(VendorId::Cursor)
618 );
619
620 let forced = run_once_with(Some(&config_path), &state_path, true, probe).unwrap();
625 assert!(forced.enabled.is_empty(), "{forced:?}");
626 assert_eq!(forced.probed, VendorId::all().len());
627 assert!(
628 !Config::load_from(&config_path)
629 .unwrap()
630 .is_enabled(VendorId::Cursor)
631 );
632 }
633
634 #[test]
635 fn plan_counts_candidates_not_enables() {
636 let config = Config::default();
637 let state = DetectState {
638 known: vec![VendorId::Grok],
639 };
640 let all = [VendorId::Anthropic, VendorId::Grok, VendorId::Cursor];
641
642 let unforced = plan(&config, &state, &all, false, |_| false);
643 assert_eq!(unforced.probed, 2, "Grok is known and skipped");
644
645 let forced = plan(&config, &state, &all, true, |_| false);
646 assert_eq!(forced.probed, 3);
647 }
648
649 #[test]
650 fn format_report_lists_display_names_and_where_they_were_written() {
651 let report = DetectReport {
652 enabled: vec![VendorId::Cursor, VendorId::Kiro],
653 known: VendorId::all().to_vec(),
654 probed: 3,
655 };
656
657 let text = format_report(&report, "/home/u/.config/ai-usagebar/config.toml");
658
659 assert_eq!(
660 text,
661 "Enabled: Cursor, Kiro\nWritten to /home/u/.config/ai-usagebar/config.toml"
662 );
663 }
664
665 #[test]
666 fn format_report_says_how_many_were_checked_when_nothing_changed() {
667 let none = DetectReport {
668 enabled: vec![],
669 known: VendorId::all().to_vec(),
670 probed: 3,
671 };
672 assert_eq!(
673 format_report(&none, "unused"),
674 "Nothing new detected (3 vendors checked)"
675 );
676
677 let one = DetectReport {
678 probed: 1,
679 ..none.clone()
680 };
681 assert_eq!(
682 format_report(&one, "unused"),
683 "Nothing new detected (1 vendor checked)"
684 );
685 }
686
687 #[test]
689 fn report_serializes_slugs_and_the_probed_count() {
690 let dir = TempDir::new().unwrap();
691 let config_path = dir.path().join("config.toml");
692 let state_path = dir.path().join("detect.json");
693 std::fs::write(&config_path, "").unwrap();
694 DetectState {
695 known: vec![VendorId::Anthropic, VendorId::Grok],
696 }
697 .save_at(&state_path)
698 .unwrap();
699 let probe = |vendor: VendorId, _: &Config| vendor == VendorId::Cursor;
700
701 let report = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
702 let json: serde_json::Value = serde_json::to_value(&report).unwrap();
703
704 assert_eq!(json["enabled"], serde_json::json!(["cursor"]));
705 assert_eq!(
706 json["probed"],
707 serde_json::json!(VendorId::all().len() - 2),
708 "the two known vendors were not candidates"
709 );
710 let known = json["known"].as_array().unwrap();
711 assert_eq!(known.len(), VendorId::all().len());
712 assert_eq!(known[0], serde_json::json!("anthropic"));
713 assert_eq!(json.as_object().unwrap().len(), 3, "{json}");
714 }
715
716 #[test]
717 fn a_panicking_probe_counts_as_absent_and_still_saves_state() {
718 let dir = TempDir::new().unwrap();
719 let config_path = dir.path().join("config.toml");
720 let state_path = dir.path().join("detect.json");
721 let probe = |vendor: VendorId, _: &Config| {
722 catch_unwind(AssertUnwindSafe(|| {
723 if vendor == VendorId::Kiro {
724 panic!("boom");
725 }
726 vendor == VendorId::Grok
727 }))
728 .unwrap_or(false)
729 };
730
731 let enabled = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
732
733 assert_eq!(enabled.enabled, vec![VendorId::Grok]);
734 assert_eq!(DetectState::load_at(&state_path).known, VendorId::all());
735 }
736}