Give edit steps multibuffer a title (#15625)

Antonio Scandurra and Nathan created

Release Notes:

- N/A

Co-authored-by: Nathan <nathan@zed.dev>

Change summary

assets/prompts/step_resolution.md       | 25 ++++++--
crates/assistant/src/assistant_panel.rs | 67 ++++++++++++++--------
crates/assistant/src/context.rs         | 80 +++++++++++++++-----------
3 files changed, 107 insertions(+), 65 deletions(-)

Detailed changes

assets/prompts/step_resolution.md 🔗

@@ -37,6 +37,7 @@ What are the operations for the step: <step>Add a new method 'calculate_area' to
 
 A (wrong):
 {
+  "title": "Add Rectangle methods",
   "operations": [
     {
       "kind": "AppendChild",
@@ -57,6 +58,7 @@ This demonstrates what NOT to do. NEVER append multiple children at the same loc
 
 A (corrected):
 {
+  "title": "Add Rectangle methods",
   "operations": [
     {
       "kind": "AppendChild",
@@ -72,6 +74,7 @@ What are the operations for the step: <step>Implement the 'Display' trait for th
 
 A:
 {
+  "title": "Implement Display for Rectangle",
   "operations": [
     {
       "kind": "InsertSiblingAfter",
@@ -110,6 +113,7 @@ What are the operations for the step: <step>Update the 'print_info' method to us
 
 A:
 {
+  "title": "Use formatted output",
   "operations": [
     {
       "kind": "Update",
@@ -125,13 +129,14 @@ What are the operations for the step: <step>Remove the 'email' field from the Us
 
 A:
 {
+  "title": "Remove email field",
   "operations": [
-    {
-      "kind": "Delete",
-      "path": "src/user.rs",
-      "symbol": "struct User email"
-    }
-  ]
+      {
+        "kind": "Delete",
+        "path": "src/user.rs",
+        "symbol": "struct User email"
+      }
+    ]
 }
 
 Example 3:
@@ -162,6 +167,7 @@ What are the operations for the step: <step>Add a 'use std::fmt;' statement at t
 
 A:
 {
+  "title": "Add use std::fmt statement",
   "operations": [
     {
       "kind": "PrependChild",
@@ -176,6 +182,7 @@ What are the operations for the step: <step>Add a new method 'start_engine' in t
 
 A:
 {
+  "title": "Add start_engine method",
   "operations": [
     {
       "kind": "InsertSiblingAfter",
@@ -219,6 +226,7 @@ What are the operations for the step: <step>Make salary an f32</step>
 
 A (wrong):
 {
+  "title": "Change salary to f32",
   "operations": [
     {
       "kind": "Update",
@@ -239,6 +247,7 @@ This example demonstrates what not to do. `struct Employee salary` is a child of
 
 A (corrected):
 {
+  "title": "Change salary to f32",
   "operations": [
     {
       "kind": "Update",
@@ -254,6 +263,7 @@ What are the correct operations for the step: <step>Remove the 'department' fiel
 
 A:
 {
+  "title": "Remove department",
   "operations": [
     {
       "kind": "Delete",
@@ -300,6 +310,7 @@ impl Game {
 
 A:
 {
+  "title": "Add level field to Player",
   "operations": [
     {
       "kind": "InsertSiblingAfter",
@@ -337,6 +348,7 @@ impl Config {
 
 A:
 {
+  "title": "Add load_from_file method",
   "operations": [
     {
       "kind": "PrependChild",
@@ -376,6 +388,7 @@ impl Database {
 
 A:
 {
+  "title": "Add error handling to query",
   "operations": [
     {
       "kind": "PrependChild",

crates/assistant/src/assistant_panel.rs 🔗

@@ -9,8 +9,8 @@ use crate::{
     },
     terminal_inline_assistant::TerminalInlineAssistant,
     Assist, ConfirmCommand, Context, ContextEvent, ContextId, ContextStore, CycleMessageRole,
-    DebugEditSteps, DeployHistory, DeployPromptLibrary, EditStep, EditStepOperations,
-    EditSuggestionGroup, InlineAssist, InlineAssistId, InlineAssistant, InsertIntoEditor,
+    DebugEditSteps, DeployHistory, DeployPromptLibrary, EditStep, EditStepState,
+    EditStepSuggestions, InlineAssist, InlineAssistId, InlineAssistant, InsertIntoEditor,
     MessageStatus, ModelSelector, PendingSlashCommand, PendingSlashCommandStatus, QuoteSelection,
     RemoteContextMetadata, SavedContextMetadata, Split, ToggleFocus, ToggleModelSelector,
 };
@@ -40,8 +40,7 @@ use gpui::{
 };
 use indexed_docs::IndexedDocsStore;
 use language::{
-    language_settings::SoftWrap, Buffer, Capability, LanguageRegistry, LspAdapterDelegate, Point,
-    ToOffset,
+    language_settings::SoftWrap, Capability, LanguageRegistry, LspAdapterDelegate, Point, ToOffset,
 };
 use language_model::{LanguageModelProvider, LanguageModelProviderId, LanguageModelRegistry, Role};
 use markdown::{Markdown, MarkdownStyle};
@@ -1454,18 +1453,19 @@ impl ContextEditor {
                     .text_for_range(step.source_range.clone())
                     .collect::<String>()
             ));
-            match &step.operations {
-                Some(EditStepOperations::Ready(operations)) => {
-                    output.push_str("Parsed Operations:\n");
-                    for op in operations {
+            match &step.state {
+                Some(EditStepState::Resolved(resolution)) => {
+                    output.push_str("Resolution:\n");
+                    output.push_str(&format!("  {:?}\n", resolution.step_title));
+                    for op in &resolution.operations {
                         output.push_str(&format!("  {:?}\n", op));
                     }
                 }
-                Some(EditStepOperations::Pending(_)) => {
-                    output.push_str("Operations: Pending\n");
+                Some(EditStepState::Pending(_)) => {
+                    output.push_str("Resolution: Pending\n");
                 }
                 None => {
-                    output.push_str("Operations: None\n");
+                    output.push_str("Resolution: None\n");
                 }
             }
             output.push('\n');
@@ -1906,12 +1906,18 @@ impl ContextEditor {
                     }
 
                     if let Some(new_active_step) = self.edit_step_for_cursor(cx) {
-                        let suggestions = new_active_step.edit_suggestions(&self.project, cx);
+                        let start = new_active_step.source_range.start;
+                        let open_editor = new_active_step
+                            .edit_suggestions(&self.project, cx)
+                            .map(|suggestions| {
+                                self.open_editor_for_edit_suggestions(suggestions, cx)
+                            })
+                            .unwrap_or_else(|| Task::ready(Ok(())));
                         self.active_edit_step = Some(ActiveEditStep {
-                            start: new_active_step.source_range.start,
+                            start,
                             assist_ids: Vec::new(),
                             editor: None,
-                            _open_editor: self.open_editor_for_edit_suggestions(suggestions, cx),
+                            _open_editor: open_editor,
                         });
                     }
                 }
@@ -1923,23 +1929,33 @@ impl ContextEditor {
 
     fn open_editor_for_edit_suggestions(
         &mut self,
-        edit_suggestions: Task<HashMap<Model<Buffer>, Vec<EditSuggestionGroup>>>,
+        edit_step_suggestions: Task<EditStepSuggestions>,
         cx: &mut ViewContext<Self>,
     ) -> Task<Result<()>> {
         let workspace = self.workspace.clone();
         let project = self.project.clone();
         let assistant_panel = self.assistant_panel.clone();
         cx.spawn(|this, mut cx| async move {
-            let edit_suggestions = edit_suggestions.await;
+            let edit_step_suggestions = edit_step_suggestions.await;
 
             let mut assist_ids = Vec::new();
-            let editor = if edit_suggestions.is_empty() {
+            let editor = if edit_step_suggestions.suggestions.is_empty() {
                 return Ok(());
-            } else if edit_suggestions.len() == 1
-                && edit_suggestions.values().next().unwrap().len() == 1
+            } else if edit_step_suggestions.suggestions.len() == 1
+                && edit_step_suggestions
+                    .suggestions
+                    .values()
+                    .next()
+                    .unwrap()
+                    .len()
+                    == 1
             {
                 // If there's only one buffer and one suggestion group, open it directly
-                let (buffer, suggestion_groups) = edit_suggestions.into_iter().next().unwrap();
+                let (buffer, suggestion_groups) = edit_step_suggestions
+                    .suggestions
+                    .into_iter()
+                    .next()
+                    .unwrap();
                 let suggestion_group = suggestion_groups.into_iter().next().unwrap();
                 let editor = workspace.update(&mut cx, |workspace, cx| {
                     let active_pane = workspace.active_pane().clone();
@@ -2004,8 +2020,9 @@ impl ContextEditor {
                 let mut inline_assist_suggestions = Vec::new();
                 let multibuffer = cx.new_model(|cx| {
                     let replica_id = project.read(cx).replica_id();
-                    let mut multibuffer = MultiBuffer::new(replica_id, Capability::ReadWrite);
-                    for (buffer, suggestion_groups) in edit_suggestions {
+                    let mut multibuffer = MultiBuffer::new(replica_id, Capability::ReadWrite)
+                        .with_title(edit_step_suggestions.title);
+                    for (buffer, suggestion_groups) in edit_step_suggestions.suggestions {
                         let excerpt_ids = multibuffer.push_excerpts(
                             buffer,
                             suggestion_groups
@@ -2358,9 +2375,9 @@ impl ContextEditor {
     fn render_send_button(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
         let focus_handle = self.focus_handle(cx).clone();
         let button_text = match self.edit_step_for_cursor(cx) {
-            Some(edit_step) => match &edit_step.operations {
-                Some(EditStepOperations::Pending(_)) => "Computing Changes...",
-                Some(EditStepOperations::Ready(_)) => "Apply Changes",
+            Some(edit_step) => match &edit_step.state {
+                Some(EditStepState::Pending(_)) => "Computing Changes...",
+                Some(EditStepState::Resolved(_)) => "Apply Changes",
                 None => "Send",
             },
             None => "Send",

crates/assistant/src/context.rs 🔗

@@ -341,7 +341,7 @@ pub struct SlashCommandId(clock::Lamport);
 #[derive(Debug)]
 pub struct EditStep {
     pub source_range: Range<language::Anchor>,
-    pub operations: Option<EditStepOperations>,
+    pub state: Option<EditStepState>,
 }
 
 #[derive(Debug)]
@@ -358,22 +358,29 @@ pub struct EditSuggestion {
     pub initial_insertion: Option<InitialInsertion>,
 }
 
+pub struct EditStepSuggestions {
+    pub title: String,
+    pub suggestions: HashMap<Model<Buffer>, Vec<EditSuggestionGroup>>,
+}
+
 impl EditStep {
     pub fn edit_suggestions(
         &self,
         project: &Model<Project>,
         cx: &AppContext,
-    ) -> Task<HashMap<Model<Buffer>, Vec<EditSuggestionGroup>>> {
-        let Some(EditStepOperations::Ready(operations)) = &self.operations else {
-            return Task::ready(HashMap::default());
+    ) -> Option<Task<EditStepSuggestions>> {
+        let Some(EditStepState::Resolved(resolution)) = &self.state else {
+            return None;
         };
 
-        let suggestion_tasks: Vec<_> = operations
+        let title = resolution.step_title.clone();
+        let suggestion_tasks: Vec<_> = resolution
+            .operations
             .iter()
             .map(|operation| operation.edit_suggestion(project.clone(), cx))
             .collect();
 
-        cx.spawn(|mut cx| async move {
+        Some(cx.spawn(|mut cx| async move {
             let suggestions = future::join_all(suggestion_tasks)
                 .await
                 .into_iter()
@@ -468,21 +475,24 @@ impl EditStep {
                 suggestion_groups_by_buffer.insert(buffer, suggestion_groups);
             }
 
-            suggestion_groups_by_buffer
-        })
+            EditStepSuggestions {
+                title,
+                suggestions: suggestion_groups_by_buffer,
+            }
+        }))
     }
 }
 
-pub enum EditStepOperations {
+pub enum EditStepState {
     Pending(Task<Option<()>>),
-    Ready(Vec<EditOperation>),
+    Resolved(EditStepResolution),
 }
 
-impl Debug for EditStepOperations {
+impl Debug for EditStepState {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            EditStepOperations::Pending(_) => write!(f, "EditStepOperations::Pending"),
-            EditStepOperations::Ready(operations) => f
+            EditStepState::Pending(_) => write!(f, "EditStepOperations::Pending"),
+            EditStepState::Resolved(operations) => f
                 .debug_struct("EditStepOperations::Parsed")
                 .field("operations", operations)
                 .finish(),
@@ -490,6 +500,25 @@ impl Debug for EditStepOperations {
     }
 }
 
+#[derive(Debug, Deserialize, JsonSchema)]
+pub struct EditStepResolution {
+    /// An extremely short title for the edit step represented by these operations.
+    pub step_title: String,
+    /// A sequence of operations to apply to the codebase.
+    /// When multiple operations are required for a step, be sure to include multiple operations in this list.
+    pub operations: Vec<EditOperation>,
+}
+
+impl LanguageModelTool for EditStepResolution {
+    fn name() -> String {
+        "edit".into()
+    }
+
+    fn description() -> String {
+        "suggest edits to one or more locations in the codebase".into()
+    }
+}
+
 /// A description of an operation to apply to one location in the codebase.
 ///
 /// This object represents a single edit operation that can be performed on a specific file
@@ -1324,7 +1353,7 @@ impl Context {
                                 ix,
                                 EditStep {
                                     source_range,
-                                    operations: None,
+                                    state: None,
                                 },
                             ));
                         }
@@ -1340,7 +1369,7 @@ impl Context {
         // Insert new steps and generate their corresponding tasks
         for (index, mut step) in new_edit_steps.into_iter().rev() {
             let task = self.generate_edit_step_operations(&step, cx);
-            step.operations = Some(EditStepOperations::Pending(task));
+            step.state = Some(EditStepState::Pending(task));
             self.edit_steps.insert(index, step);
         }
 
@@ -1353,23 +1382,6 @@ impl Context {
         edit_step: &EditStep,
         cx: &mut ModelContext<Self>,
     ) -> Task<Option<()>> {
-        #[derive(Debug, Deserialize, JsonSchema)]
-        struct EditTool {
-            /// A sequence of operations to apply to the codebase.
-            /// When multiple operations are required for a step, be sure to include multiple operations in this list.
-            operations: Vec<EditOperation>,
-        }
-
-        impl LanguageModelTool for EditTool {
-            fn name() -> String {
-                "edit".into()
-            }
-
-            fn description() -> String {
-                "suggest edits to one or more locations in the codebase".into()
-            }
-        }
-
         let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
             return Task::ready(Err(anyhow!("no active model")).log_err());
         };
@@ -1394,7 +1406,7 @@ impl Context {
                     content: prompt,
                 });
 
-                let tool_use = model.use_tool::<EditTool>(request, &cx).await?;
+                let resolution = model.use_tool::<EditStepResolution>(request, &cx).await?;
 
                 this.update(&mut cx, |this, cx| {
                     let step_index = this
@@ -1405,7 +1417,7 @@ impl Context {
                         })
                         .map_err(|_| anyhow!("edit step not found"))?;
                     if let Some(edit_step) = this.edit_steps.get_mut(step_index) {
-                        edit_step.operations = Some(EditStepOperations::Ready(tool_use.operations));
+                        edit_step.state = Some(EditStepState::Resolved(resolution));
                         cx.emit(ContextEvent::EditStepsChanged);
                     }
                     anyhow::Ok(())