feat(runtime): implement JSON repair function for truncated tool-call arguments

src/dto/chat/tool.rs

@@ -11,6 +11,95 @@
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn repair_json_closes_string() {
+        assert_eq!(repair_json("{\"a\": \"bc"), "{\"a\": \"bc\"}");
+    }
+
+    #[test]
+    fn repair_json_closes_brace() {
+        assert_eq!(repair_json("{\"a\": 1"), "{\"a\": 1}");
+    }
+
+    #[test]
+    fn repair_json_closes_bracket() {
+        assert_eq!(repair_json("{\"a\": [1, 2"), "{\"a\": [1, 2]}");
+    }
+
+    #[test]
+    fn repair_json_nested() {
+        assert_eq!(
+            repair_json("{\"a\": {\"b\": [1, 2"),
+            "{\"a\": {\"b\": [1, 2]}}"
+        );
+    }
+
+    #[test]
+    fn repair_json_bracket_then_brace() {
+        // `[` opened first → `]` must close first, then `}`
+        assert_eq!(
+            repair_json("[[1, 2, {\"a\": 3"),
+            "[[1, 2, {\"a\": 3}]]"
+        );
+    }
+
+    #[test]
+    fn repair_json_handles_escape() {
+        assert_eq!(repair_json("{\"a\": \"hello\\"), "{\"a\": \"hello\"}");
+    }
+
+    #[test]
+    fn repair_json_handles_escaped_quote() {
+        assert_eq!(
+            repair_json("{\"a\": \"he said \\\"hi\\\""),
+            "{\"a\": \"he said \\\"hi\\\"\"}"
+        );
+    }
+
+    #[test]
+    fn repair_json_handles_nested_brackets_and_braces() {
+        assert_eq!(
+            repair_json("{\"a\": [1, {\"b\": 2"),
+            "{\"a\": [1, {\"b\": 2}]}"
+        );
+    }
+
+    #[test]
+    fn repair_json_unchanged_for_valid() {
+        let v = "{\"a\": 1, \"b\": [2, 3]}";
+        assert_eq!(repair_json(v), v);
+    }
+
+    #[test]
+    fn sanitize_repairs_truncated_string() {
+        let args = Value::String("{\"path\": \"a.txt\", \"content\": \"short\"}".to_string());
+        let result = sanitize_tool_arguments(&args);
+        assert!(result.is_object());
+        assert_eq!(result.get("path").and_then(|v| v.as_str()), Some("a.txt"));
+    }
+
+    #[test]
+    fn sanitize_passes_object_through() {
+        let args = serde_json::json!({"path": "a.txt"});
+        let result = sanitize_tool_arguments(&args);
+        assert_eq!(result, args);
+    }
+
+    #[test]
+    fn sanitize_falls_back_to_raw_on_unrepairable() {
+        // Completely garbage — not even close to JSON
+        let args = Value::String("not even close".to_string());
+        let result = sanitize_tool_arguments(&args);
+        assert!(result.is_object());
+        assert!(result.get("_raw").is_some());
+        assert!(result.get("_parse_error").is_some());
+    }
+}
+
 /// A single tool-call request emitted by the model in an assistant message.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ToolCall {
@@ -27,6 +116,80 @@ pub struct ToolFunction {
     pub arguments: Value,
 }
 
+/// Normalize tool-call arguments into a JSON object/value.
+///
+/// Flow: some providers send `arguments` as a JSON-encoded string rather
+/// than a nested object; if `args` is a string, attempt to parse it as
+/// JSON. Objects and other value types pass through unchanged.
+///
+/// Security: on parse failure we wrap the raw string in `{ "_raw": "..." }`
+/// instead of passing it through as a raw string, so tools that expect a
+/// JSON object (via `args.get("key")`) get `None` rather than unexpectedly
+/// receiving a plain string value.
+///
+/// Return: the parsed `Value`, or a wrapper object on parse failure.
+/// Attempt to fix truncated JSON by closing open strings, braces and brackets.
+///
+/// Flow: single-pass character scan tracking string/escape state with a
+/// LIFO stack for `{`/`[` → append missing `"`, `]`, `}` in the right
+/// (reverse nesting) order.
+///
+/// Why: LLM output can be cut off mid‑JSON (max_tokens hit, connection
+/// drop).  This gives tools a chance to act on whatever was emitted.
+///
+/// Why LIFO vs. depth counters: `{` inside `[` must close with `}` before
+/// `]`. Simple depth counters get the nesting order wrong.
+fn repair_json(s: &str) -> String {
+    let mut stack: Vec<char> = Vec::new();
+    let mut in_string = false;
+    let mut prev_was_backslash = false;
+    let mut ends_with_unclosed_escape = false;
+
+    for c in s.chars() {
+        if prev_was_backslash {
+            prev_was_backslash = false;
+            ends_with_unclosed_escape = false;
+            continue;
+        }
+        if c == '\\' && in_string {
+            prev_was_backslash = true;
+            ends_with_unclosed_escape = true;
+            continue;
+        }
+        ends_with_unclosed_escape = false;
+        if c == '"' {
+            in_string = !in_string;
+            continue;
+        }
+        if in_string {
+            continue;
+        }
+        match c {
+            '{' | '[' => stack.push(c),
+            '}' | ']' => {
+                stack.pop();
+            }
+            _ => {}
+        }
+    }
+
+    let mut result = s.to_string();
+    if ends_with_unclosed_escape {
+        result.pop();
+    }
+    if in_string {
+        result.push('"');
+    }
+    for &opener in stack.iter().rev() {
+        match opener {
+            '{' => result.push('}'),
+            '[' => result.push(']'),
+            _ => {}
+        }
+    }
+    result
+}
+
 /// Normalize tool-call arguments into a JSON object/value.
 ///
 /// Flow: some providers send `arguments` as a JSON-encoded string rather
@@ -45,14 +208,26 @@ pub fn sanitize_tool_arguments(args: &Value) -> Value {
             match serde_json::from_str::<Value>(s) {
                 Ok(v) => v,
                 Err(e) => {
-                    tracing::error!(
-                        "tool argument is a JSON string but failed to parse: {}. \
-                         Wrapping in object to prevent tool misbehaviour. Raw was: {}",
-                        e, s.chars().take(200).collect::<String>(),
-                    );
-                    // Wrap in a safe object so tools don't receive a raw
-                    // string that could be misinterpreted as an object key.
-                    serde_json::json!({"_raw": s, "_parse_error": e.to_string()})
+                    // Try to repair truncated JSON before giving up.
+                    let repaired = repair_json(s);
+                    match serde_json::from_str::<Value>(&repaired) {
+                        Ok(v) => {
+                            tracing::warn!(
+                                "tool argument string was truncated — repaired \
+                                 successfully: {}",
+                                e,
+                            );
+                            v
+                        }
+                        Err(e2) => {
+                            tracing::error!(
+                                "tool argument is a JSON string but failed to parse: {} \
+                                 (after repair: {}). Wrapping in object. Raw (first 200): {}",
+                                e, e2, s.chars().take(200).collect::<String>(),
+                            );
+                            serde_json::json!({"_raw": s, "_parse_error": e.to_string()})
+                        }
+                    }
                 }
             }
         }