refactor: improve code quanity (#238)

5 files changed, 39 insertions, 53 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 0fb62f9..1844466 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -42,7 +42,6 @@ dependencies = [
  "clap",
  "crossbeam",
  "crossterm 0.26.1",
- "ctrlc",
  "dirs",
  "fancy-regex",
  "futures-util",
@@ -496,16 +495,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "ctrlc"
-version = "3.4.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "82e95fbd621905b854affdc67943b043a0fbb6ed7385fd5a25650d19a8a6cfdf"
-dependencies = [
- "nix 0.27.1",
- "windows-sys 0.48.0",
-]
-
-[[package]]
 name = "deranged"
 version = "0.3.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1089,17 +1078,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "nix"
-version = "0.27.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2eb04e9c688eff1c89d72b407f168cf79bb9e867a9d3323ed6c01519eb9cc053"
-dependencies = [
- "bitflags 2.4.1",
- "cfg-if",
- "libc",
-]
-
-[[package]]
 name = "nom"
 version = "7.1.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -2454,7 +2432,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "592b4883219f345e712b3209c62654ebda0bb50887f330cbd018d0f654bfd507"
 dependencies = [
  "gethostname",
- "nix 0.24.3",
+ "nix",
  "winapi",
  "winapi-wsapoll",
  "x11rb-protocol",
@@ -2466,5 +2444,5 @@ version = "0.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "56b245751c0ac9db0e006dc812031482784e434630205a93c73cfefcaabeac67"
 dependencies = [
- "nix 0.24.3",
+ "nix",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 569d9e1..e432353 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -27,7 +27,6 @@ crossbeam = "0.8.2"
 crossterm = "0.26.1"
 chrono = "0.4.23"
 bincode = "1.3.3"
-ctrlc = "3.2.5"
 parking_lot = "0.12.1"
 lazy_static = "1.4.0"
 fancy-regex = "0.11.0"
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') {