From c1c767a5bdc6a8f572ff643c3d0ea42fffca7865 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Mon, 13 Jan 2025 14:09:27 -0500
Subject: [PATCH] assistant2: Make `Esc` cancel current completion (#23076)

This PR makes it so pressing `Esc` in Assistant2 will cancel the current
completion.

Release Notes:

- N/A
---
 crates/assistant2/src/active_thread.rs   |  6 ++++++
 crates/assistant2/src/assistant_panel.rs |  6 ++++++
 crates/assistant2/src/thread.rs          | 11 +++++++++++
 3 files changed, 23 insertions(+)

diff --git a/crates/assistant2/src/active_thread.rs b/crates/assistant2/src/active_thread.rs
index e56d766ea1..7633d044a1 100644
--- a/crates/assistant2/src/active_thread.rs
+++ b/crates/assistant2/src/active_thread.rs
@@ -80,6 +80,12 @@ impl ActiveThread {
         self.thread.read(cx).summary_or_default()
     }
 
+    pub fn cancel_last_completion(&mut self, cx: &mut AppContext) -> bool {
+        self.last_error.take();
+        self.thread
+            .update(cx, |thread, _cx| thread.cancel_last_completion())
+    }
+
     pub fn last_error(&self) -> Option<ThreadError> {
         self.last_error.clone()
     }
diff --git a/crates/assistant2/src/assistant_panel.rs b/crates/assistant2/src/assistant_panel.rs
index 894064842a..4c6993312e 100644
--- a/crates/assistant2/src/assistant_panel.rs
+++ b/crates/assistant2/src/assistant_panel.rs
@@ -143,6 +143,11 @@ impl AssistantPanel {
         &self.thread_store
     }
 
+    fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
+        self.thread
+            .update(cx, |thread, cx| thread.cancel_last_completion(cx));
+    }
+
     fn new_thread(&mut self, cx: &mut ViewContext<Self>) {
         let thread = self
             .thread_store
@@ -611,6 +616,7 @@ impl Render for AssistantPanel {
             .key_context("AssistantPanel2")
             .justify_between()
             .size_full()
+            .on_action(cx.listener(Self::cancel))
             .on_action(cx.listener(|this, _: &NewThread, cx| {
                 this.new_thread(cx);
             }))
diff --git a/crates/assistant2/src/thread.rs b/crates/assistant2/src/thread.rs
index c3933cf459..707d8be514 100644
--- a/crates/assistant2/src/thread.rs
+++ b/crates/assistant2/src/thread.rs
@@ -502,6 +502,17 @@ impl Thread {
             };
         }
     }
+
+    /// Cancels the last pending completion, if there are any pending.
+    ///
+    /// Returns whether a completion was canceled.
+    pub fn cancel_last_completion(&mut self) -> bool {
+        if let Some(_last_completion) = self.pending_completions.pop() {
+            true
+        } else {
+            false
+        }
+    }
 }
 
 #[derive(Debug, Clone)]