From e13f04083313f3544bdb1a76b5ecd535ecf59e4f Mon Sep 17 00:00:00 2001 From: asepharyana Date: Wed, 15 Jul 2026 01:42:15 +0700 Subject: [PATCH] feat(runtime): implement JSON repair function for truncated tool-call arguments --- src/app/runtime/stream/turn.rs | 179 +++++++++++++++++++++++++++--- src/dto/chat/tool.rs | 191 +++++++++++++++++++++++++++++++-- src/model/app_config.rs | 21 ++++ 3 files changed, 371 insertions(+), 20 deletions(-) diff --git a/src/app/runtime/stream/turn.rs b/src/app/runtime/stream/turn.rs index 47994ee..ffad389 100644 --- a/src/app/runtime/stream/turn.rs +++ b/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 = 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::(&turn.tool_calls[0].arguments).is_err()); + } } diff --git a/src/dto/chat/tool.rs b/src/dto/chat/tool.rs index 62076c6..0d7d7b7 100644 --- a/src/dto/chat/tool.rs +++ b/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 = 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::(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::(), - ); - // 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::(&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::(), + ); + serde_json::json!({"_raw": s, "_parse_error": e.to_string()}) + } + } } } } diff --git a/src/model/app_config.rs b/src/model/app_config.rs index 101f850..fc3fd78 100644 --- a/src/model/app_config.rs +++ b/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"); + } +}