use std::io::{IsTerminal, Write};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};
const FRAMES: [char; 10] = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];
const START_DELAY: Duration = Duration::from_millis(200);
const FRAME_INTERVAL: Duration = Duration::from_millis(90);
const POLL_INTERVAL: Duration = Duration::from_millis(10);
const COUNTER_THROTTLE: Duration = Duration::from_millis(120);
pub(crate) fn should_show_spinner(
no_color: bool,
quiet: bool,
term_dumb: bool,
stderr_is_tty: bool,
) -> bool {
!(no_color || quiet || term_dumb || !stderr_is_tty)
}
fn should_show_spinner_now(quiet: bool) -> bool {
should_show_spinner(
std::env::var_os("NO_COLOR").is_some(),
quiet,
std::env::var("TERM").ok().as_deref() == Some("dumb"),
std::io::stderr().is_terminal(),
)
}
pub(crate) fn env_quiet() -> bool {
std::env::var_os("SUPERCODE_QUIET").is_some()
}
pub struct Spinner {
enabled: bool,
running: AtomicBool,
stop_flag: Arc<AtomicBool>,
handle: Mutex<Option<JoinHandle<()>>>,
turn_bytes: AtomicU64,
session_bytes: AtomicU64,
counter_drawn: AtomicBool,
last_counter_render: Mutex<Instant>,
}
impl Spinner {
pub fn new(quiet: bool) -> Self {
Self::with_enabled(should_show_spinner_now(quiet))
}
pub(crate) fn with_enabled(enabled: bool) -> Self {
Self {
enabled,
running: AtomicBool::new(false),
stop_flag: Arc::new(AtomicBool::new(false)),
handle: Mutex::new(None),
turn_bytes: AtomicU64::new(0),
session_bytes: AtomicU64::new(0),
counter_drawn: AtomicBool::new(false),
last_counter_render: Mutex::new(
Instant::now()
.checked_sub(COUNTER_THROTTLE)
.unwrap_or_else(Instant::now),
),
}
}
pub fn start(&self, label: &str) {
if !self.enabled {
return;
}
if self.running.swap(true, Ordering::AcqRel) {
return;
}
self.stop_flag.store(false, Ordering::Release);
let stop_flag = self.stop_flag.clone();
let label = label.to_string();
let handle = thread::spawn(move || spin(&stop_flag, &label));
*self.handle.lock().unwrap() = Some(handle);
}
pub fn stop(&self) {
if !self.running.swap(false, Ordering::AcqRel) {
return;
}
self.stop_flag.store(true, Ordering::Release);
if let Some(handle) = self.handle.lock().unwrap().take() {
let _ = handle.join();
}
}
pub(crate) fn reset_turn(&self) {
if !self.enabled {
return;
}
self.turn_bytes.store(0, Ordering::Relaxed);
}
pub(crate) fn record_delta(&self, text: &str) {
if !self.enabled {
return;
}
let bytes = text.len() as u64;
self.turn_bytes.fetch_add(bytes, Ordering::Relaxed);
self.session_bytes.fetch_add(bytes, Ordering::Relaxed);
}
pub(crate) fn render_counter(&self) {
if !self.enabled {
return;
}
{
let mut last = self.last_counter_render.lock().unwrap();
if last.elapsed() < COUNTER_THROTTLE {
return;
}
*last = Instant::now();
}
let turn_tok = self.turn_bytes.load(Ordering::Relaxed).div_ceil(4);
let session_tok = self.session_bytes.load(Ordering::Relaxed).div_ceil(4);
let mut out = std::io::stderr();
let _ = write!(
out,
"\rthis turn: ~{turn_tok} tok · session: ~{session_tok} tok\x1b[K"
);
let _ = out.flush();
self.counter_drawn.store(true, Ordering::Release);
}
pub(crate) fn clear_counter(&self) {
if !self.counter_drawn.swap(false, Ordering::AcqRel) {
return;
}
let mut out = std::io::stderr();
let _ = write!(out, "\r\x1b[2K");
let _ = out.flush();
}
}
impl Drop for Spinner {
fn drop(&mut self) {
self.stop();
self.clear_counter();
}
}
fn spin(stop_flag: &AtomicBool, label: &str) {
let started = Instant::now();
while started.elapsed() < START_DELAY {
if stop_flag.load(Ordering::Acquire) {
return; }
thread::sleep(POLL_INTERVAL);
}
let mut out = std::io::stderr();
let mut frame = 0usize;
loop {
if stop_flag.load(Ordering::Acquire) {
break;
}
let elapsed = started.elapsed().as_secs_f32();
let _ = write!(
out,
"\r{} {label} ({elapsed:.1}s)",
FRAMES[frame % FRAMES.len()]
);
let _ = out.flush();
frame += 1;
let deadline = Instant::now() + FRAME_INTERVAL;
while Instant::now() < deadline {
if stop_flag.load(Ordering::Acquire) {
break;
}
thread::sleep(POLL_INTERVAL);
}
}
let _ = write!(out, "\r\x1b[2K");
let _ = out.flush();
}
pub fn with_spinner<T>(quiet: bool, f: impl FnOnce() -> T) -> T {
let spinner = Spinner::new(quiet);
spinner.start("Working…");
let out = f();
spinner.stop();
out
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn shows_on_a_plain_interactive_tty() {
assert!(should_show_spinner(false, false, false, true));
}
#[test]
fn suppressed_when_no_color_is_set() {
assert!(!should_show_spinner(true, false, false, true));
}
#[test]
fn suppressed_in_quiet_mode() {
assert!(!should_show_spinner(false, true, false, true));
}
#[test]
fn quiet_flag_wins_even_on_a_perfectly_normal_interactive_tty() {
assert!(should_show_spinner(false, false, false, true)); assert!(!should_show_spinner(
false, true, false,
true,
));
}
#[test]
fn suppressed_on_dumb_terminal() {
assert!(!should_show_spinner(false, false, true, true));
}
#[test]
fn suppressed_when_stderr_is_not_a_tty() {
assert!(!should_show_spinner(false, false, false, false));
}
#[test]
fn suppressed_when_stderr_is_not_a_tty_even_with_everything_else_favorable() {
assert!(!should_show_spinner(false, false, false, false));
}
#[test]
fn multiple_suppressors_still_suppress() {
assert!(!should_show_spinner(true, true, true, false));
}
#[test]
fn disabled_spinner_start_stop_is_a_no_op() {
let s = Spinner::with_enabled(false);
s.start("Thinking…");
s.stop();
}
#[test]
fn enabled_spinner_stopped_before_delay_elapses_draws_nothing() {
let s = Spinner::with_enabled(true);
s.start("Thinking…");
s.stop();
s.start("Thinking…");
s.stop();
}
#[test]
fn start_is_idempotent_while_running() {
let s = Spinner::with_enabled(true);
s.start("Thinking…");
s.start("Thinking…"); s.stop();
}
#[test]
fn record_delta_accumulates_turn_and_session_bytes() {
let s = Spinner::with_enabled(true);
s.record_delta("hello"); s.record_delta(" world"); assert_eq!(s.turn_bytes.load(Ordering::Relaxed), 11);
assert_eq!(s.session_bytes.load(Ordering::Relaxed), 11);
}
#[test]
fn record_delta_is_a_true_no_op_when_disabled() {
let s = Spinner::with_enabled(false);
s.record_delta("hello world");
assert_eq!(s.turn_bytes.load(Ordering::Relaxed), 0);
assert_eq!(s.session_bytes.load(Ordering::Relaxed), 0);
}
#[test]
fn reset_turn_zeroes_this_turn_but_not_session() {
let first = "first turn output";
let second = "second turn";
let s = Spinner::with_enabled(true);
s.record_delta(first);
assert_eq!(s.turn_bytes.load(Ordering::Relaxed), first.len() as u64);
s.reset_turn();
assert_eq!(s.turn_bytes.load(Ordering::Relaxed), 0);
assert_eq!(
s.session_bytes.load(Ordering::Relaxed),
first.len() as u64,
"session total must survive a turn reset"
);
s.record_delta(second);
assert_eq!(s.turn_bytes.load(Ordering::Relaxed), second.len() as u64);
assert_eq!(
s.session_bytes.load(Ordering::Relaxed),
(first.len() + second.len()) as u64
);
}
#[test]
fn reset_turn_is_a_no_op_when_disabled() {
let s = Spinner::with_enabled(false);
s.reset_turn(); assert_eq!(s.turn_bytes.load(Ordering::Relaxed), 0);
}
#[test]
fn render_counter_draws_when_enabled_and_due() {
let s = Spinner::with_enabled(true);
s.record_delta("some streamed text");
assert!(!s.counter_drawn.load(Ordering::Acquire));
s.render_counter();
assert!(
s.counter_drawn.load(Ordering::Acquire),
"first render after construction must draw immediately (backdated throttle)"
);
}
#[test]
fn render_counter_is_a_no_op_when_disabled() {
let s = Spinner::with_enabled(false);
s.record_delta("text");
s.render_counter();
assert!(!s.counter_drawn.load(Ordering::Acquire));
}
#[test]
fn render_counter_is_throttled_within_the_window() {
let s = Spinner::with_enabled(true);
s.record_delta("a");
s.render_counter();
assert!(s.counter_drawn.load(Ordering::Acquire));
s.clear_counter();
assert!(!s.counter_drawn.load(Ordering::Acquire));
s.record_delta("b");
s.render_counter(); assert!(
!s.counter_drawn.load(Ordering::Acquire),
"a redraw inside COUNTER_THROTTLE must be suppressed"
);
}
#[test]
fn clear_counter_is_idempotent_when_nothing_was_drawn() {
let s = Spinner::with_enabled(true);
s.clear_counter(); assert!(!s.counter_drawn.load(Ordering::Acquire));
s.clear_counter(); assert!(!s.counter_drawn.load(Ordering::Acquire));
}
#[test]
fn clear_counter_resets_the_drawn_flag() {
let s = Spinner::with_enabled(true);
s.record_delta("hi");
s.render_counter();
assert!(s.counter_drawn.load(Ordering::Acquire));
s.clear_counter();
assert!(!s.counter_drawn.load(Ordering::Acquire));
}
#[test]
fn token_math_matches_the_shared_estimator() {
for sample in ["", "a", "abcd", "abcde", "hello, world! this streams."] {
let bytes = sample.len() as u64;
let from_bytes = bytes.div_ceil(4);
let from_estimator = supercode::tokens::estimate_tokens(sample);
assert_eq!(
from_bytes, from_estimator,
"byte-count formula must match estimate_tokens for {sample:?}"
);
}
}
}