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