1use crate::agent::ui::colors::ansi;
9use parking_lot::RwLock;
10use std::io::{self, Write};
11use std::sync::Arc;
12use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
13use std::time::{Duration, Instant};
14use tokio::sync::mpsc;
15
16const INDICATOR_FRAMES: &[&str] = &["✱", "✳", "✱", "✴", "✱", "✳"];
18
19const ANIMATION_INTERVAL_MS: u64 = 300;
21
22#[derive(Debug, Clone)]
24pub enum ProgressMessage {
25 UpdateTokens { input: u64, output: u64 },
27 Action(String),
29 Focus(String),
31 ClearFocus,
33 Stop,
35}
36
37#[derive(Debug)]
39pub struct ProgressState {
40 pub input_tokens: AtomicU64,
41 pub output_tokens: AtomicU64,
42 pub is_running: AtomicBool,
43 pub is_paused: AtomicBool,
45 pub interrupt_requested: AtomicBool,
47 pub action: RwLock<String>,
49 pub focus: RwLock<Option<String>>,
51 pub start_time: std::time::Instant,
53 pub layout_state: RwLock<Option<std::sync::Arc<super::layout::LayoutState>>>,
55}
56
57impl Default for ProgressState {
58 fn default() -> Self {
59 Self {
60 input_tokens: AtomicU64::new(0),
61 output_tokens: AtomicU64::new(0),
62 is_running: AtomicBool::new(true),
63 is_paused: AtomicBool::new(false),
64 interrupt_requested: AtomicBool::new(false),
65 action: RwLock::new("Generating".to_string()),
66 focus: RwLock::new(None),
67 start_time: std::time::Instant::now(),
68 layout_state: RwLock::new(None),
69 }
70 }
71}
72
73impl ProgressState {
74 pub fn new() -> Arc<Self> {
75 Arc::new(Self::default())
76 }
77
78 pub fn update_tokens(&self, input: u64, output: u64) {
79 self.input_tokens.fetch_add(input, Ordering::SeqCst);
80 self.output_tokens.fetch_add(output, Ordering::SeqCst);
81 }
82
83 pub fn get_tokens(&self) -> (u64, u64) {
84 (
85 self.input_tokens.load(Ordering::SeqCst),
86 self.output_tokens.load(Ordering::SeqCst),
87 )
88 }
89
90 pub fn set_action(&self, action: &str) {
91 *self.action.write() = action.to_string();
92 }
93
94 pub fn get_action(&self) -> String {
95 self.action.read().clone()
96 }
97
98 pub fn set_focus(&self, focus: &str) {
99 *self.focus.write() = Some(focus.to_string());
100 }
101
102 pub fn clear_focus(&self) {
103 *self.focus.write() = None;
104 }
105
106 pub fn get_focus(&self) -> Option<String> {
107 self.focus.read().clone()
108 }
109
110 pub fn stop(&self) {
111 self.is_running.store(false, Ordering::SeqCst);
112 }
113
114 pub fn is_running(&self) -> bool {
115 self.is_running.load(Ordering::SeqCst)
116 }
117
118 pub fn pause(&self) {
120 self.is_paused.store(true, Ordering::SeqCst);
121 }
122
123 pub fn resume(&self) {
125 self.is_paused.store(false, Ordering::SeqCst);
126 }
127
128 pub fn is_paused(&self) -> bool {
129 self.is_paused.load(Ordering::SeqCst)
130 }
131
132 pub fn elapsed(&self) -> std::time::Duration {
134 self.start_time.elapsed()
135 }
136
137 pub fn set_layout(&self, layout: std::sync::Arc<super::layout::LayoutState>) {
139 *self.layout_state.write() = Some(layout);
140 }
141
142 pub fn has_layout(&self) -> bool {
144 self.layout_state
145 .read()
146 .as_ref()
147 .map(|l| l.is_active())
148 .unwrap_or(false)
149 }
150
151 pub fn get_layout(&self) -> Option<std::sync::Arc<super::layout::LayoutState>> {
153 self.layout_state.read().clone()
154 }
155
156 pub fn request_interrupt(&self) {
158 self.interrupt_requested.store(true, Ordering::SeqCst);
159 }
160
161 pub fn is_interrupted(&self) -> bool {
163 self.interrupt_requested.load(Ordering::SeqCst)
164 }
165
166 pub fn clear_interrupt(&self) {
168 self.interrupt_requested.store(false, Ordering::SeqCst);
169 }
170}
171
172pub struct GenerationIndicator {
174 sender: mpsc::Sender<ProgressMessage>,
175 state: Arc<ProgressState>,
176}
177
178impl GenerationIndicator {
179 pub fn new() -> Self {
181 Self::with_action("Generating")
182 }
183
184 pub fn with_action(action: &str) -> Self {
186 let (sender, receiver) = mpsc::channel(32);
187 let state = ProgressState::new();
188 state.set_action(action);
189 let state_clone = state.clone();
190
191 tokio::spawn(async move {
192 run_progress_indicator(receiver, state_clone).await;
193 });
194
195 Self { sender, state }
196 }
197
198 pub async fn update_tokens(&self, input: u64, output: u64) {
200 self.state.update_tokens(input, output);
201 let _ = self
202 .sender
203 .send(ProgressMessage::UpdateTokens { input, output })
204 .await;
205 }
206
207 pub async fn set_action(&self, action: &str) {
209 self.state.set_action(action);
210 let _ = self
211 .sender
212 .send(ProgressMessage::Action(action.to_string()))
213 .await;
214 }
215
216 pub async fn set_focus(&self, focus: &str) {
218 self.state.set_focus(focus);
219 let _ = self
220 .sender
221 .send(ProgressMessage::Focus(focus.to_string()))
222 .await;
223 }
224
225 pub async fn clear_focus(&self) {
227 self.state.clear_focus();
228 let _ = self.sender.send(ProgressMessage::ClearFocus).await;
229 }
230
231 pub async fn stop(&self) {
233 self.state.stop();
234 let _ = self.sender.send(ProgressMessage::Stop).await;
235 tokio::time::sleep(Duration::from_millis(50)).await;
237 }
238
239 pub async fn pause(&self) {
241 self.state.pause();
242 print!("\r{}", ansi::CLEAR_LINE);
244 print!("{}", ansi::SHOW_CURSOR);
245 let _ = io::stdout().flush();
246 }
247
248 pub async fn resume(&self) {
250 self.state.resume();
251 print!("{}", ansi::HIDE_CURSOR);
252 let _ = io::stdout().flush();
253 }
254
255 pub fn state(&self) -> Arc<ProgressState> {
257 self.state.clone()
258 }
259}
260
261impl Default for GenerationIndicator {
262 fn default() -> Self {
263 Self::new()
264 }
265}
266
267fn format_tokens(tokens: u64) -> String {
269 if tokens >= 100_000 {
270 format!("{:.1}k", tokens as f64 / 1000.0)
271 } else if tokens >= 10_000 {
272 format!("{:.0}k", tokens as f64 / 1000.0)
273 } else {
274 tokens.to_string()
275 }
276}
277
278const CORAL: &str = "\x1b[38;5;209m";
280
281async fn run_progress_indicator(
286 mut receiver: mpsc::Receiver<ProgressMessage>,
287 state: Arc<ProgressState>,
288) {
289 let start_time = Instant::now();
290 let mut frame_index = 0;
291 let mut had_focus = false;
292 let mut interval = tokio::time::interval(Duration::from_millis(ANIMATION_INTERVAL_MS));
293
294 if !state.has_layout() {
296 print!("{}", ansi::HIDE_CURSOR);
297 let _ = io::stdout().flush();
298 }
299
300 let mut was_rendering = false;
302
303 loop {
304 tokio::select! {
305 _ = interval.tick() => {
306 if !state.is_running() {
307 break;
308 }
309
310 let use_layout = state.has_layout();
311
312 if state.is_paused() {
314 if was_rendering && !use_layout {
315 if had_focus {
317 print!("{}{}", ansi::CURSOR_UP, ansi::CLEAR_LINE);
318 }
319 print!("\r{}", ansi::CLEAR_LINE);
320 print!("{}", ansi::SHOW_CURSOR);
321 let _ = io::stdout().flush();
322 was_rendering = false;
323 had_focus = false;
324 }
325 continue;
326 }
327
328 if !was_rendering && !use_layout {
330 print!("{}", ansi::HIDE_CURSOR);
331 let _ = io::stdout().flush();
332 }
333 was_rendering = true;
334
335 let elapsed = start_time.elapsed();
336 let indicator = INDICATOR_FRAMES[frame_index % INDICATOR_FRAMES.len()];
337 frame_index += 1;
338
339 let action = state.get_action();
340 let focus = state.get_focus();
341 let (input_tokens, output_tokens) = state.get_tokens();
342 let total_tokens = input_tokens + output_tokens;
343
344 let elapsed_secs = elapsed.as_secs_f64();
346 let elapsed_str = if elapsed_secs >= 60.0 {
347 format!("{:.0}m {:.0}s", elapsed_secs / 60.0, elapsed_secs % 60.0)
348 } else {
349 format!("{:.1}s", elapsed_secs)
350 };
351
352 let stats = if total_tokens > 0 {
353 format!(
354 "{}(^C to stop · {} · ↓ {} tokens){}",
355 ansi::DIM,
356 elapsed_str,
357 format_tokens(total_tokens),
358 ansi::RESET
359 )
360 } else {
361 format!(
362 "{}(^C to stop · {}){}",
363 ansi::DIM,
364 elapsed_str,
365 ansi::RESET
366 )
367 };
368
369 let status_content = format!(
371 "{}{}{} {}{}…{} {}",
372 CORAL,
373 indicator,
374 ansi::RESET,
375 CORAL,
376 action,
377 ansi::RESET,
378 stats,
379 );
380
381 if use_layout {
383 if let Some(layout_state) = state.get_layout() {
384 render_to_layout(&layout_state, &status_content, focus.as_deref());
386 }
387 } else {
388 if had_focus {
391 print!("{}{}", ansi::CURSOR_UP, ansi::CLEAR_LINE);
392 }
393 print!("\r{}", ansi::CLEAR_LINE);
394
395 print!("{}", status_content);
397
398 if let Some(ref focus_text) = focus {
400 print!(
401 "\n{}└{} {}{}{}",
402 ansi::DIM,
403 ansi::RESET,
404 ansi::GRAY,
405 focus_text,
406 ansi::RESET
407 );
408 had_focus = true;
409 } else {
410 had_focus = false;
411 }
412
413 let _ = io::stdout().flush();
414 }
415 }
416 Some(msg) = receiver.recv() => {
417 match msg {
418 ProgressMessage::UpdateTokens { .. } => {
419 }
421 ProgressMessage::Action(action) => {
422 state.set_action(&action);
423 }
424 ProgressMessage::Focus(focus) => {
425 state.set_focus(&focus);
426 }
427 ProgressMessage::ClearFocus => {
428 state.clear_focus();
429 }
430 ProgressMessage::Stop => {
431 state.stop();
432 break;
433 }
434 }
435 }
436 }
437 }
438
439 if !state.has_layout() {
441 if had_focus {
442 print!("{}{}", ansi::CURSOR_UP, ansi::CLEAR_LINE);
443 }
444 print!("\r{}", ansi::CLEAR_LINE);
445 print!("{}", ansi::SHOW_CURSOR);
446 let _ = io::stdout().flush();
447 }
448}
449
450fn render_to_layout(layout_state: &super::layout::LayoutState, status: &str, focus: Option<&str>) {
452 use super::layout::escape;
453
454 if !layout_state.is_active() {
455 return;
456 }
457
458 let mut stdout = io::stdout();
459 let status_line = layout_state.status_line();
460 let focus_line = layout_state.focus_line();
461
462 print!("{}", escape::SAVE_CURSOR);
464 print!("{}", escape::move_to_line(status_line));
465 print!("{}", ansi::CLEAR_LINE);
466 print!("{}", status);
467
468 print!("{}", escape::move_to_line(focus_line));
470 print!("{}", ansi::CLEAR_LINE);
471 if let Some(focus_text) = focus {
472 print!(
473 "{}└{} {}{}{}",
474 ansi::DIM,
475 ansi::RESET,
476 ansi::GRAY,
477 focus_text,
478 ansi::RESET
479 );
480 }
481
482 print!("{}", escape::RESTORE_CURSOR);
483 let _ = stdout.flush();
484}
485
486#[cfg(test)]
487mod tests {
488 use super::*;
489
490 #[test]
491 fn test_format_tokens() {
492 assert_eq!(format_tokens(0), "0");
493 assert_eq!(format_tokens(999), "999");
494 assert_eq!(format_tokens(1000), "1000");
495 assert_eq!(format_tokens(9999), "9999");
496 assert_eq!(format_tokens(10000), "10k");
497 assert_eq!(format_tokens(10499), "10k");
498 assert_eq!(format_tokens(10999), "11k");
499 assert_eq!(format_tokens(100000), "100.0k");
500 assert_eq!(format_tokens(150000), "150.0k");
501 }
502
503 #[test]
504 fn test_progress_state() {
505 let state = ProgressState::new();
506 assert!(state.is_running());
507 assert_eq!(state.get_tokens(), (0, 0));
508 assert_eq!(state.get_action(), "Generating");
509 assert!(state.get_focus().is_none());
510
511 state.update_tokens(100, 50);
512 assert_eq!(state.get_tokens(), (100, 50));
513
514 state.set_action("Analyzing");
515 assert_eq!(state.get_action(), "Analyzing");
516
517 state.set_focus("Reading file.rs");
518 assert_eq!(state.get_focus(), Some("Reading file.rs".to_string()));
519
520 state.clear_focus();
521 assert!(state.get_focus().is_none());
522
523 state.stop();
524 assert!(!state.is_running());
525 }
526}