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