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

src/app/runtime/stream/turn.rs

@@ -7,6 +7,79 @@ use crate::dto::chat::tool::{ToolCall, ToolFunction};
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
+/// Try to repair truncated JSON by closing open strings, braces, and brackets.
+///
+/// Flow: scan character-by-character tracking string/escape state.  For
+/// every `{` or `[` seen outside a string, push onto a LIFO stack; on
+/// `}`/`]` pop the matching opener (tracking remaining depth only).
+/// At the end, if the last char was a backslash (start of an escape
+/// sequence), remove it; if inside a string, append `"`; then close
+/// every unclosed opener in reverse (LIFO) order.
+///
+/// Why: LLM responses can be cut off (max_tokens, network) mid‑JSON
+/// string, but we want tools to receive whatever arguments were already
+/// emitted so the partial work can proceed.
+///
+/// Why LIFO vs. depth counters: `{` inside `[` must be closed with `}`
+/// *before* the `]`, not after it.  Simple depth counters get the order
+/// wrong for nested heterogenous structures.
+fn repair_incomplete_json(s: &str) -> String {
+    let mut stack: Vec<char> = Vec::new();
+    let mut in_string = false;
+    let mut prev_was_backslash = false;
+    // `true` only when the very last character consumed was a bare `\`
+    // inside a string (i.e. the start of an escape that was never completed).
+    let mut ends_with_unclosed_escape = false;
+
+    for c in s.chars() {
+        if prev_was_backslash {
+            // Consume the character that was being escaped — the escape is
+            // complete, so clear the unclosed-escape flag.
+            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 {
+        // The last character is a dangling backslash that started an escape
+        // but got cut off before the escaped char — remove it.
+        result.pop();
+    }
+    if in_string {
+        result.push('"');
+    }
+    for &opener in stack.iter().rev() {
+        match opener {
+            '{' => result.push('}'),
+            '[' => result.push(']'),
+            _ => {}
+        }
+    }
+    result
+}
+
 /// Accumulates a single streaming assistant turn into its final
 /// `ChatMessage` form, including tool-call deltas and content/reasoning.
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -117,16 +190,32 @@ impl StreamedTurn {
                 .iter()
                 .filter(|tc| !tc.name.is_empty())
                 .map(|tc| {
-                    let args_value: serde_json::Value = serde_json::from_str(&tc.arguments)
-                        .unwrap_or_else(|e| {
-                            tracing::warn!(
-                                "[stream] tool call '{}' has invalid JSON arguments: {} — \
-                                 arguments will be double-stringified, which may cause \
-                                 tool execution to fail",
-                                tc.name, e,
-                            );
-                            serde_json::Value::String(tc.arguments.clone())
-                        });
+                    let args_value: serde_json::Value = match serde_json::from_str(&tc.arguments)
+                    {
+                        Ok(v) => v,
+                        Err(e) => {
+                            let repaired = repair_incomplete_json(&tc.arguments);
+                            match serde_json::from_str(&repaired) {
+                                Ok(v) => {
+                                    tracing::warn!(
+                                        "[stream] tool call '{}' had truncated JSON \
+                                         arguments — repaired successfully: {}",
+                                        tc.name, e,
+                                    );
+                                    v
+                                }
+                                Err(e2) => {
+                                    tracing::warn!(
+                                        "[stream] tool call '{}' has invalid JSON \
+                                         arguments: {} (after repair: {}) — falling \
+                                         back to raw string",
+                                        tc.name, e, e2,
+                                    );
+                                    serde_json::Value::String(tc.arguments.clone())
+                                }
+                            }
+                        }
+                    };
                     ToolCall {
                         id: tc.id.clone(),
                         type_: "function".to_string(),
@@ -212,6 +301,63 @@ mod tests {
         }
     }
 
+    #[test]
+    fn repair_closes_unclosed_string() {
+        let result = repair_incomplete_json("{\"key\": \"value");
+        assert_eq!(result, "{\"key\": \"value\"}");
+    }
+
+    #[test]
+    fn repair_closes_unclosed_object() {
+        let result = repair_incomplete_json("{\"key\": \"value\"");
+        assert_eq!(result, "{\"key\": \"value\"}");
+    }
+
+    #[test]
+    fn repair_closes_nested_structures() {
+        let result = repair_incomplete_json("{\"a\": [1, 2, {\"b\": 3");
+        assert_eq!(result, "{\"a\": [1, 2, {\"b\": 3}]}");
+    }
+
+    #[test]
+    fn repair_leaves_complete_json_unchanged() {
+        let s = "{\"a\": 1, \"b\": \"hello\"}";
+        assert_eq!(repair_incomplete_json(s), s);
+    }
+
+    #[test]
+    fn repair_handles_trailing_backslash_before_cut() {
+        // Truncated inside an escape sequence like "hello\"
+        let result = repair_incomplete_json("{\"text\": \"hello\\");
+        assert_eq!(result, "{\"text\": \"hello\"}");
+    }
+
+    #[test]
+    fn repair_handles_escaped_quotes_inside_string() {
+        // Input ends with `\"` where the `"` is the escaped character
+        // (consumed by the backslash handler), so the string is still
+        // unterminated. Repair adds `"` to close the string and `}` to
+        // close the object.
+        let result = repair_incomplete_json("{\"msg\": \"he said \\\"hello\\\"");
+        assert_eq!(result, "{\"msg\": \"he said \\\"hello\\\"\"}");
+    }
+
+    #[test]
+    fn build_assistant_message_repairs_truncated_tool_call() {
+        let mut turn = StreamedTurn::new();
+        turn.tool_calls.push(tool_call(
+            "write",
+            "{\"path\": \"a.txt\", \"content\": \"short\", \"reason\": \"trunc",
+        ));
+        let msg = turn.build_assistant_message();
+        let tcs = msg.tool_calls.expect("should produce tool calls");
+        assert_eq!(tcs.len(), 1);
+        let args = &tcs[0].function.arguments;
+        assert!(args.is_object(), "args should be an object after repair: {args:?}");
+        assert_eq!(args.get("path").and_then(|v| v.as_str()), Some("a.txt"));
+        assert_eq!(args.get("content").and_then(|v| v.as_str()), Some("short"));
+    }
+
     #[test]
     fn incomplete_tool_call_flags_truncated_json() {
         let mut turn = StreamedTurn::new();
@@ -229,10 +375,19 @@ mod tests {
 
     #[test]
     fn incomplete_tool_call_ignores_calls_without_a_name() {
-        // A slot reserved by `apply_event` (via index padding) but never
-        // filled with a name shouldn't be mistaken for a truncated call.
         let mut turn = StreamedTurn::new();
         turn.tool_calls.push(tool_call("", "not json at all"));
         assert!(turn.incomplete_tool_call().is_none());
     }
+
+    #[test]
+    fn incomplete_tool_call_accepts_repaired_json() {
+        // `incomplete_tool_call` uses raw `serde_json::from_str` (no repair)
+        // so it should still flag truncated JSON even though
+        // `build_assistant_message` will later repair it.
+        let mut turn = StreamedTurn::new();
+        turn.tool_calls.push(tool_call("write", "{\"path\": \"a.txt\", \"content\": \"unterm"));
+        // Even though it's repairable, raw parse should still fail
+        assert!(serde_json::from_str::<Value>(&turn.tool_calls[0].arguments).is_err());
+    }
 }