use std::{ cell::RefCell, io::{self, IsTerminal, Write}, sync::{Arc, Condvar, Mutex}, thread::{self, JoinHandle}, time::Duration, }; use ccusage_terminal::{terminal_width, truncate_to_width}; const SPINNER_FRAMES: [&str; 10] = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]; const SPINNER_INTERVAL: Duration = Duration::from_millis(80); /// Display columns the spinner frame and the space after it take up. const SPINNER_PREFIX_WIDTH: usize = 2; thread_local! { static ACTIVE_PROGRESS: RefCell> = const { RefCell::new(None) }; } /// Display label for one agent's load step. /// /// A newtype rather than an enum so the roster lives with the adapters: adding /// an agent does not touch this crate, which every adapter depends on. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub struct UsageLoadAgent(pub &'static str); #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum LoadProgressState { Loading, Succeeded, Failed, } pub fn should_show_usage_load_progress(json: bool, output_is_tty: bool) -> bool { !json && output_is_tty } fn agent_label(agent: UsageLoadAgent) -> &'static str { agent.0 } fn format_usage_load_progress_text( states: &[(UsageLoadAgent, LoadProgressState)], status: Option<&str>, ) -> String { if states.is_empty() { return status.unwrap_or("Loading usage logs").to_string(); } let base = { let completed = states .iter() .filter(|(_, state)| !matches!(state, LoadProgressState::Loading)) .count(); let loading_agents = states .iter() .filter_map(|(agent, state)| { matches!(state, LoadProgressState::Loading).then_some(agent_label(*agent)) }) .collect::>() .join(", "); if loading_agents.is_empty() { format!("Loading usage logs ({}/{})", completed, states.len()) } else { format!( "Loading usage logs ({}/{}) :: {}", completed, states.len(), loading_agents ) } }; match status { Some(status) => format!("{status} :: {base}"), None => base, } } /// Shorten the status text so the spinner line always occupies a single row. /// /// Each frame clears its own line with `\r\x1b[K`, which only erases the row the /// cursor sits on. A status wide enough to wrap therefore leaves the overflow /// row on screen and the redraw reads as flicker. One column is left free on top /// of the spinner prefix so the cursor never ends a frame in the wrap-pending /// state some terminals enter at the last column. fn fit_status_to_width(text: &str, terminal_width: usize) -> String { truncate_to_width( text, terminal_width.saturating_sub(SPINNER_PREFIX_WIDTH + 1), ) } pub fn usage_load_output_is_tty() -> bool { io::stdout().is_terminal() } pub struct UsageLoadProgress { enabled: bool, controller: Option, owns_session: bool, running: Option, Condvar)>>, worker: Option>, stopped: bool, } #[derive(Clone)] struct ProgressController { state: Arc>, } #[derive(Default)] struct ProgressState { status: Option, states: Vec<(UsageLoadAgent, LoadProgressState)>, frame: usize, rendered: bool, } impl UsageLoadProgress { pub fn new(enabled: bool) -> Self { if !enabled { return Self { enabled, controller: None, owns_session: false, running: None, worker: None, stopped: false, }; } if let Some(controller) = ACTIVE_PROGRESS.with(|active| active.borrow().clone()) { return Self { enabled, controller: Some(controller), owns_session: false, running: None, worker: None, stopped: false, }; } let state = Arc::new(Mutex::new(ProgressState::default())); let controller = ProgressController { state: Arc::clone(&state), }; let running = Arc::new((Mutex::new(true), Condvar::new())); let worker = { let running = Arc::clone(&running); Some(thread::spawn(move || { loop { if let Ok(mut state) = state.lock() && state.has_content() { state.render(); } let (lock, cvar) = &*running; let Ok(guard) = lock.lock() else { break; }; if !*guard { break; } let Ok((guard, _)) = cvar.wait_timeout(guard, SPINNER_INTERVAL) else { break; }; if !*guard { break; } } })) }; ACTIVE_PROGRESS.with(|active| { *active.borrow_mut() = Some(controller.clone()); }); Self { enabled, controller: Some(controller), owns_session: true, running: Some(running), worker, stopped: false, } } pub fn start(&mut self, agent: UsageLoadAgent) { self.set_state(agent, LoadProgressState::Loading); } pub fn succeed(&mut self, agent: UsageLoadAgent) { self.set_state(agent, LoadProgressState::Succeeded); } pub fn fail(&mut self, agent: UsageLoadAgent) { self.set_state(agent, LoadProgressState::Failed); } fn stop(&mut self) { if !self.enabled || self.stopped { return; } self.stopped = true; if !self.owns_session { return; } if let Some(running) = self.running.take() { let (lock, cvar) = &*running; if let Ok(mut is_running) = lock.lock() { *is_running = false; cvar.notify_all(); } } if let Some(worker) = self.worker.take() { let _ = worker.join(); } if let Some(controller) = self.controller.as_ref() && let Ok(mut state) = controller.state.lock() { if state.rendered { let _ = write!(io::stderr(), "\r\x1b[K\x1b[?25h"); let _ = io::stderr().flush(); } state.status = None; state.states.clear(); } ACTIVE_PROGRESS.with(|active| { *active.borrow_mut() = None; }); } fn set_status(&mut self, status: Option) { let Some(controller) = self.controller.as_ref() else { return; }; let Ok(mut state) = controller.state.lock() else { return; }; state.status = status; } fn set_state(&mut self, agent: UsageLoadAgent, state: LoadProgressState) { let Some(controller) = self.controller.as_ref() else { return; }; let Ok(mut progress_state) = controller.state.lock() else { return; }; if let Some((_, current)) = progress_state .states .iter_mut() .find(|(current_agent, _)| *current_agent == agent) { *current = state; } else { progress_state.states.push((agent, state)); } } } impl Drop for UsageLoadProgress { fn drop(&mut self) { self.stop(); } } impl ProgressState { fn has_content(&self) -> bool { self.status.is_some() || !self.states.is_empty() } fn render(&mut self) { if !self.has_content() { return; } let text = fit_status_to_width( &format_usage_load_progress_text(&self.states, self.status.as_deref()), terminal_width(), ); let frame = SPINNER_FRAMES[self.frame % SPINNER_FRAMES.len()]; self.frame = self.frame.wrapping_add(1); let _ = write!( io::stderr(), "\r\x1b[K\x1b[?25l\x1b[36m{frame}\x1b[39m {text}" ); let _ = io::stderr().flush(); self.rendered = true; } } pub fn track_usage_load( agent: UsageLoadAgent, json: bool, load: impl FnOnce() -> std::result::Result, ) -> std::result::Result { let enabled = crate::log_level() != Some(0) && should_show_usage_load_progress(json, usage_load_output_is_tty()); let mut progress = UsageLoadProgress::new(enabled); progress.start(agent); let result = load(); match &result { Ok(_) => progress.succeed(agent), Err(_) => progress.fail(agent), } progress.stop(); result } pub(crate) fn track_status( enabled: bool, status: impl Into, run: impl FnOnce() -> T, ) -> T { let mut progress = UsageLoadProgress::new(enabled); progress.set_status(Some(status.into())); let result = run(); progress.set_status(None); progress.stop(); result } #[cfg(test)] mod tests { use super::*; #[test] fn renders_active_agent_progress_with_completed_count() { let states = [ (UsageLoadAgent("Claude"), LoadProgressState::Succeeded), (UsageLoadAgent("Codex"), LoadProgressState::Loading), (UsageLoadAgent("OpenCode"), LoadProgressState::Loading), ]; assert_eq!( format_usage_load_progress_text(&states, None), "Loading usage logs (1/3) :: Codex, OpenCode" ); } #[test] fn includes_pricing_status_in_progress_text() { let states = [ (UsageLoadAgent("Claude"), LoadProgressState::Loading), (UsageLoadAgent("Codex"), LoadProgressState::Loading), ]; assert_eq!( format_usage_load_progress_text( &states, Some("Refreshing model pricing from LiteLLM...") ), "Refreshing model pricing from LiteLLM... :: Loading usage logs (0/2) :: Claude, Codex" ); } #[test] fn renders_standalone_status_without_usage_suffix() { assert_eq!( format_usage_load_progress_text(&[], Some("Refreshing model pricing from LiteLLM...")), "Refreshing model pricing from LiteLLM..." ); } #[test] fn fits_the_status_within_a_narrow_terminal() { let text = format_usage_load_progress_text( &[ (UsageLoadAgent("Claude"), LoadProgressState::Loading), (UsageLoadAgent("Codex"), LoadProgressState::Loading), ], Some("Refreshing model pricing from LiteLLM..."), ); let fitted = fit_status_to_width(&text, 80); assert_eq!( fitted, "Refreshing model pricing from LiteLLM... :: Loading usage logs (0/2) :: Clau…" ); assert_eq!(fitted.chars().count() + SPINNER_PREFIX_WIDTH, 79); } #[test] fn keeps_the_status_intact_on_a_wide_terminal() { let text = format_usage_load_progress_text(&[], Some("Refreshing model pricing")); assert_eq!(fit_status_to_width(&text, 120), "Refreshing model pricing"); } #[test] fn drops_the_status_when_the_terminal_has_no_room_at_all() { assert_eq!(fit_status_to_width("Loading usage logs", 3), ""); assert_eq!(fit_status_to_width("Loading usage logs", 4), "…"); } #[test] fn hides_progress_for_json_or_non_tty_output() { assert!(!should_show_usage_load_progress(true, true)); assert!(!should_show_usage_load_progress(false, false)); assert!(should_show_usage_load_progress(false, true)); } }