Don't show deleted hunks when agent overwrites file (#29918)

Antonio Scandurra created 7 months ago

Release Notes:

- Improved display of diffs when the agent rewrites a file from scratch.

Change summary

crates/agent/src/agent_diff.rs                 |  6 +-
crates/agent/src/thread.rs                     |  2 
crates/assistant_tool/src/action_log.rs        | 46 ++++++++++---------
crates/assistant_tool/src/outline.rs           |  2 
crates/assistant_tools/src/create_file_tool.rs |  2 
crates/assistant_tools/src/edit_agent.rs       | 18 +++---
crates/assistant_tools/src/edit_file_tool.rs   |  3 
crates/assistant_tools/src/read_file_tool.rs   |  4 
8 files changed, 43 insertions(+), 40 deletions(-)

Detailed changes

crates/agent/src/agent_diff.rs

@@ -1816,7 +1816,7 @@ mod tests {
             .await
             .unwrap();
         cx.update(|_, cx| {
-            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| {
                 buffer
                     .edit(
@@ -2031,7 +2031,7 @@ mod tests {
 
         // Make changes
         cx.update(|_, cx| {
-            action_log.update(cx, |log, cx| log.track_buffer(buffer1.clone(), cx));
+            action_log.update(cx, |log, cx| log.buffer_read(buffer1.clone(), cx));
             buffer1.update(cx, |buffer, cx| {
                 buffer
                     .edit(
@@ -2048,7 +2048,7 @@ mod tests {
             });
             action_log.update(cx, |log, cx| log.buffer_edited(buffer1.clone(), cx));
 
-            action_log.update(cx, |log, cx| log.track_buffer(buffer2.clone(), cx));
+            action_log.update(cx, |log, cx| log.buffer_read(buffer2.clone(), cx));
             buffer2.update(cx, |buffer, cx| {
                 buffer
                     .edit(

crates/agent/src/thread.rs

@@ -898,7 +898,7 @@ impl Thread {
         if !loaded_context.referenced_buffers.is_empty() {
             self.action_log.update(cx, |log, cx| {
                 for buffer in loaded_context.referenced_buffers {
-                    log.track_buffer(buffer, cx);
+                    log.buffer_read(buffer, cx);
                 }
             });
         }

crates/assistant_tool/src/action_log.rs

@@ -46,6 +46,7 @@ impl ActionLog {
     fn track_buffer_internal(
         &mut self,
         buffer: Entity<Buffer>,
+        is_created: bool,
         cx: &mut Context<Self>,
     ) -> &mut TrackedBuffer {
         let tracked_buffer = self
@@ -62,11 +63,7 @@ impl ActionLog {
                 let base_text;
                 let status;
                 let unreviewed_changes;
-                if buffer
-                    .read(cx)
-                    .file()
-                    .map_or(true, |file| !file.disk_state().exists())
-                {
+                if is_created {
                     base_text = Rope::default();
                     status = TrackedBufferStatus::Created;
                     unreviewed_changes = Patch::new(vec![Edit {
@@ -153,7 +150,7 @@ impl ActionLog {
                     // resurrected externally, we want to clear the changes we
                     // were tracking and reset the buffer's state.
                     self.tracked_buffers.remove(&buffer);
-                    self.track_buffer_internal(buffer, cx);
+                    self.track_buffer_internal(buffer, false, cx);
                 }
                 cx.notify();
             }
@@ -267,15 +264,22 @@ impl ActionLog {
     }
 
     /// Track a buffer as read, so we can notify the model about user edits.
-    pub fn track_buffer(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
-        self.track_buffer_internal(buffer, cx);
+    pub fn buffer_read(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
+        self.track_buffer_internal(buffer, false, cx);
+    }
+
+    /// Mark a buffer as edited, so we can refresh it in the context
+    pub fn buffer_created(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
+        self.edited_since_project_diagnostics_check = true;
+        self.tracked_buffers.remove(&buffer);
+        self.track_buffer_internal(buffer.clone(), true, cx);
     }
 
     /// Mark a buffer as edited, so we can refresh it in the context
     pub fn buffer_edited(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
         self.edited_since_project_diagnostics_check = true;
 
-        let tracked_buffer = self.track_buffer_internal(buffer.clone(), cx);
+        let tracked_buffer = self.track_buffer_internal(buffer.clone(), false, cx);
         if let TrackedBufferStatus::Deleted = tracked_buffer.status {
             tracked_buffer.status = TrackedBufferStatus::Modified;
         }
@@ -283,7 +287,7 @@ impl ActionLog {
     }
 
     pub fn will_delete_buffer(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
-        let tracked_buffer = self.track_buffer_internal(buffer.clone(), cx);
+        let tracked_buffer = self.track_buffer_internal(buffer.clone(), false, cx);
         match tracked_buffer.status {
             TrackedBufferStatus::Created => {
                 self.tracked_buffers.remove(&buffer);
@@ -393,7 +397,7 @@ impl ActionLog {
 
                 // Clear all tracked changes for this buffer and start over as if we just read it.
                 self.tracked_buffers.remove(&buffer);
-                self.track_buffer_internal(buffer.clone(), cx);
+                self.buffer_read(buffer.clone(), cx);
                 cx.notify();
                 save
             }
@@ -704,7 +708,7 @@ mod tests {
             .unwrap();
 
         cx.update(|cx| {
-            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| {
                 buffer
                     .edit([(Point::new(1, 1)..Point::new(1, 2), "E")], None, cx)
@@ -785,7 +789,7 @@ mod tests {
             .unwrap();
 
         cx.update(|cx| {
-            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| {
                 buffer
                     .edit([(Point::new(1, 0)..Point::new(2, 0), "")], None, cx)
@@ -867,7 +871,7 @@ mod tests {
             .unwrap();
 
         cx.update(|cx| {
-            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| {
                 buffer
                     .edit([(Point::new(1, 2)..Point::new(2, 3), "F\nGHI")], None, cx)
@@ -963,7 +967,7 @@ mod tests {
             .await
             .unwrap();
         cx.update(|cx| {
-            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.buffer_created(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| buffer.set_text("lorem", cx));
             action_log.update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
         });
@@ -1086,7 +1090,7 @@ mod tests {
             .update(cx, |project, cx| project.open_buffer(file2_path, cx))
             .await
             .unwrap();
-        action_log.update(cx, |log, cx| log.track_buffer(buffer2.clone(), cx));
+        action_log.update(cx, |log, cx| log.buffer_read(buffer2.clone(), cx));
         buffer2.update(cx, |buffer, cx| buffer.set_text("IPSUM", cx));
         action_log.update(cx, |log, cx| log.buffer_edited(buffer2.clone(), cx));
         project
@@ -1133,7 +1137,7 @@ mod tests {
             .unwrap();
 
         cx.update(|cx| {
-            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| {
                 buffer
                     .edit([(Point::new(1, 1)..Point::new(1, 2), "E\nXYZ")], None, cx)
@@ -1268,7 +1272,7 @@ mod tests {
             .unwrap();
 
         cx.update(|cx| {
-            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| {
                 buffer
                     .edit([(Point::new(1, 1)..Point::new(1, 2), "E\nXYZ")], None, cx)
@@ -1401,7 +1405,7 @@ mod tests {
             .await
             .unwrap();
         cx.update(|cx| {
-            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.buffer_created(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| buffer.set_text("content", cx));
             action_log.update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
         });
@@ -1459,7 +1463,7 @@ mod tests {
             .await
             .unwrap();
 
-        action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
+        action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
 
         for _ in 0..operations {
             match rng.gen_range(0..100) {
@@ -1511,7 +1515,7 @@ mod tests {
             log::info!("quiescing...");
             cx.run_until_parked();
             action_log.update(cx, |log, cx| {
-                let tracked_buffer = log.track_buffer_internal(buffer.clone(), cx);
+                let tracked_buffer = log.tracked_buffers.get(&buffer).unwrap();
                 let mut old_text = tracked_buffer.base_text.clone();
                 let new_text = buffer.read(cx).as_rope();
                 for edit in tracked_buffer.unreviewed_changes.edits() {

crates/assistant_tool/src/outline.rs

@@ -31,7 +31,7 @@ pub async fn file_outline(
     };
 
     action_log.update(cx, |action_log, cx| {
-        action_log.track_buffer(buffer.clone(), cx);
+        action_log.buffer_read(buffer.clone(), cx);
     })?;
 
     // Wait until the buffer has been fully parsed, so that we can read its outline.

crates/assistant_tools/src/create_file_tool.rs

@@ -118,7 +118,7 @@ impl Tool for CreateFileTool {
                 .map_err(|err| anyhow!("Unable to open buffer for {destination_path}: {err}"))?;
             cx.update(|cx| {
                 action_log.update(cx, |action_log, cx| {
-                    action_log.track_buffer(buffer.clone(), cx)
+                    action_log.buffer_created(buffer.clone(), cx)
                 });
                 buffer.update(cx, |buffer, cx| buffer.set_text(contents, cx));
                 action_log.update(cx, |action_log, cx| {

crates/assistant_tools/src/edit_agent.rs

@@ -101,7 +101,7 @@ impl EditAgent {
             .render(&this.templates)?;
             let new_chunks = this.request(previous_messages, prompt, cx).await?;
 
-            let (output, mut inner_events) = this.replace_text_with_chunks(buffer, new_chunks, cx);
+            let (output, mut inner_events) = this.overwrite_with_chunks(buffer, new_chunks, cx);
             while let Some(event) = inner_events.next().await {
                 events_tx.unbounded_send(event).ok();
             }
@@ -110,7 +110,7 @@ impl EditAgent {
         (output, events_rx)
     }
 
-    fn replace_text_with_chunks(
+    fn overwrite_with_chunks(
         &self,
         buffer: Entity<Buffer>,
         edit_chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
@@ -123,9 +123,9 @@ impl EditAgent {
         let this = self.clone();
         let task = cx.spawn(async move |cx| {
             this.action_log
-                .update(cx, |log, cx| log.track_buffer(buffer.clone(), cx))?;
+                .update(cx, |log, cx| log.buffer_created(buffer.clone(), cx))?;
             let output = this
-                .replace_text_with_chunks_internal(buffer, edit_chunks, output_events_tx, cx)
+                .overwrite_with_chunks_internal(buffer, edit_chunks, output_events_tx, cx)
                 .await;
             this.project
                 .update(cx, |project, cx| project.set_agent_location(None, cx))?;
@@ -134,7 +134,7 @@ impl EditAgent {
         (task, output_events_rx)
     }
 
-    async fn replace_text_with_chunks_internal(
+    async fn overwrite_with_chunks_internal(
         &self,
         buffer: Entity<Buffer>,
         edit_chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
@@ -246,9 +246,9 @@ impl EditAgent {
         let this = self.clone();
         let task = cx.spawn(async move |mut cx| {
             this.action_log
-                .update(cx, |log, cx| log.track_buffer(buffer.clone(), cx))?;
+                .update(cx, |log, cx| log.buffer_read(buffer.clone(), cx))?;
             let output = this
-                .apply_edits_internal(buffer, edit_chunks, output_events_tx, &mut cx)
+                .apply_edit_chunks_internal(buffer, edit_chunks, output_events_tx, &mut cx)
                 .await;
             this.project
                 .update(cx, |project, cx| project.set_agent_location(None, cx))?;
@@ -257,7 +257,7 @@ impl EditAgent {
         (task, output_events_rx)
     }
 
-    async fn apply_edits_internal(
+    async fn apply_edit_chunks_internal(
         &self,
         buffer: Entity<Buffer>,
         edit_chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
@@ -1027,7 +1027,7 @@ mod tests {
             .read_with(cx, |log, _| log.project().clone());
         let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
         let (chunks_tx, chunks_rx) = mpsc::unbounded();
-        let (apply, mut events) = agent.replace_text_with_chunks(
+        let (apply, mut events) = agent.overwrite_with_chunks(
             buffer.clone(),
             chunks_rx.map(|chunk: &str| Ok(chunk.to_string())),
             &mut cx.to_async(),

crates/assistant_tools/src/edit_file_tool.rs

@@ -239,8 +239,7 @@ impl Tool for EditFileTool {
             };
 
             let snapshot = cx.update(|cx| {
-                action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
-
+                action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
                 let base_version = diff.base_version.clone();
                 let snapshot = buffer.update(cx, |buffer, cx| {
                     buffer.finalize_last_transaction();

crates/assistant_tools/src/read_file_tool.rs

@@ -152,7 +152,7 @@ impl Tool for ReadFileTool {
                 })?;
 
                 action_log.update(cx, |log, cx| {
-                    log.track_buffer(buffer.clone(), cx);
+                    log.buffer_read(buffer.clone(), cx);
                 })?;
 
                 if let Some(anchor) = anchor {
@@ -177,7 +177,7 @@ impl Tool for ReadFileTool {
                     let result = buffer.read_with(cx, |buffer, _cx| buffer.text())?;
 
                     action_log.update(cx, |log, cx| {
-                        log.track_buffer(buffer, cx);
+                        log.buffer_read(buffer, cx);
                     })?;
 
                     Ok(result)