Allow pressing `escape` to cancel the current assistant generation (#10987)

Antonio Scandurra created 2 years ago

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

Change summary

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

Detailed changes

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"
     }
   },
   {

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())