From 530224527d68809becc86191b2a85b89fc1110a2 Mon Sep 17 00:00:00 2001
From: Antonio Scandurra <me@as-cii.com>
Date: Thu, 25 Apr 2024 15:12:58 +0200
Subject: [PATCH] Allow pressing `escape` to cancel the current assistant
 generation (#10987)

If the assistant has already emitted some text, we will leave the
assistant message but maintain the cursor on the previous user message,
so that the user can easily discard the message by submitting again.

If no output was emitted yet, we simply delete the empty assistant
message.

Release Notes:

- N/A
---
 assets/keymaps/default-macos.json   |  3 ++-
 crates/assistant2/src/assistant2.rs | 32 ++++++++++++++++++++++++++---
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/assets/keymaps/default-macos.json b/assets/keymaps/default-macos.json
index f4da3078add322..4650df181d3369 100644
--- a/assets/keymaps/default-macos.json
+++ b/assets/keymaps/default-macos.json
@@ -212,7 +212,8 @@
     "context": "AssistantChat > Editor", // Used in the assistant2 crate
     "bindings": {
       "enter": ["assistant2::Submit", "Simple"],
-      "cmd-enter": ["assistant2::Submit", "Codebase"]
+      "cmd-enter": ["assistant2::Submit", "Codebase"],
+      "escape": "assistant2::Cancel"
     }
   },
   {
diff --git a/crates/assistant2/src/assistant2.rs b/crates/assistant2/src/assistant2.rs
index 683ce911af7afb..b89291bd133e52 100644
--- a/crates/assistant2/src/assistant2.rs
+++ b/crates/assistant2/src/assistant2.rs
@@ -34,8 +34,6 @@ pub use assistant_settings::AssistantSettings;
 
 const MAX_COMPLETION_CALLS_PER_SUBMISSION: usize = 5;
 
-// gpui::actions!(assistant, [Submit]);
-
 #[derive(Eq, PartialEq, Copy, Clone, Deserialize)]
 pub struct Submit(SubmitMode);
 
@@ -50,7 +48,7 @@ pub enum SubmitMode {
     Codebase,
 }
 
-gpui::actions!(assistant2, [ToggleFocus]);
+gpui::actions!(assistant2, [Cancel, ToggleFocus]);
 gpui::impl_actions!(assistant2, [Submit]);
 
 pub fn init(client: Arc<Client>, cx: &mut AppContext) {
@@ -256,6 +254,21 @@ impl AssistantChat {
         })
     }
 
+    fn cancel(&mut self, _: &Cancel, cx: &mut ViewContext<Self>) {
+        if self.pending_completion.take().is_none() {
+            cx.propagate();
+            return;
+        }
+
+        if let Some(ChatMessage::Assistant(message)) = self.messages.last() {
+            if message.body.text.is_empty() {
+                self.pop_message(cx);
+            } else {
+                self.push_new_user_message(false, cx);
+            }
+        }
+    }
+
     fn submit(&mut self, Submit(mode): &Submit, cx: &mut ViewContext<Self>) {
         let Some(focused_message_id) = self.focused_message_id(cx) else {
             log::error!("unexpected state: no user message editor is focused.");
@@ -282,6 +295,7 @@ impl AssistantChat {
                     .focus_handle(cx)
                     .contains_focused(cx);
                 this.push_new_user_message(focus, cx);
+                this.pending_completion = None;
             })
             .context("Failed to push new user message")
             .log_err();
@@ -453,6 +467,17 @@ impl AssistantChat {
         cx.notify();
     }
 
+    fn pop_message(&mut self, cx: &mut ViewContext<Self>) {
+        if self.messages.is_empty() {
+            return;
+        }
+
+        self.messages.pop();
+        self.list_state
+            .splice(self.messages.len()..self.messages.len() + 1, 0);
+        cx.notify();
+    }
+
     fn truncate_messages(&mut self, last_message_id: MessageId, cx: &mut ViewContext<Self>) {
         if let Some(index) = self.messages.iter().position(|message| match message {
             ChatMessage::User(message) => message.id == last_message_id,
@@ -677,6 +702,7 @@ impl Render for AssistantChat {
             .flex_1()
             .v_flex()
             .key_context("AssistantChat")
+            .on_action(cx.listener(Self::cancel))
             .text_color(Color::Default.color(cx))
             .child(self.render_model_dropdown(cx))
             .child(list(self.list_state.clone()).flex_1())