agent: Support providers streaming edits out of order (#51286)

Bennet Bo Fenner created

Release Notes:

- N/A

Change summary

crates/agent/src/tools/streaming_edit_file_tool.rs | 127 +++++++++++++++
1 file changed, 124 insertions(+), 3 deletions(-)

Detailed changes

crates/agent/src/tools/streaming_edit_file_tool.rs 🔗

@@ -118,7 +118,7 @@ pub struct Edit {
     pub new_text: String,
 }
 
-#[derive(Default, Debug, Deserialize)]
+#[derive(Clone, Default, Debug, Deserialize)]
 struct StreamingEditFileToolPartialInput {
     #[serde(default)]
     display_description: Option<String>,
@@ -132,7 +132,7 @@ struct StreamingEditFileToolPartialInput {
     edits: Option<Vec<PartialEdit>>,
 }
 
-#[derive(Default, Debug, Deserialize)]
+#[derive(Clone, Default, Debug, Deserialize)]
 pub struct PartialEdit {
     #[serde(default)]
     pub old_text: Option<String>,
@@ -314,12 +314,19 @@ impl AgentTool for StreamingEditFileTool {
     ) -> Task<Result<Self::Output, Self::Output>> {
         cx.spawn(async move |cx: &mut AsyncApp| {
             let mut state: Option<EditSession> = None;
+            let mut last_partial: Option<StreamingEditFileToolPartialInput> = None;
             loop {
                 futures::select! {
                     partial = input.recv_partial().fuse() => {
                         let Some(partial_value) = partial else { break };
                         if let Ok(parsed) = serde_json::from_value::<StreamingEditFileToolPartialInput>(partial_value) {
+                            let path_complete = parsed.path.is_some()
+                                && parsed.path.as_ref() == last_partial.as_ref().and_then(|p| p.path.as_ref());
+
+                            last_partial = Some(parsed.clone());
+
                             if state.is_none()
+                                && path_complete
                                 && let StreamingEditFileToolPartialInput {
                                     path: Some(path),
                                     display_description: Some(display_description),
@@ -1907,6 +1914,13 @@ mod tests {
         let task = cx.update(|cx| tool.clone().run(input, event_stream, cx));
 
         // Setup + single edit that stays in-progress (no second edit to prove completion)
+        sender.send_partial(json!({
+            "display_description": "Single edit",
+            "path": "root/file.txt",
+            "mode": "edit",
+        }));
+        cx.run_until_parked();
+
         sender.send_partial(json!({
             "display_description": "Single edit",
             "path": "root/file.txt",
@@ -3475,6 +3489,12 @@ mod tests {
         let task = cx.update(|cx| tool.clone().run(input, event_stream, cx));
 
         // Transition to BufferResolved
+        sender.send_partial(json!({
+            "display_description": "Overwrite file",
+            "path": "root/file.txt",
+        }));
+        cx.run_until_parked();
+
         sender.send_partial(json!({
             "display_description": "Overwrite file",
             "path": "root/file.txt",
@@ -3550,8 +3570,9 @@ mod tests {
         // Verify buffer still has old content (no content partial yet)
         let buffer = project.update(cx, |project, cx| {
             let path = project.find_project_path("root/file.txt", cx).unwrap();
-            project.get_open_buffer(&path, cx).unwrap()
+            project.open_buffer(path, cx)
         });
+        let buffer = buffer.await.unwrap();
         assert_eq!(
             buffer.read_with(cx, |b, _| b.text()),
             "old line 1\nold line 2\nold line 3\n"
@@ -3735,6 +3756,106 @@ mod tests {
         );
     }
 
+    #[gpui::test]
+    async fn test_streaming_edit_file_tool_fields_out_of_order_in_write_mode(
+        cx: &mut TestAppContext,
+    ) {
+        let (tool, _project, _action_log, _fs, _thread) =
+            setup_test(cx, json!({"file.txt": "old_content"})).await;
+        let (sender, input) = ToolInput::<StreamingEditFileToolInput>::test();
+        let (event_stream, _receiver) = ToolCallEventStream::test();
+        let task = cx.update(|cx| tool.clone().run(input, event_stream, cx));
+
+        sender.send_partial(json!({
+            "display_description": "Overwrite file",
+            "mode": "write"
+        }));
+        cx.run_until_parked();
+
+        sender.send_partial(json!({
+            "display_description": "Overwrite file",
+            "mode": "write",
+            "content": "new_content"
+        }));
+        cx.run_until_parked();
+
+        sender.send_partial(json!({
+            "display_description": "Overwrite file",
+            "mode": "write",
+            "content": "new_content",
+            "path": "root"
+        }));
+        cx.run_until_parked();
+
+        // Send final.
+        sender.send_final(json!({
+            "display_description": "Overwrite file",
+            "mode": "write",
+            "content": "new_content",
+            "path": "root/file.txt"
+        }));
+
+        let result = task.await;
+        let StreamingEditFileToolOutput::Success { new_text, .. } = result.unwrap() else {
+            panic!("expected success");
+        };
+        assert_eq!(new_text, "new_content");
+    }
+
+    #[gpui::test]
+    async fn test_streaming_edit_file_tool_fields_out_of_order_in_edit_mode(
+        cx: &mut TestAppContext,
+    ) {
+        let (tool, _project, _action_log, _fs, _thread) =
+            setup_test(cx, json!({"file.txt": "old_content"})).await;
+        let (sender, input) = ToolInput::<StreamingEditFileToolInput>::test();
+        let (event_stream, _receiver) = ToolCallEventStream::test();
+        let task = cx.update(|cx| tool.clone().run(input, event_stream, cx));
+
+        sender.send_partial(json!({
+            "display_description": "Overwrite file",
+            "mode": "edit"
+        }));
+        cx.run_until_parked();
+
+        sender.send_partial(json!({
+            "display_description": "Overwrite file",
+            "mode": "edit",
+            "edits": [{"old_text": "old_content"}]
+        }));
+        cx.run_until_parked();
+
+        sender.send_partial(json!({
+            "display_description": "Overwrite file",
+            "mode": "edit",
+            "edits": [{"old_text": "old_content", "new_text": "new_content"}]
+        }));
+        cx.run_until_parked();
+
+        sender.send_partial(json!({
+            "display_description": "Overwrite file",
+            "mode": "edit",
+            "edits": [{"old_text": "old_content", "new_text": "new_content"}],
+            "path": "root"
+        }));
+        cx.run_until_parked();
+
+        // Send final.
+        sender.send_final(json!({
+            "display_description": "Overwrite file",
+            "mode": "edit",
+            "edits": [{"old_text": "old_content", "new_text": "new_content"}],
+            "path": "root/file.txt"
+        }));
+        cx.run_until_parked();
+
+        let result = task.await;
+        let StreamingEditFileToolOutput::Success { new_text, .. } = result.unwrap() else {
+            panic!("expected success");
+        };
+        assert_eq!(new_text, "new_content");
+    }
+
     async fn setup_test_with_fs(
         cx: &mut TestAppContext,
         fs: Arc<project::FakeFs>,