From 9d4dd5c42b76d4f894d13ab739eff5141e531377 Mon Sep 17 00:00:00 2001
From: Antonio Scandurra <me@as-cii.com>
Date: Mon, 26 Jun 2023 15:57:36 +0200
Subject: [PATCH] Insert empty user message when assisting with the current
 last message

Co-Authored-By: Nathan Sobo <nathan@zed.dev>
---
 crates/ai/src/assistant.rs | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/crates/ai/src/assistant.rs b/crates/ai/src/assistant.rs
index 11a82aac4b..7dcd8830ed 100644
--- a/crates/ai/src/assistant.rs
+++ b/crates/ai/src/assistant.rs
@@ -1023,6 +1023,14 @@ impl Conversation {
     ) -> Vec<MessageAnchor> {
         let mut user_messages = Vec::new();
         let mut tasks = Vec::new();
+
+        let last_message_id = self.message_anchors.iter().rev().find_map(|message| {
+            message
+                .start
+                .is_valid(self.buffer.read(cx))
+                .then_some(message.id)
+        });
+
         for selected_message_id in selected_messages {
             let selected_message_role =
                 if let Some(metadata) = self.messages_metadata.get(&selected_message_id) {
@@ -1085,6 +1093,19 @@ impl Conversation {
                     )
                     .unwrap();
 
+                // Queue up the user's next reply
+                if Some(selected_message_id) == last_message_id {
+                    let user_message = self
+                        .insert_message_after(
+                            assistant_message.id,
+                            Role::User,
+                            MessageStatus::Done,
+                            cx,
+                        )
+                        .unwrap();
+                    user_messages.push(user_message);
+                }
+
                 tasks.push(cx.spawn_weak({
                     |this, mut cx| async move {
                         let assistant_message_id = assistant_message.id;