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