refactor: improve code quanity (#238)

3 files changed, 37 insertions, 28 deletions

diff --git a/src/client/common.rs b/src/client/common.rs
index 27dd81d..cf5ba9b 100644
--- a/src/client/common.rs
+++ b/src/client/common.rs
@@ -219,7 +219,7 @@ pub trait Client {
                         let content = global_config.read().echo_messages(content);
                         let tokens = tokenize(&content);
                         for token in tokens {
-                            tokio::time::sleep(Duration::from_millis(25)).await;
+                            tokio::time::sleep(Duration::from_millis(10)).await;
                             handler.text(&token)?;
                         }
                         return Ok(());
@@ -235,10 +235,6 @@ pub trait Client {
                     handler.done()?;
                     Ok(())
                  },
-                _ =  tokio::signal::ctrl_c() => {
-                    abort.set_ctrlc();
-                    Ok(())
-                }
             }
         })
     }
diff --git a/src/main.rs b/src/main.rs
index dd9f9cf..bb6c8a0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -117,11 +117,6 @@ fn start_directive(config: &GlobalConfig, input: &str, no_stream: bool) -> Resul
         output
     } else {
         let abort = create_abort_signal();
-        let abort_clone = abort.clone();
-        ctrlc::set_handler(move || {
-            abort_clone.set_ctrlc();
-        })
-        .expect("Failed to setting Ctrl-C handler");
         render_stream(input, client.as_ref(), config, abort)?
     };
     config.write().save_message(input, &output)
diff --git a/src/render/stream.rs b/src/render/stream.rs
index 2e05c3a..ad35dec 100644
--- a/src/render/stream.rs
+++ b/src/render/stream.rs
@@ -12,6 +12,7 @@ use crossterm::{
 };
 use std::{
     io::{self, Stdout, Write},
+    ops::Div,
     time::{Duration, Instant},
 };
 use textwrap::core::display_width;
@@ -66,19 +67,17 @@ fn markdown_stream_inner(
 
     let mut spinner = Spinner::new(" Generating");
 
-    loop {
+    'outer: loop {
         if abort.aborted() {
             return Ok(());
         }
         spinner.step(writer)?;
 
-        if let Ok(evt) = rx.try_recv() {
-            match evt {
-                ReplyEvent::Text(text) => {
-                    if spinner.is_running() {
-                        spinner.stop(writer)?;
-                    }
+        for reply_event in gather_events(rx) {
+            spinner.stop(writer)?;
 
+            match reply_event {
+                ReplyEvent::Text(text) => {
                     let (col, mut row) = cursor::position()?;
 
                     // Fix unexpected duplicate lines on kitty, see https://github.com/sigoden/aichat/issues/105
@@ -130,25 +129,24 @@ fn markdown_stream_inner(
                     writer.flush()?;
                 }
                 ReplyEvent::Done => {
-                    break;
+                    break 'outer;
                 }
             }
-            continue;
         }
 
         let timeout = tick_rate
             .checked_sub(last_tick.elapsed())
-            .unwrap_or_else(|| Duration::from_secs(0));
+            .unwrap_or_else(|| tick_rate.div(2));
         if crossterm::event::poll(timeout)? {
             if let Event::Key(key) = event::read()? {
                 match key.code {
                     KeyCode::Char('c') if key.modifiers == KeyModifiers::CONTROL => {
                         abort.set_ctrlc();
-                        return Ok(());
+                        break;
                     }
                     KeyCode::Char('d') if key.modifiers == KeyModifiers::CONTROL => {
                         abort.set_ctrld();
-                        return Ok(());
+                        break;
                     }
                     _ => {}
                 }
@@ -159,6 +157,9 @@ fn markdown_stream_inner(
             last_tick = Instant::now();
         }
     }
+
+    spinner.stop(writer)?;
+
     Ok(())
 }
 
@@ -179,17 +180,13 @@ impl Spinner {
         }
     }
 
-    fn is_running(&self) -> bool {
-        !self.stopped
-    }
-
     fn step(&mut self, writer: &mut Stdout) -> Result<()> {
         if self.stopped {
             return Ok(());
         }
         let frame = Self::DATA[self.index % Self::DATA.len()];
-        let dots = ".".repeat((self.index / 8) % 4);
-        let line = format!("{frame}{}{dots}", self.message);
+        let dots = ".".repeat((self.index / 5) % 4);
+        let line = format!("{frame}{}{:<3}", self.message, dots);
         queue!(writer, cursor::MoveToColumn(0), style::Print(line),)?;
         if self.index == 0 {
             queue!(writer, cursor::Hide)?;
@@ -215,6 +212,27 @@ impl Spinner {
     }
 }
 
+fn gather_events(rx: &Receiver<ReplyEvent>) -> Vec<ReplyEvent> {
+    let mut texts = vec![];
+    let mut done = false;
+    for reply_event in rx.try_iter() {
+        match reply_event {
+            ReplyEvent::Text(v) => texts.push(v),
+            ReplyEvent::Done => {
+                done = true;
+            }
+        }
+    }
+    let mut events = vec![];
+    if !texts.is_empty() {
+        events.push(ReplyEvent::Text(texts.join("")))
+    }
+    if done {
+        events.push(ReplyEvent::Done)
+    }
+    events
+}
+
 fn print_block(writer: &mut Stdout, text: &str, columns: u16) -> Result<u16> {
     let mut num = 0;
     for line in text.split('\n') {