agent: Allow model to respond with stringified array when it edits (#52303)

Bennet Bo Fenner and Gaauwe Rombouts created 2 weeks ago

## Context

Sometimes the model likes to respond with `"[...]"` for the edits array,
instead of just `[...]`, which causes deserialisation to fail:
```
Failed to receive tool input: invalid type: string "[{"old_text": ..., "new_text": ...}]", expected a sequence
```

## How to Review

<!-- Help reviewers focus their attention:
- For small PRs: note what to focus on (e.g., "error handling in
foo.rs")
- For large PRs (>400 LOC): provide a guided tour — numbered list of
files/commits to read in order. (The `large-pr` label is applied
automatically.)
     - See the review process guidelines for comment conventions -->

## Self-Review Checklist

<!-- Check before requesting review: -->
- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content is consistent with the [UI/UX
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Release Notes:

- N/A

---------

Co-authored-by: Gaauwe Rombouts <[e-mail]>

Change summary

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

Detailed changes

crates/agent/src/tools/streaming_edit_file_tool.rs

@@ -22,7 +22,10 @@ use language_model::LanguageModelToolResultContent;
 use project::lsp_store::{FormatTrigger, LspFormatTarget};
 use project::{AgentLocation, Project, ProjectPath};
 use schemars::JsonSchema;
-use serde::{Deserialize, Serialize};
+use serde::{
+    Deserialize, Deserializer, Serialize,
+    de::{DeserializeOwned, Error as _},
+};
 use std::ops::Range;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -89,7 +92,11 @@ pub struct StreamingEditFileToolInput {
 
     /// List of edit operations to apply sequentially (required for 'edit' mode).
     /// Each edit finds `old_text` in the file and replaces it with `new_text`.
-    #[serde(default, skip_serializing_if = "Option::is_none")]
+    #[serde(
+        default,
+        skip_serializing_if = "Option::is_none",
+        deserialize_with = "deserialize_optional_vec_or_json_string"
+    )]
     pub edits: Option<Vec<Edit>>,
 }
 
@@ -128,7 +135,7 @@ struct StreamingEditFileToolPartialInput {
     mode: Option<StreamingEditFileMode>,
     #[serde(default)]
     content: Option<String>,
-    #[serde(default)]
+    #[serde(default, deserialize_with = "deserialize_optional_vec_or_json_string")]
     edits: Option<Vec<PartialEdit>>,
 }
 
@@ -140,6 +147,33 @@ pub struct PartialEdit {
     pub new_text: Option<String>,
 }
 
+/// Sometimes the model responds with a stringified JSON array of edits (`"[...]"`) instead of a regular array (`[...]`)
+fn deserialize_optional_vec_or_json_string<'de, T, D>(
+    deserializer: D,
+) -> Result<Option<Vec<T>>, D::Error>
+where
+    T: DeserializeOwned,
+    D: Deserializer<'de>,
+{
+    #[derive(Deserialize)]
+    #[serde(untagged)]
+    enum VecOrJsonString<T> {
+        Vec(Vec<T>),
+        String(String),
+    }
+
+    let value = Option::<VecOrJsonString<T>>::deserialize(deserializer)?;
+    match value {
+        None => Ok(None),
+        Some(VecOrJsonString::Vec(items)) => Ok(Some(items)),
+        Some(VecOrJsonString::String(string)) => serde_json::from_str::<Vec<T>>(&string)
+            .map(Some)
+            .map_err(|error| {
+                D::Error::custom(format!("failed to parse stringified edits array: {error}"))
+            }),
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(untagged)]
 pub enum StreamingEditFileToolOutput {
@@ -3673,6 +3707,35 @@ mod tests {
         assert_eq!(new_text, "HELLO\nWORLD\nfoo\n");
     }
 
+    #[gpui::test]
+    async fn test_streaming_final_input_stringified_edits_succeeds(cx: &mut TestAppContext) {
+        let (tool, _project, _action_log, _fs, _thread) =
+            setup_test(cx, json!({"file.txt": "hello\nworld\n"})).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": "Edit",
+            "path": "root/file.txt",
+            "mode": "edit"
+        }));
+        cx.run_until_parked();
+
+        sender.send_final(json!({
+            "display_description": "Edit",
+            "path": "root/file.txt",
+            "mode": "edit",
+            "edits": "[{\"old_text\": \"hello\\nworld\", \"new_text\": \"HELLO\\nWORLD\"}]"
+        }));
+
+        let result = task.await;
+        let StreamingEditFileToolOutput::Success { new_text, .. } = result.unwrap() else {
+            panic!("expected success");
+        };
+        assert_eq!(new_text, "HELLO\nWORLD\n");
+    }
+
     // Verifies that after streaming_edit_file_tool edits a file, the action log
     // reports changed buffers so that the Accept All / Reject All review UI appears.
     #[gpui::test]