agent: Delay edit tool buffer clearing until the first chunk is sent (#49633)

Lukas Wirth created 1 month ago

Release Notes:

- The agent edit tool no longer clears files until the first edit comes
in, preventing a buffer being empty for prolonged time if the agent is
slow in reporting the first text chunk

Change summary

crates/agent/src/edit_agent.rs | 106 ++++++++++++++++++++++++-----------
1 file changed, 71 insertions(+), 35 deletions(-)

Detailed changes

crates/agent/src/edit_agent.rs

@@ -166,56 +166,69 @@ impl EditAgent {
         output_events_tx: mpsc::UnboundedSender<EditAgentOutputEvent>,
         cx: &mut AsyncApp,
     ) -> Result<()> {
-        cx.update(|cx| {
-            buffer.update(cx, |buffer, cx| buffer.set_text("", cx));
-            self.action_log.update(cx, |log, cx| {
-                log.buffer_edited(buffer.clone(), cx);
-            });
+        let buffer_id = cx.update(|cx| {
+            let buffer_id = buffer.read(cx).remote_id();
             self.project.update(cx, |project, cx| {
                 project.set_agent_location(
                     Some(AgentLocation {
                         buffer: buffer.downgrade(),
-                        position: language::Anchor::max_for_buffer(buffer.read(cx).remote_id()),
+                        position: language::Anchor::min_for_buffer(buffer_id),
                     }),
                     cx,
                 )
             });
+            buffer_id
+        });
+
+        let send_edit_event = || {
             output_events_tx
                 .unbounded_send(EditAgentOutputEvent::Edited(
-                    Anchor::min_max_range_for_buffer(buffer.read(cx).remote_id()),
+                    Anchor::min_max_range_for_buffer(buffer_id),
                 ))
-                .ok();
-        });
-
+                .ok()
+        };
+        let set_agent_location = |cx: &mut _| {
+            self.project.update(cx, |project, cx| {
+                project.set_agent_location(
+                    Some(AgentLocation {
+                        buffer: buffer.downgrade(),
+                        position: language::Anchor::max_for_buffer(buffer_id),
+                    }),
+                    cx,
+                )
+            })
+        };
+        let mut first_chunk = true;
         while let Some(event) = parse_rx.next().await {
             match event? {
                 CreateFileParserEvent::NewTextChunk { chunk } => {
-                    let buffer_id = cx.update(|cx| {
-                        buffer.update(cx, |buffer, cx| buffer.append(chunk, cx));
+                    cx.update(|cx| {
+                        buffer.update(cx, |buffer, cx| {
+                            if mem::take(&mut first_chunk) {
+                                buffer.set_text(chunk, cx)
+                            } else {
+                                buffer.append(chunk, cx)
+                            }
+                        });
                         self.action_log
                             .update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
-                        self.project.update(cx, |project, cx| {
-                            project.set_agent_location(
-                                Some(AgentLocation {
-                                    buffer: buffer.downgrade(),
-                                    position: language::Anchor::max_for_buffer(
-                                        buffer.read(cx).remote_id(),
-                                    ),
-                                }),
-                                cx,
-                            )
-                        });
-                        buffer.read(cx).remote_id()
+                        set_agent_location(cx);
                     });
-                    output_events_tx
-                        .unbounded_send(EditAgentOutputEvent::Edited(
-                            Anchor::min_max_range_for_buffer(buffer_id),
-                        ))
-                        .ok();
+                    send_edit_event();
                 }
             }
         }
 
+        if first_chunk {
+            cx.update(|cx| {
+                buffer.update(cx, |buffer, cx| buffer.set_text("", cx));
+                self.action_log
+                    .update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
+                set_agent_location(cx);
+            });
+            send_edit_event();
+        }
+
         Ok(())
     }
 
@@ -1194,19 +1207,16 @@ mod tests {
         );
 
         cx.run_until_parked();
-        assert_matches!(
-            drain_events(&mut events).as_slice(),
-            [EditAgentOutputEvent::Edited(_)]
-        );
+        assert_eq!(drain_events(&mut events).as_slice(), []);
         assert_eq!(
             buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
-            ""
+            "abc\ndef\nghi"
         );
         assert_eq!(
             project.read_with(cx, |project, _| project.agent_location()),
             Some(AgentLocation {
                 buffer: buffer.downgrade(),
-                position: language::Anchor::max_for_buffer(
+                position: language::Anchor::min_for_buffer(
                     cx.update(|cx| buffer.read(cx).remote_id())
                 ),
             })
@@ -1290,6 +1300,32 @@ mod tests {
         );
     }
 
+    #[gpui::test]
+    async fn test_overwrite_no_content(cx: &mut TestAppContext) {
+        let agent = init_test(cx).await;
+        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
+        let (chunks_tx, chunks_rx) = mpsc::unbounded::<&str>();
+        let (apply, mut events) = agent.overwrite_with_chunks(
+            buffer.clone(),
+            chunks_rx.map(|chunk| Ok(chunk.to_string())),
+            &mut cx.to_async(),
+        );
+
+        drop(chunks_tx);
+        cx.run_until_parked();
+
+        let result = apply.await;
+        assert!(result.is_ok(),);
+        assert_matches!(
+            drain_events(&mut events).as_slice(),
+            [EditAgentOutputEvent::Edited { .. }]
+        );
+        assert_eq!(
+            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
+            ""
+        );
+    }
+
     #[gpui::test(iterations = 100)]
     async fn test_indent_new_text_chunks(mut rng: StdRng) {
         let chunks = to_random_chunks(&mut rng, "    abc\n  def\n      ghi");