Allow for multi-cursor assist and cycle_role actions

Co-Authored-By: Nathan Sobo <nathan@zed.dev>
Co-Authored-By: Kyle Caverly <kyle@zed.dev>

crates/ai/src/assistant.rs

@@ -22,7 +22,7 @@ use gpui::{
     Subscription, Task, View, ViewContext, ViewHandle, WeakViewHandle, WindowContext,
 };
 use isahc::{http::StatusCode, Request, RequestExt};
-use language::{language_settings::SoftWrap, Buffer, LanguageRegistry, Selection, ToOffset as _};
+use language::{language_settings::SoftWrap, Buffer, LanguageRegistry, ToOffset as _};
 use serde::Deserialize;
 use settings::SettingsStore;
 use std::{
@@ -591,106 +591,129 @@ impl Assistant {
 
     fn assist(
         &mut self,
-        selection: Selection<usize>,
+        selected_messages: HashSet<MessageId>,
         cx: &mut ModelContext<Self>,
-    ) -> Option<(MessageAnchor, MessageAnchor)> {
+    ) -> Vec<MessageAnchor> {
-        let request = OpenAIRequest {
+        let mut user_messages = Vec::new();
-            model: self.model.clone(),
+        for selected_message_id in selected_messages {
-            messages: self
+            let selected_message_role =
-                .messages(cx)
+                if let Some(metadata) = self.messages_metadata.get(&selected_message_id) {
-                .map(|message| message.to_open_ai_message(self.buffer.read(cx)))
+                    metadata.role
-                .collect(),
+                } else {
-            stream: true,
+                    continue;
-        };
+                };
+            let Some(user_message) = self.insert_message_after(selected_message_id, Role::User, cx) else {
+                continue;
+            };
+            user_messages.push(user_message);
+            if selected_message_role == Role::User {
+                let request = OpenAIRequest {
+                    model: self.model.clone(),
+                    messages: self
+                        .messages(cx)
+                        .map(|message| message.to_open_ai_message(self.buffer.read(cx)))
+                        .chain(Some(RequestMessage {
+                            role: Role::System,
+                            content: format!(
+                                "Direct your reply to message with id {}. Do not include a [Message X] header.",
+                                selected_message_id.0
+                            ),
+                        }))
+                        .collect(),
+                    stream: true,
+                };
 
-        let api_key = self.api_key.borrow().clone()?;
+                let Some(api_key) = self.api_key.borrow().clone() else { continue };
-        let stream = stream_completion(api_key, cx.background().clone(), request);
+                let stream = stream_completion(api_key, cx.background().clone(), request);
-        let assistant_message = self.insert_message_after(
+                let assistant_message = self
-            self.message_for_offset(selection.head(), cx)?.id,
+                    .insert_message_after(selected_message_id, Role::Assistant, cx)
-            Role::Assistant,
+                    .unwrap();
-            cx,
-        )?;
-        let user_message = self.insert_message_after(assistant_message.id, Role::User, cx)?;
 
-        let task = cx.spawn_weak({
+                let task = cx.spawn_weak({
-            |this, mut cx| async move {
+                    |this, mut cx| async move {
-                let assistant_message_id = assistant_message.id;
+                        let assistant_message_id = assistant_message.id;
-                let stream_completion = async {
+                        let stream_completion = async {
-                    let mut messages = stream.await?;
+                            let mut messages = stream.await?;
+
+                            while let Some(message) = messages.next().await {
+                                let mut message = message?;
+                                if let Some(choice) = message.choices.pop() {
+                                    this.upgrade(&cx)
+                                        .ok_or_else(|| anyhow!("assistant was dropped"))?
+                                        .update(&mut cx, |this, cx| {
+                                            let text: Arc<str> = choice.delta.content?.into();
+                                            let message_ix = this.message_anchors.iter().position(
+                                                |message| message.id == assistant_message_id,
+                                            )?;
+                                            this.buffer.update(cx, |buffer, cx| {
+                                                let offset = if message_ix + 1
+                                                    == this.message_anchors.len()
+                                                {
+                                                    buffer.len()
+                                                } else {
+                                                    this.message_anchors[message_ix + 1]
+                                                        .start
+                                                        .to_offset(buffer)
+                                                        .saturating_sub(1)
+                                                };
+                                                buffer.edit([(offset..offset, text)], None, cx);
+                                            });
+                                            cx.emit(AssistantEvent::StreamedCompletion);
+
+                                            Some(())
+                                        });
+                                }
+                            }
 
-                    while let Some(message) = messages.next().await {
-                        let mut message = message?;
-                        if let Some(choice) = message.choices.pop() {
                             this.upgrade(&cx)
                                 .ok_or_else(|| anyhow!("assistant was dropped"))?
                                 .update(&mut cx, |this, cx| {
-                                    let text: Arc<str> = choice.delta.content?.into();
+                                    this.pending_completions.retain(|completion| {
-                                    let message_ix = this
+                                        completion.id != this.completion_count
-                                        .message_anchors
-                                        .iter()
-                                        .position(|message| message.id == assistant_message_id)?;
-                                    this.buffer.update(cx, |buffer, cx| {
-                                        let offset = if message_ix + 1 == this.message_anchors.len()
-                                        {
-                                            buffer.len()
-                                        } else {
-                                            this.message_anchors[message_ix + 1]
-                                                .start
-                                                .to_offset(buffer)
-                                                .saturating_sub(1)
-                                        };
-                                        buffer.edit([(offset..offset, text)], None, cx);
                                     });
-                                    cx.emit(AssistantEvent::StreamedCompletion);
+                                    this.summarize(cx);
-
-                                    Some(())
                                 });
+
+                            anyhow::Ok(())
+                        };
+
+                        let result = stream_completion.await;
+                        if let Some(this) = this.upgrade(&cx) {
+                            this.update(&mut cx, |this, cx| {
+                                if let Err(error) = result {
+                                    if let Some(metadata) =
+                                        this.messages_metadata.get_mut(&assistant_message.id)
+                                    {
+                                        metadata.error = Some(error.to_string().trim().into());
+                                        cx.notify();
+                                    }
+                                }
+                            });
                         }
                     }
+                });
 
-                    this.upgrade(&cx)
+                self.pending_completions.push(PendingCompletion {
-                        .ok_or_else(|| anyhow!("assistant was dropped"))?
+                    id: post_inc(&mut self.completion_count),
-                        .update(&mut cx, |this, cx| {
+                    _task: task,
-                            this.pending_completions
+                });
-                                .retain(|completion| completion.id != this.completion_count);
-                            this.summarize(cx);
-                        });
-
-                    anyhow::Ok(())
-                };
-
-                let result = stream_completion.await;
-                if let Some(this) = this.upgrade(&cx) {
-                    this.update(&mut cx, |this, cx| {
-                        if let Err(error) = result {
-                            if let Some(metadata) =
-                                this.messages_metadata.get_mut(&assistant_message.id)
-                            {
-                                metadata.error = Some(error.to_string().trim().into());
-                                cx.notify();
-                            }
-                        }
-                    });
-                }
             }
-        });
+        }
 
-        self.pending_completions.push(PendingCompletion {
+        user_messages
-            id: post_inc(&mut self.completion_count),
-            _task: task,
-        });
-        Some((assistant_message, user_message))
     }
 
     fn cancel_last_assist(&mut self) -> bool {
         self.pending_completions.pop().is_some()
     }
 
-    fn cycle_message_role(&mut self, id: MessageId, cx: &mut ModelContext<Self>) {
+    fn cycle_message_roles(&mut self, ids: HashSet<MessageId>, cx: &mut ModelContext<Self>) {
-        if let Some(metadata) = self.messages_metadata.get_mut(&id) {
+        for id in ids {
-            metadata.role.cycle();
+            if let Some(metadata) = self.messages_metadata.get_mut(&id) {
-            cx.emit(AssistantEvent::MessagesEdited);
+                metadata.role.cycle();
-            cx.notify();
+                cx.emit(AssistantEvent::MessagesEdited);
+                cx.notify();
+            }
         }
     }
 
@@ -884,14 +907,39 @@ impl Assistant {
         }
     }
 
-    fn message_for_offset<'a>(&'a self, offset: usize, cx: &'a AppContext) -> Option<Message> {
+    fn message_for_offset(&self, offset: usize, cx: &AppContext) -> Option<Message> {
+        self.messages_for_offsets([offset], cx).pop()
+    }
+
+    fn messages_for_offsets(
+        &self,
+        offsets: impl IntoIterator<Item = usize>,
+        cx: &AppContext,
+    ) -> Vec<Message> {
+        let mut result = Vec::new();
+
+        let buffer_len = self.buffer.read(cx).len();
         let mut messages = self.messages(cx).peekable();
-        while let Some(message) = messages.next() {
+        let mut offsets = offsets.into_iter().peekable();
-            if message.range.contains(&offset) || messages.peek().is_none() {
+        while let Some(offset) = offsets.next() {
-                return Some(message);
+            // Skip messages that start after the offset.
+            while messages.peek().map_or(false, |message| {
+                message.range.end < offset || (message.range.end == offset && offset < buffer_len)
+            }) {
+                messages.next();
             }
+            let Some(message) = messages.peek() else { continue };
+
+            // Skip offsets that are in the same message.
+            while offsets.peek().map_or(false, |offset| {
+                message.range.contains(offset) || message.range.end == buffer_len
+            }) {
+                offsets.next();
+            }
+
+            result.push(message.clone());
         }
-        None
+        result
     }
 
     fn messages<'a>(&'a self, cx: &'a AppContext) -> impl 'a + Iterator<Item = Message> {
@@ -983,24 +1031,32 @@ impl AssistantEditor {
     }
 
     fn assist(&mut self, _: &Assist, cx: &mut ViewContext<Self>) {
-        let selection = self.editor.read(cx).selections.newest(cx);
+        let cursors = self.cursors(cx);
-        let user_message = self.assistant.update(cx, |assistant, cx| {
-            let (_, user_message) = assistant.assist(selection, cx)?;
-            Some(user_message)
-        });
 
-        if let Some(user_message) = user_message {
+        let user_messages = self.assistant.update(cx, |assistant, cx| {
-            let cursor = user_message
+            let selected_messages = assistant
-                .start
+                .messages_for_offsets(cursors, cx)
-                .to_offset(&self.assistant.read(cx).buffer.read(cx));
+                .into_iter()
-            self.editor.update(cx, |editor, cx| {
+                .map(|message| message.id)
-                editor.change_selections(
+                .collect();
-                    Some(Autoscroll::Strategy(AutoscrollStrategy::Fit)),
+            assistant.assist(selected_messages, cx)
-                    cx,
+        });
-                    |selections| selections.select_ranges([cursor..cursor]),
+        let new_selections = user_messages
-                );
+            .iter()
-            });
+            .map(|message| {
-        }
+                let cursor = message
+                    .start
+                    .to_offset(self.assistant.read(cx).buffer.read(cx));
+                cursor..cursor
+            })
+            .collect::<Vec<_>>();
+        self.editor.update(cx, |editor, cx| {
+            editor.change_selections(
+                Some(Autoscroll::Strategy(AutoscrollStrategy::Fit)),
+                cx,
+                |selections| selections.select_ranges(new_selections),
+            );
+        });
     }
 
     fn cancel_last_assist(&mut self, _: &editor::Cancel, cx: &mut ViewContext<Self>) {
@@ -1013,14 +1069,25 @@ impl AssistantEditor {
     }
 
     fn cycle_message_role(&mut self, _: &CycleMessageRole, cx: &mut ViewContext<Self>) {
-        let cursor_offset = self.editor.read(cx).selections.newest(cx).head();
+        let cursors = self.cursors(cx);
         self.assistant.update(cx, |assistant, cx| {
-            if let Some(message) = assistant.message_for_offset(cursor_offset, cx) {
+            let messages = assistant
-                assistant.cycle_message_role(message.id, cx);
+                .messages_for_offsets(cursors, cx)
-            }
+                .into_iter()
+                .map(|message| message.id)
+                .collect();
+            assistant.cycle_message_roles(messages, cx)
         });
     }
 
+    fn cursors(&self, cx: &AppContext) -> Vec<usize> {
+        let selections = self.editor.read(cx).selections.all::<usize>(cx);
+        selections
+            .into_iter()
+            .map(|selection| selection.head())
+            .collect()
+    }
+
     fn handle_assistant_event(
         &mut self,
         _: ModelHandle<Assistant>,
@@ -1149,7 +1216,10 @@ impl AssistantEditor {
                                 let assistant = assistant.clone();
                                 move |_, _, cx| {
                                     assistant.update(cx, |assistant, cx| {
-                                        assistant.cycle_message_role(message_id, cx)
+                                        assistant.cycle_message_roles(
+                                            HashSet::from_iter(Some(message_id)),
+                                            cx,
+                                        )
                                     })
                                 }
                             });
@@ -1444,9 +1514,11 @@ pub struct Message {
 
 impl Message {
     fn to_open_ai_message(&self, buffer: &Buffer) -> RequestMessage {
+        let mut content = format!("[Message {}]\n", self.id.0).to_string();
+        content.extend(buffer.text_for_range(self.range.clone()));
         RequestMessage {
             role: self.role,
-            content: buffer.text_for_range(self.range.clone()).collect(),
+            content,
         }
     }
 }
@@ -1761,6 +1833,66 @@ mod tests {
         );
     }
 
+    #[gpui::test]
+    fn test_messages_for_offsets(cx: &mut AppContext) {
+        let registry = Arc::new(LanguageRegistry::test());
+        let assistant = cx.add_model(|cx| Assistant::new(Default::default(), registry, cx));
+        let buffer = assistant.read(cx).buffer.clone();
+
+        let message_1 = assistant.read(cx).message_anchors[0].clone();
+        assert_eq!(
+            messages(&assistant, cx),
+            vec![(message_1.id, Role::User, 0..0)]
+        );
+
+        buffer.update(cx, |buffer, cx| buffer.edit([(0..0, "aaa")], None, cx));
+        let message_2 = assistant
+            .update(cx, |assistant, cx| {
+                assistant.insert_message_after(message_1.id, Role::User, cx)
+            })
+            .unwrap();
+        buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "bbb")], None, cx));
+
+        let message_3 = assistant
+            .update(cx, |assistant, cx| {
+                assistant.insert_message_after(message_2.id, Role::User, cx)
+            })
+            .unwrap();
+        buffer.update(cx, |buffer, cx| buffer.edit([(8..8, "ccc")], None, cx));
+
+        assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc");
+        assert_eq!(
+            messages(&assistant, cx),
+            vec![
+                (message_1.id, Role::User, 0..4),
+                (message_2.id, Role::User, 4..8),
+                (message_3.id, Role::User, 8..11)
+            ]
+        );
+
+        assert_eq!(
+            message_ids_for_offsets(&assistant, &[0, 4, 9], cx),
+            [message_1.id, message_2.id, message_3.id]
+        );
+        assert_eq!(
+            message_ids_for_offsets(&assistant, &[0, 1, 11], cx),
+            [message_1.id, message_3.id]
+        );
+
+        fn message_ids_for_offsets(
+            assistant: &ModelHandle<Assistant>,
+            offsets: &[usize],
+            cx: &AppContext,
+        ) -> Vec<MessageId> {
+            assistant
+                .read(cx)
+                .messages_for_offsets(offsets.iter().copied(), cx)
+                .into_iter()
+                .map(|message| message.id)
+                .collect()
+        }
+    }
+
     fn messages(
         assistant: &ModelHandle<Assistant>,
         cx: &AppContext,