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());
+    }
 }

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()})
+                        }
+                    }
                 }
             }
         }

src/model/app_config.rs

@@ -182,3 +182,24 @@ fn claude_credentials_from_env() -> Option<(String, String)> {
     let key = std::env::var("ANTHROPIC_API_KEY").ok()?;
     Some((base_url, key))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn claude_credentials_from_env_resolves_real_env_vars() {
+        // In the test runner's environment ANTHROPIC_BASE_URL and
+        // ANTHROPIC_API_KEY may or may not be set — we only verify that
+        // the function returns Some(..) when both are present.
+        let (b, k) = match claude_credentials_from_env() {
+            Some(v) => v,
+            None => {
+                // Not an error: CI / local without the vars.
+                return;
+            }
+        };
+        assert!(!b.is_empty(), "ANTHROPIC_BASE_URL must not be empty");
+        assert!(!k.is_empty(), "ANTHROPIC_API_KEY must not be empty");
+    }
+}