From 740f4060f1272063ec89f2b93ed334d3ed1cbfe4 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Tue, 29 Oct 2024 14:31:04 +0800
Subject: fix: unexpected Ctrl-C/Ctrl-D handling in non-stream REPL Chat (#957)

---
 src/client/common.rs | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

(limited to 'src/client/common.rs')

diff --git a/src/client/common.rs b/src/client/common.rs
index fe12593..1275d0b 100644
--- a/src/client/common.rs
+++ b/src/client/common.rs
@@ -403,9 +403,15 @@ pub async fn call_chat_completions(
     input: &Input,
     extract_code: bool,
     client: &dyn Client,
+    abort_signal: AbortSignal,
 ) -> Result<(String, Vec<ToolResult>)> {
-    let task = client.chat_completions(input.clone());
-    let ret = run_with_spinner(task, "Generating").await;
+    let ret = abortable_run_with_spinner(
+        client.chat_completions(input.clone()),
+        "Generating",
+        abort_signal,
+    )
+    .await;
+
     match ret {
         Ok(ret) => {
             let ChatCompletionsOutput {
@@ -438,6 +444,10 @@ pub async fn call_chat_completions_streaming(
         render_stream(rx, client.global_config(), abort_signal.clone()),
     );
 
+    if handler.abort().aborted() {
+        bail!("Aborted.");
+    }
+
     render_ret?;
 
     let (text, tool_calls) = handler.take();
-- 
cgit v1.2.3