Restore OpenAI tool-argument repair, keep morph_edit_file strict

Author: alexylon

src/api/openai.rs

@@ -1,6 +1,7 @@
 use super::types::*;
 use super::utils::{self, REQUEST_TIMEOUT};
 use crate::error::{Result, SofosError};
+use crate::tools::tool_name::ToolName;
 use reqwest::header::{AUTHORIZATION, CONTENT_TYPE, HeaderMap, HeaderValue};
 use serde::Deserialize;
 use serde_json::json;
@@ -181,7 +182,7 @@ impl OpenAIClient {
 
                     if let Some(tool_calls) = item.tool_calls {
                         for call in tool_calls {
-                            let input = utils::parse_tool_arguments(&call.arguments);
+                            let input = parse_openai_tool_arguments(&call.name, &call.arguments);
                             content_blocks.push(ContentBlock::ToolUse {
                                 id: call.id,
                                 name: call.name,
@@ -194,7 +195,7 @@ impl OpenAIClient {
                     if let (Some(name), Some(arguments), Some(call_id)) =
                         (item.name, item.arguments, item.call_id)
                     {
-                        let input = utils::parse_tool_arguments(&arguments);
+                        let input = parse_openai_tool_arguments(&name, &arguments);
                         content_blocks.push(ContentBlock::ToolUse {
                             id: call_id,
                             name,
@@ -244,6 +245,56 @@ impl OpenAIClient {
     }
 }
 
+/// Parse tool-call arguments emitted by OpenAI into a JSON value.
+///
+/// `morph_edit_file` uses strict parsing because its `code_edit` field carries
+/// arbitrary source code; the repair heuristics below can "succeed" on a
+/// truncated payload and silently merge corrupted code into a file. For every
+/// other tool we attempt a small repair ladder (trim, drop trailing commas,
+/// close a missing brace) and fall back to `{"raw_arguments": args}` so the
+/// model can see its own malformed output and self-correct on the next turn —
+/// rather than receiving an empty `{}` that round-trips back to OpenAI as a
+/// `function_call` with no args.
+fn parse_openai_tool_arguments(name: &str, args: &str) -> serde_json::Value {
+    if name == ToolName::MorphEditFile.as_str() {
+        return serde_json::from_str(args)
+            .unwrap_or_else(|_| serde_json::Value::Object(serde_json::Map::new()));
+    }
+
+    if let Ok(v) = serde_json::from_str::<serde_json::Value>(args) {
+        return v;
+    }
+
+    let trimmed = args.trim();
+    if trimmed != args {
+        if let Ok(v) = serde_json::from_str::<serde_json::Value>(trimmed) {
+            return v;
+        }
+    }
+
+    let fixed = trimmed.replace(",}", "}").replace(",]", "]");
+    if fixed != trimmed {
+        if let Ok(v) = serde_json::from_str::<serde_json::Value>(&fixed) {
+            return v;
+        }
+    }
+
+    if trimmed.starts_with('{') && !trimmed.ends_with('}') {
+        let braced = format!("{}}}", trimmed.trim_end_matches(','));
+        if let Ok(v) = serde_json::from_str::<serde_json::Value>(&braced) {
+            return v;
+        }
+    }
+
+    let preview_end = utils::truncate_at_char_boundary(args, 200);
+    eprintln!(
+        "  \x1b[33m⚠\x1b[0m Failed to parse tool arguments as JSON for {}: {}",
+        name,
+        &args[..preview_end]
+    );
+    json!({"raw_arguments": args})
+}
+
 fn build_response_input(request: &CreateMessageRequest) -> Vec<serde_json::Value> {
     let mut input = Vec::new();
 
@@ -432,3 +483,75 @@ struct OpenAIResponseUsage {
     #[serde(default)]
     output_tokens: Option<u32>,
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn parse_args_valid_object_round_trips() {
+        let v = parse_openai_tool_arguments("read_file", r#"{"path":"src/main.rs"}"#);
+        assert_eq!(v["path"], "src/main.rs");
+    }
+
+    #[test]
+    fn parse_args_repairs_trailing_comma() {
+        let v = parse_openai_tool_arguments("read_file", r#"{"path":"src/main.rs",}"#);
+        assert_eq!(v["path"], "src/main.rs");
+    }
+
+    #[test]
+    fn parse_args_repairs_missing_closing_brace() {
+        let v = parse_openai_tool_arguments("read_file", r#"{"path":"src/main.rs""#);
+        assert_eq!(v["path"], "src/main.rs");
+    }
+
+    #[test]
+    fn parse_args_unrepairable_falls_back_to_raw_arguments() {
+        let v = parse_openai_tool_arguments("read_file", "not json at all");
+        assert_eq!(v["raw_arguments"], "not json at all");
+    }
+
+    #[test]
+    fn parse_args_morph_edit_strict_returns_empty_object_on_failure() {
+        // Truncated code_edit must NOT be silently "repaired" — that would
+        // merge corrupted source into the user's file. Strict parse → empty
+        // object → tool dispatch reports a clear "missing parameter" error.
+        let v = parse_openai_tool_arguments(
+            ToolName::MorphEditFile.as_str(),
+            r#"{"target_filepath":"src/lib.rs","code_edit":"fn x() { let y = [1,2,"#,
+        );
+        assert!(v.is_object());
+        assert_eq!(v.as_object().unwrap().len(), 0);
+    }
+
+    #[test]
+    fn parse_args_empty_string_falls_back_to_raw_arguments() {
+        let v = parse_openai_tool_arguments("read_file", "");
+        assert_eq!(v["raw_arguments"], "");
+    }
+
+    #[test]
+    fn parse_args_whitespace_only_falls_back_to_raw_arguments() {
+        let v = parse_openai_tool_arguments("read_file", "   \n\t");
+        assert_eq!(v["raw_arguments"], "   \n\t");
+    }
+
+    #[test]
+    fn parse_args_array_root_returned_as_is() {
+        // Non-object root: dispatcher will surface a missing-parameter error,
+        // which the model can self-correct from. We just need to not panic.
+        let v = parse_openai_tool_arguments("read_file", "[1,2,3]");
+        assert!(v.is_array());
+    }
+
+    #[test]
+    fn parse_args_morph_edit_valid_round_trips() {
+        let v = parse_openai_tool_arguments(
+            ToolName::MorphEditFile.as_str(),
+            r#"{"target_filepath":"src/lib.rs","instructions":"add fn","code_edit":"fn x() {}"}"#,
+        );
+        assert_eq!(v["target_filepath"], "src/lib.rs");
+        assert_eq!(v["code_edit"], "fn x() {}");
+    }
+}