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 }
97}
98
99fn key_present(config: &Config, vendor: VendorId) -> bool {
104 crate::config::optional_api_key(
105 config.api_key_env_for(vendor),
106 config.inline_api_key(vendor),
107 )
108 .is_some()
109}
110
111fn anthropic_present(config: &Config) -> bool {
115 use crate::anthropic::creds::{CredsTarget, default_path, resolve};
116 let target = match config.anthropic.credentials_path.clone() {
117 Some(path) => CredsTarget::Explicit(path),
118 None => match default_path() {
119 Ok(path) => CredsTarget::Default(path),
120 Err(_) => return false,
121 },
122 };
123 resolve(&target).is_ok()
124}
125
126fn copilot_present() -> bool {
136 if std::env::var_os("GITHUB_COPILOT_TOKEN").is_some_and(|value| !value.is_empty()) {
137 return true;
138 }
139 crate::copilot::credentials::default_hosts_path()
140 .is_ok_and(|path| copilot_hosts_present_at(&path))
141}
142
143pub(crate) fn copilot_hosts_present_at(path: &Path) -> bool {
146 std::fs::metadata(path).is_ok_and(|meta| meta.is_file() && meta.len() > 0)
147}
148
149fn antigravity_present() -> bool {
153 if std::env::var_os("ANTIGRAVITY_LS_ADDRESS").is_some_and(|value| !value.is_empty()) {
154 return true;
155 }
156 if !crate::antigravity::fetch::discover_ls_ports().is_empty() {
157 return true;
158 }
159 crate::cache::Cache::for_vendor(crate::vendor::VendorId::Antigravity.slug()).is_ok_and(
166 |cache| {
167 crate::antigravity::cloud::has_persisted_session(
168 &crate::antigravity::cloud::oauth_cache_path(&cache),
169 )
170 },
171 )
172}
173
174fn cursor_present(config: &Config) -> bool {
175 let db_path = match config.cursor.db_path.clone() {
176 Some(path) => path,
177 None => match crate::cursor::db::default_db_path() {
178 Ok(path) => path,
179 Err(_) => return false,
180 },
181 };
182 let agent_auth_path = match config.cursor.agent_auth_path.clone() {
183 Some(path) => path,
184 None => match crate::cursor::db::default_agent_auth_path() {
185 Ok(path) => path,
186 Err(_) => return false,
187 },
188 };
189 crate::cursor::db::resolve_access_token(&db_path, &agent_auth_path).is_ok()
190}
191
192#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)]
196pub struct DetectState {
197 #[serde(default)]
198 pub known: Vec<VendorId>,
199}
200
201impl DetectState {
202 pub fn load_at(path: &Path) -> DetectState {
206 std::fs::read(path)
207 .ok()
208 .and_then(|bytes| serde_json::from_slice(&bytes).ok())
209 .unwrap_or_default()
210 }
211
212 pub fn save_at(&self, path: &Path) -> Result<()> {
214 let bytes = serde_json::to_vec_pretty(self)?;
215 crate::cache::atomic_write(path, &bytes)
216 }
217}
218
219pub fn default_state_path() -> Result<PathBuf> {
222 Ok(crate::cache::xdg_cache_dir()?
223 .join("ai-usagebar")
224 .join("detect.json"))
225}
226
227#[derive(Debug, Default, Clone, PartialEq, Eq)]
229pub struct DetectPlan {
230 pub enable: Vec<VendorId>,
232 pub known: Vec<VendorId>,
234 pub probed: usize,
237}
238
239#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize)]
242pub struct DetectReport {
243 pub enabled: Vec<VendorId>,
245 pub known: Vec<VendorId>,
247 pub probed: usize,
249}
250
251pub fn plan(
257 config: &Config,
258 state: &DetectState,
259 all: &[VendorId],
260 force: bool,
261 probe: impl Fn(VendorId) -> bool,
262) -> DetectPlan {
263 let candidates: Vec<VendorId> = all
264 .iter()
265 .copied()
266 .filter(|vendor| force || !state.known.contains(vendor))
267 .collect();
268 let probed = candidates.len();
269 let enable = candidates
270 .into_iter()
271 .filter(|vendor| !config.is_enabled(*vendor))
272 .filter(|vendor| probe(*vendor))
273 .collect();
274 let known = VendorId::all()
275 .iter()
276 .copied()
277 .filter(|vendor| state.known.contains(vendor) || all.contains(vendor))
278 .collect();
279 DetectPlan {
280 enable,
281 known,
282 probed,
283 }
284}
285
286pub fn run_once(
295 config_path: Option<&Path>,
296 state_path: &Path,
297 force: bool,
298) -> Result<Vec<VendorId>> {
299 run_once_report(config_path, state_path, force).map(|report| report.enabled)
300}
301
302pub fn run_once_report(
305 config_path: Option<&Path>,
306 state_path: &Path,
307 force: bool,
308) -> Result<DetectReport> {
309 run_once_with(config_path, state_path, force, |vendor, config| {
310 catch_unwind(AssertUnwindSafe(|| has_local_credentials(vendor, config))).unwrap_or(false)
311 })
312}
313
314pub fn run_cli(all: bool, json: bool) -> i32 {
319 let report =
320 default_state_path().and_then(|state_path| run_once_report(None, &state_path, all));
321 match report {
322 Ok(report) if json => match serde_json::to_string(&report) {
323 Ok(text) => {
324 println!("{text}");
325 0
326 }
327 Err(error) => {
328 eprintln!("ai-usagebar detect: {error}");
329 1
330 }
331 },
332 Ok(report) => {
333 println!(
334 "{}",
335 format_report(&report, &crate::config::config_path_hint())
336 );
337 0
338 }
339 Err(error) => {
340 eprintln!("ai-usagebar detect: {}", error.user_message());
341 1
342 }
343 }
344}
345
346pub fn format_report(report: &DetectReport, config_hint: &str) -> String {
349 if report.enabled.is_empty() {
350 let noun = if report.probed == 1 {
351 "vendor"
352 } else {
353 "vendors"
354 };
355 return format!("Nothing new detected ({} {noun} checked)", report.probed);
356 }
357 let names: Vec<&str> = report
358 .enabled
359 .iter()
360 .map(|vendor| vendor.display_name())
361 .collect();
362 format!("Enabled: {}\nWritten to {config_hint}", names.join(", "))
363}
364
365pub fn run_once_with(
369 config_path: Option<&Path>,
370 state_path: &Path,
371 force: bool,
372 probe: impl Fn(VendorId, &Config) -> bool,
373) -> Result<DetectReport> {
374 let resolved = match config_path {
375 Some(path) => Some(path.to_path_buf()),
376 None => crate::config::resolved_path(),
377 };
378 let config = match &resolved {
379 Some(path) => Config::load_from(path)?,
380 None => Config::default(),
381 };
382 let state = DetectState::load_at(state_path);
383 let plan = plan(&config, &state, VendorId::all(), force, |vendor| {
384 probe(vendor, &config)
385 });
386 let enabled = if plan.enable.is_empty() {
390 Vec::new()
391 } else {
392 let path = resolved.ok_or_else(|| {
393 AppError::Other("could not resolve the config.toml path to enable vendors in".into())
394 })?;
395 crate::config::enable_vendors_in(&path, &plan.enable)?
396 };
397 DetectState {
398 known: plan.known.clone(),
399 }
400 .save_at(state_path)?;
401 Ok(DetectReport {
402 enabled,
403 known: plan.known,
404 probed: plan.probed,
405 })
406}
407
408#[cfg(test)]
409mod tests {
410 use super::*;
411 use tempfile::TempDir;
412
413 fn probe_in(present: &[VendorId]) -> impl Fn(VendorId) -> bool + '_ {
414 move |vendor| present.contains(&vendor)
415 }
416
417 #[test]
418 fn state_round_trips_through_json_by_slug() {
419 let dir = TempDir::new().unwrap();
420 let path = dir.path().join("nested").join("detect.json");
421 let state = DetectState {
422 known: vec![
423 VendorId::Cursor,
424 VendorId::OpenCodeGo,
425 VendorId::NousResearch,
426 ],
427 };
428
429 state.save_at(&path).unwrap();
430
431 let text = std::fs::read_to_string(&path).unwrap();
432 assert!(text.contains("\"opencode-go\""), "{text}");
433 assert!(text.contains("\"nous\""), "{text}");
434 assert_eq!(DetectState::load_at(&path), state);
435 }
436
437 #[test]
438 fn missing_or_corrupt_state_is_the_default() {
439 let dir = TempDir::new().unwrap();
440 assert_eq!(
441 DetectState::load_at(&dir.path().join("absent.json")),
442 DetectState::default()
443 );
444
445 let corrupt = dir.path().join("corrupt.json");
446 std::fs::write(&corrupt, "{\"known\": [\"not-a-vendor\"").unwrap();
447 assert_eq!(DetectState::load_at(&corrupt), DetectState::default());
448
449 let unknown_slug = dir.path().join("unknown.json");
450 std::fs::write(&unknown_slug, "{\"known\": [\"not-a-vendor\"]}").unwrap();
451 assert_eq!(DetectState::load_at(&unknown_slug), DetectState::default());
452 }
453
454 #[test]
455 fn plan_enables_only_unknown_probed_vendors_that_are_off() {
456 let config = Config::default(); let state = DetectState {
458 known: vec![VendorId::Grok],
459 };
460 let all = [
461 VendorId::Anthropic, VendorId::Grok, VendorId::Cursor, VendorId::Kiro, ];
466 let present = [VendorId::Anthropic, VendorId::Grok, VendorId::Cursor];
467
468 let plan = plan(&config, &state, &all, false, probe_in(&present));
469
470 assert_eq!(plan.enable, vec![VendorId::Cursor]);
471 }
472
473 #[test]
474 fn force_reconsiders_known_vendors_but_never_enabled_ones() {
475 let config = Config::default();
476 let state = DetectState {
477 known: vec![VendorId::Grok, VendorId::Zai],
478 };
479 let all = [VendorId::Zai, VendorId::Grok];
480 let present = [VendorId::Zai, VendorId::Grok];
481
482 let plan = plan(&config, &state, &all, true, probe_in(&present));
483
484 assert_eq!(plan.enable, vec![VendorId::Grok]);
485 }
486
487 #[test]
488 fn plan_orders_enable_by_the_candidate_list() {
489 let config = Config::default();
490 let all = [VendorId::Kiro, VendorId::Cursor, VendorId::Grok];
491 let present = [VendorId::Grok, VendorId::Cursor, VendorId::Kiro];
492
493 let plan = plan(
494 &config,
495 &DetectState::default(),
496 &all,
497 false,
498 probe_in(&present),
499 );
500
501 assert_eq!(
502 plan.enable,
503 vec![VendorId::Kiro, VendorId::Cursor, VendorId::Grok]
504 );
505 }
506
507 #[test]
508 fn known_becomes_the_union_in_canonical_order_without_duplicates() {
509 let config = Config::default();
510 let state = DetectState {
511 known: vec![VendorId::Grok, VendorId::Cursor],
512 };
513 let all = [VendorId::Cursor, VendorId::Anthropic, VendorId::Cursor];
514
515 let plan = plan(&config, &state, &all, false, |_| false);
516
517 assert_eq!(
518 plan.known,
519 vec![VendorId::Anthropic, VendorId::Grok, VendorId::Cursor]
520 );
521 assert!(plan.enable.is_empty());
522 }
523
524 #[test]
525 fn probe_is_not_consulted_for_skipped_vendors() {
526 let config = Config::default();
527 let state = DetectState {
528 known: vec![VendorId::Grok],
529 };
530 let all = [VendorId::Grok, VendorId::Anthropic];
531
532 let plan = plan(&config, &state, &all, false, |vendor| {
533 panic!("probe called for {}", vendor.slug())
534 });
535
536 assert!(plan.enable.is_empty());
537 }
538
539 #[test]
540 fn copilot_hosts_file_must_be_a_non_empty_regular_file() {
541 let dir = TempDir::new().unwrap();
542 let hosts = dir.path().join("hosts.yml");
543 assert!(!copilot_hosts_present_at(&hosts));
544
545 std::fs::write(&hosts, "").unwrap();
546 assert!(!copilot_hosts_present_at(&hosts));
547
548 std::fs::write(&hosts, "github.com:\n user: octocat\n").unwrap();
549 assert!(copilot_hosts_present_at(&hosts));
550
551 assert!(!copilot_hosts_present_at(dir.path()));
552 }
553
554 #[test]
558 fn run_once_writes_enables_into_the_config_and_marks_everything_known() {
559 let dir = TempDir::new().unwrap();
560 let config_path = dir.path().join("config.toml");
561 let state_path = dir.path().join("detect.json");
562 std::fs::write(
563 &config_path,
564 "# mine
565[zai]
566enabled = false
567",
568 )
569 .unwrap();
570 let present = [VendorId::Zai, VendorId::Cursor, VendorId::Anthropic];
571 let probe = |vendor: VendorId, _: &Config| present.contains(&vendor);
572
573 let report = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
574
575 assert_eq!(report.enabled, vec![VendorId::Cursor]);
579 assert_eq!(report.known, VendorId::all());
580 assert_eq!(report.probed, VendorId::all().len());
581 let after = Config::load_from(&config_path).unwrap();
582 assert!(!after.is_enabled(VendorId::Zai), "an opt-out must survive");
583 assert!(after.is_enabled(VendorId::Cursor));
584 let text = std::fs::read_to_string(&config_path).unwrap();
585 assert!(
586 text.starts_with(
587 "# mine
588"
589 ),
590 "{text}"
591 );
592 assert_eq!(DetectState::load_at(&state_path).known, VendorId::all());
593
594 let again = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
595 assert!(again.enabled.is_empty(), "{again:?}");
596 assert_eq!(again.probed, 0, "everything is known: nothing to check");
597
598 std::fs::write(
600 &config_path,
601 "[cursor]
602enabled = false
603",
604 )
605 .unwrap();
606 let third = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
607 assert!(third.enabled.is_empty(), "{third:?}");
608 assert!(
609 !Config::load_from(&config_path)
610 .unwrap()
611 .is_enabled(VendorId::Cursor)
612 );
613
614 let forced = run_once_with(Some(&config_path), &state_path, true, probe).unwrap();
619 assert!(forced.enabled.is_empty(), "{forced:?}");
620 assert_eq!(forced.probed, VendorId::all().len());
621 assert!(
622 !Config::load_from(&config_path)
623 .unwrap()
624 .is_enabled(VendorId::Cursor)
625 );
626 }
627
628 #[test]
629 fn plan_counts_candidates_not_enables() {
630 let config = Config::default();
631 let state = DetectState {
632 known: vec![VendorId::Grok],
633 };
634 let all = [VendorId::Anthropic, VendorId::Grok, VendorId::Cursor];
635
636 let unforced = plan(&config, &state, &all, false, |_| false);
637 assert_eq!(unforced.probed, 2, "Grok is known and skipped");
638
639 let forced = plan(&config, &state, &all, true, |_| false);
640 assert_eq!(forced.probed, 3);
641 }
642
643 #[test]
644 fn format_report_lists_display_names_and_where_they_were_written() {
645 let report = DetectReport {
646 enabled: vec![VendorId::Cursor, VendorId::Kiro],
647 known: VendorId::all().to_vec(),
648 probed: 3,
649 };
650
651 let text = format_report(&report, "/home/u/.config/ai-usagebar/config.toml");
652
653 assert_eq!(
654 text,
655 "Enabled: Cursor, Kiro\nWritten to /home/u/.config/ai-usagebar/config.toml"
656 );
657 }
658
659 #[test]
660 fn format_report_says_how_many_were_checked_when_nothing_changed() {
661 let none = DetectReport {
662 enabled: vec![],
663 known: VendorId::all().to_vec(),
664 probed: 3,
665 };
666 assert_eq!(
667 format_report(&none, "unused"),
668 "Nothing new detected (3 vendors checked)"
669 );
670
671 let one = DetectReport {
672 probed: 1,
673 ..none.clone()
674 };
675 assert_eq!(
676 format_report(&one, "unused"),
677 "Nothing new detected (1 vendor checked)"
678 );
679 }
680
681 #[test]
683 fn report_serializes_slugs_and_the_probed_count() {
684 let dir = TempDir::new().unwrap();
685 let config_path = dir.path().join("config.toml");
686 let state_path = dir.path().join("detect.json");
687 std::fs::write(&config_path, "").unwrap();
688 DetectState {
689 known: vec![VendorId::Anthropic, VendorId::Grok],
690 }
691 .save_at(&state_path)
692 .unwrap();
693 let probe = |vendor: VendorId, _: &Config| vendor == VendorId::Cursor;
694
695 let report = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
696 let json: serde_json::Value = serde_json::to_value(&report).unwrap();
697
698 assert_eq!(json["enabled"], serde_json::json!(["cursor"]));
699 assert_eq!(
700 json["probed"],
701 serde_json::json!(VendorId::all().len() - 2),
702 "the two known vendors were not candidates"
703 );
704 let known = json["known"].as_array().unwrap();
705 assert_eq!(known.len(), VendorId::all().len());
706 assert_eq!(known[0], serde_json::json!("anthropic"));
707 assert_eq!(json.as_object().unwrap().len(), 3, "{json}");
708 }
709
710 #[test]
711 fn a_panicking_probe_counts_as_absent_and_still_saves_state() {
712 let dir = TempDir::new().unwrap();
713 let config_path = dir.path().join("config.toml");
714 let state_path = dir.path().join("detect.json");
715 let probe = |vendor: VendorId, _: &Config| {
716 catch_unwind(AssertUnwindSafe(|| {
717 if vendor == VendorId::Kiro {
718 panic!("boom");
719 }
720 vendor == VendorId::Grok
721 }))
722 .unwrap_or(false)
723 };
724
725 let enabled = run_once_with(Some(&config_path), &state_path, false, probe).unwrap();
726
727 assert_eq!(enabled.enabled, vec![VendorId::Grok]);
728 assert_eq!(DetectState::load_at(&state_path).known, VendorId::all());
729 }
730}