diff --git a/crates/g3-cli/src/filter_json.rs b/crates/g3-cli/src/filter_json.rs index 83306de..0e4650a 100644 --- a/crates/g3-cli/src/filter_json.rs +++ b/crates/g3-cli/src/filter_json.rs @@ -1,78 +1,160 @@ //! JSON tool call filtering for streaming LLM responses. //! //! This module filters out JSON tool calls from LLM output streams while preserving -//! regular text content. It uses a state machine to handle streaming chunks. +//! regular text content. It uses a simple state machine optimized for streaming. //! //! # Design //! -//! The filter detects tool calls by looking for JSON objects that start with `{"tool":` -//! at the beginning of a line. It uses brace counting to find the complete JSON object -//! and removes it from the output stream. +//! The filter uses three states: +//! - **Streaming**: Normal pass-through mode. Watches for newline + whitespace + `{` +//! - **Buffering**: Saw potential tool call start, buffering to confirm/deny +//! - **Suppressing**: Confirmed tool call, counting braces (string-aware) to find end //! -//! # Known Edge Cases -//! -//! 1. **Brace counting without string awareness in main loop**: The main filtering loop -//! counts braces without considering whether they're inside JSON strings. This can -//! cause premature exit from suppression mode if a string contains `}`. -//! -//! 2. **Tool calls not at line start**: Tool calls that don't start at the beginning -//! of a line (after optional whitespace) won't be detected. -//! -//! 3. **Streaming chunk boundaries**: If a tool call pattern is split across chunks -//! (e.g., `{"to` in one chunk and `ol":` in the next), detection may fail. +//! The key insight is that we only need to buffer a small amount (around 12 chars) +//! to confirm whether `{` starts a tool call pattern like `{"tool":`. -use regex::Regex; use std::cell::RefCell; use tracing::debug; +/// Maximum chars needed to confirm/deny a tool call pattern. +/// Pattern is: { + optional whitespace + "tool" + optional whitespace + : + optional whitespace + " +/// Realistically: `{"tool":"` = 9 chars, with whitespace maybe 15 max +const MAX_BUFFER_FOR_DETECTION: usize = 20; + // Thread-local state for tracking JSON tool call suppression thread_local! { - static JSON_TOOL_STATE: RefCell = RefCell::new(JsonToolState::new()); + static JSON_TOOL_STATE: RefCell = RefCell::new(FilterState::new()); } -/// Internal state for tracking JSON tool call filtering across streaming chunks. +/// The three possible states of the filter +#[derive(Debug, Clone, PartialEq)] +enum State { + /// Normal streaming - pass through content, watch for newline + whitespace + { + Streaming, + /// Saw potential start, buffering to confirm/deny tool pattern + Buffering, + /// Confirmed tool call, suppressing until braces balance + Suppressing, +} + +/// Internal state for the filter #[derive(Debug, Clone)] -struct JsonToolState { - /// True when actively suppressing a confirmed tool call - suppression_mode: bool, - /// True when buffering potential JSON (saw { but not yet confirmed as tool call) - potential_json_mode: bool, - /// Tracks nesting depth of braces within JSON - brace_depth: i32, +struct FilterState { + state: State, + /// Buffer for potential tool call detection (Buffering state) buffer: String, - json_start_in_buffer: Option, // Position where confirmed JSON tool call starts - content_returned_up_to: usize, // Track how much content we've already returned - potential_json_start: Option, // Where the potential JSON started + /// Brace depth for JSON tracking (Suppressing state) - string-aware + brace_depth: i32, + /// Are we inside a JSON string? (for proper brace counting) + in_string: bool, + /// Was the previous char a backslash? (for escape handling) + escape_next: bool, + /// Track if we just saw a newline (to detect line-start patterns) + at_line_start: bool, + /// Whitespace seen after newline (before potential {) + pending_whitespace: String, } -impl JsonToolState { +impl FilterState { fn new() -> Self { Self { - suppression_mode: false, - potential_json_mode: false, - brace_depth: 0, + state: State::Streaming, buffer: String::new(), - json_start_in_buffer: None, - content_returned_up_to: 0, - potential_json_start: None, + brace_depth: 0, + in_string: false, + escape_next: false, + at_line_start: true, // Start of input counts as line start + pending_whitespace: String::new(), } } fn reset(&mut self) { - self.suppression_mode = false; - self.potential_json_mode = false; - self.brace_depth = 0; + self.state = State::Streaming; self.buffer.clear(); - self.json_start_in_buffer = None; - self.content_returned_up_to = 0; - self.potential_json_start = None; + self.brace_depth = 0; + self.in_string = false; + self.escape_next = false; + self.at_line_start = true; + self.pending_whitespace.clear(); } } +/// Check if buffer matches the tool call pattern. +/// Pattern: `{` followed by optional whitespace, `"tool"`, optional whitespace, `:`, optional whitespace, `"` +/// +/// Returns: +/// - Some(true) if confirmed as tool call +/// - Some(false) if confirmed NOT a tool call +/// - None if need more data +fn check_tool_pattern(buffer: &str) -> Option { + // Must start with { + if !buffer.starts_with('{') { + return Some(false); + } + + let after_brace = &buffer[1..]; + + // Skip leading whitespace after { + let trimmed = after_brace.trim_start(); + + // Need at least `"tool":"` = 8 chars after whitespace + if trimmed.len() < 8 { + // Not enough data yet - but check for early rejection + if trimmed.starts_with('"') { + let after_quote = &trimmed[1..]; + // If we have chars after the quote, check if it starts with 't' + if !after_quote.is_empty() && !after_quote.starts_with('t') { + return Some(false); // Definitely not "tool + } + if after_quote.len() >= 2 && !after_quote.starts_with("to") { + return Some(false); + } + if after_quote.len() >= 3 && !after_quote.starts_with("too") { + return Some(false); + } + if after_quote.len() >= 4 && !after_quote.starts_with("tool") { + return Some(false); + } + } else if !trimmed.is_empty() && !trimmed.starts_with('"') { + // First non-whitespace char after { is not " - not a tool call + return Some(false); + } + return None; // Need more data + } + + // We have enough data - check the full pattern + // Must be: "tool" followed by optional whitespace, :, optional whitespace, " + if !trimmed.starts_with("\"tool\"") { + return Some(false); + } + + let after_tool = trimmed[6..].trim_start(); // 6 = len of "tool" + + if after_tool.is_empty() { + return None; // Need more data + } + + if !after_tool.starts_with(':') { + return Some(false); + } + + let after_colon = after_tool[1..].trim_start(); + + if after_colon.is_empty() { + return None; // Need more data + } + + if after_colon.starts_with('"') { + return Some(true); // Confirmed tool call! + } + + Some(false) // Has : but not followed by " +} + /// Filters JSON tool calls from streaming LLM content. /// -/// Processes content chunks and removes JSON tool calls while preserving regular text. -/// Maintains state across calls to handle tool calls spanning multiple chunks. +/// Processes content character-by-character and removes JSON tool calls +/// while preserving regular text. Maintains state across calls. /// /// # Arguments /// * `content` - A chunk of streaming content from the LLM @@ -86,383 +168,165 @@ pub fn filter_json_tool_calls(content: &str) -> String { JSON_TOOL_STATE.with(|state| { let mut state = state.borrow_mut(); - - // Add new content to buffer - state.buffer.push_str(content); - - // If we're already in suppression mode, continue brace counting - if state.suppression_mode { - // Count braces in the new content only - for ch in content.chars() { - match ch { - '{' => state.brace_depth += 1, - '}' => { - state.brace_depth -= 1; - // Exit suppression mode when all braces are closed - if state.brace_depth <= 0 { - debug!("JSON tool call completed - exiting suppression mode"); - - // Extract the complete result with JSON filtered out - let result = extract_content_without_json( - &state.buffer, - state.json_start_in_buffer.unwrap_or(0), - ); - - // Return only the part we haven't returned yet - let new_content = if result.len() > state.content_returned_up_to { - result[state.content_returned_up_to..].to_string() - } else { - String::new() - }; - - state.reset(); - return new_content; - } - } - _ => {} - } - } - - // After counting braces, if still in suppression mode, - // check if a new tool call pattern appears. This handles truncated JSON - // followed by complete JSON. - if state.suppression_mode { - let current_json_start = state.json_start_in_buffer.unwrap(); - // Don't require newline - the new JSON might be concatenated directly - let tool_call_regex = Regex::new(r#"\{\s*"tool"\s*:\s*""#).unwrap(); - - // Look for new tool call patterns after the current one - if let Some(captures) = tool_call_regex.find(&state.buffer[current_json_start + 1..]) { - let new_json_start = current_json_start + 1 + captures.start() + captures.as_str().find('{').unwrap(); - - debug!("Detected new tool call at position {} while processing incomplete one at {} - discarding old", new_json_start, current_json_start); - - // The previous JSON was incomplete/malformed - // Return content before the old JSON (if any) - let content_before_old_json = if current_json_start > state.content_returned_up_to { - state.buffer[state.content_returned_up_to..current_json_start].to_string() - } else { - String::new() - }; - - // Update state to skip the incomplete JSON and position at the new one - // We'll process the new JSON on the next call - state.content_returned_up_to = new_json_start; - state.suppression_mode = false; - state.json_start_in_buffer = None; - state.brace_depth = 0; - - return content_before_old_json; - } - } - - // Still in suppression mode, return empty string (content is being accumulated) - return String::new(); - } - - // Check if we're in potential JSON mode (saw { but waiting to confirm it's a tool call) - if state.potential_json_mode { - // Check if the buffer contains a confirmed tool call pattern - let tool_call_regex = Regex::new(r#"(?m)^\s*\{\s*"tool"\s*:\s*""#).unwrap(); - - if let Some(captures) = tool_call_regex.find(&state.buffer) { - // Confirmed! This is a tool call - enter suppression mode - let match_text = captures.as_str(); - if let Some(brace_offset) = match_text.find('{') { - let json_start = captures.start() + brace_offset; - - debug!("Confirmed JSON tool call at position {} - entering suppression mode", json_start); - - state.potential_json_mode = false; - state.suppression_mode = true; - state.brace_depth = 0; - state.json_start_in_buffer = Some(json_start); - - // Count braces from json_start to see if JSON is complete - let buffer_slice = state.buffer[json_start..].to_string(); - for ch in buffer_slice.chars() { - match ch { - '{' => state.brace_depth += 1, - '}' => { - state.brace_depth -= 1; - if state.brace_depth <= 0 { - debug!("JSON tool call completed immediately"); - let result = extract_content_without_json(&state.buffer, json_start); - let new_content = if result.len() > state.content_returned_up_to { - result[state.content_returned_up_to..].to_string() - } else { - String::new() - }; - state.reset(); - return new_content; - } - } - _ => {} - } - } - // JSON incomplete, stay in suppression mode, return nothing - return String::new(); - } - } - - // Check if we can rule out this being a tool call - // If we have enough content after the { and it doesn't match the pattern, release it - if let Some(potential_start) = state.potential_json_start { - let content_after_brace = &state.buffer[potential_start..]; - - // Rule out as a tool call if: - // 1. Closing } appears before we see the full pattern - // 2. Content clearly doesn't match the tool call pattern - // 3. Newline appears after the opening brace (tool calls should be compact) - - let has_closing_brace = content_after_brace.contains('}'); - let has_newline = content_after_brace[1..].contains('\n'); // Skip first char which is { - let long_enough = content_after_brace.len() >= 10; - - // Detect non-tool JSON patterns: - // - { followed by " and a key that doesn't start with "tool" - // - { followed by "t" but not "to" - // - { followed by "to" but not "too", etc. - let not_tool_pattern = Regex::new(r#"^\{\s*"(?:[^t]|t(?:[^o]|o(?:[^o]|o(?:[^l]|l[^"\s:]))))"#).unwrap(); - let definitely_not_tool = not_tool_pattern.is_match(content_after_brace); - - if has_closing_brace || has_newline || (long_enough && definitely_not_tool) { - debug!("Potential JSON ruled out - not a tool call"); - state.potential_json_mode = false; - state.potential_json_start = None; - - // Return the buffered content we've been holding - let new_content = if state.buffer.len() > state.content_returned_up_to { - state.buffer[state.content_returned_up_to..].to_string() - } else { - String::new() - }; - state.content_returned_up_to = state.buffer.len(); - return new_content; - } - } - - // Still in potential mode, keep buffering - return String::new(); - } - - // Detect potential JSON start: { at the beginning of a line - let potential_json_regex = Regex::new(r"(?m)^\s*\{\s*").unwrap(); + let mut output = String::new(); - if let Some(captures) = potential_json_regex.find(&state.buffer[state.content_returned_up_to..]) { - let match_start = state.content_returned_up_to + captures.start(); - let brace_pos = match_start + captures.as_str().find('{').unwrap(); - - debug!("Potential JSON detected at position {} - entering buffering mode", brace_pos); - - // Fast path: check if this is already a confirmed tool call - let tool_call_regex = Regex::new(r#"(?m)^\s*\{\s*"tool"\s*:\s*""#).unwrap(); - if tool_call_regex.is_match(&state.buffer[brace_pos..]) { - // This is a confirmed tool call! Process it immediately - let json_start = brace_pos; - debug!("Immediately confirmed tool call at position {}", json_start); - - // Return content before JSON - let content_before = if json_start > state.content_returned_up_to { - state.buffer[state.content_returned_up_to..json_start].to_string() - } else { - String::new() - }; - - state.content_returned_up_to = json_start; - state.suppression_mode = true; - state.brace_depth = 0; - state.json_start_in_buffer = Some(json_start); - - // Count braces to see if JSON is complete - let buffer_slice = state.buffer[json_start..].to_string(); - for ch in buffer_slice.chars() { - match ch { - '{' => state.brace_depth += 1, - '}' => { - state.brace_depth -= 1; - if state.brace_depth <= 0 { - debug!("JSON tool call completed in same chunk"); - let result = extract_content_without_json(&state.buffer, json_start); - let content_after = if result.len() > json_start { - &result[json_start..] - } else { - "" - }; - let final_result = format!("{}{}", content_before, content_after); - state.reset(); - return final_result; - } - } - _ => {} - } + for ch in content.chars() { + match state.state { + State::Streaming => { + handle_streaming_char(&mut state, ch, &mut output); } - // JSON incomplete, return content before and stay in suppression mode - return content_before; - } - - // Return content before the potential JSON - let content_before = if brace_pos > state.content_returned_up_to { - state.buffer[state.content_returned_up_to..brace_pos].to_string() - } else { - String::new() - }; - - state.content_returned_up_to = brace_pos; - state.potential_json_mode = true; - state.potential_json_start = Some(brace_pos); - - // Optimization: immediately check if we can rule this out for single-chunk processing - let content_after_brace = &state.buffer[brace_pos..]; - let has_closing_brace = content_after_brace.contains('}'); - let has_newline = content_after_brace.len() > 1 && content_after_brace[1..].contains('\n'); - let long_enough = content_after_brace.len() >= 10; - - let not_tool_pattern = Regex::new(r#"^\{\s*"(?:[^t]|t(?:[^o]|o(?:[^o]|o(?:[^l]|l[^"\s:]))))"#).unwrap(); - let definitely_not_tool = not_tool_pattern.is_match(content_after_brace); - - if has_closing_brace || has_newline || (long_enough && definitely_not_tool) { - debug!("Immediately ruled out as not a tool call"); - state.potential_json_mode = false; - state.potential_json_start = None; - - // Return all the buffered content - let new_content = if state.buffer.len() > state.content_returned_up_to { - state.buffer[state.content_returned_up_to..].to_string() - } else { - String::new() - }; - state.content_returned_up_to = state.buffer.len(); - return format!("{}{}", content_before, new_content); - } - - return content_before; - } - - // Check for tool call pattern using corrected regex - let tool_call_regex = Regex::new(r#"(?m)^\s*\{\s*"tool"\s*:\s*"[^"]*""#).unwrap(); - - if let Some(captures) = tool_call_regex.find(&state.buffer) { - let match_text = captures.as_str(); - - // Find the position of the opening brace in the match - if let Some(brace_offset) = match_text.find('{') { - let json_start = captures.start() + brace_offset; - - debug!( - "Detected JSON tool call at position {} - entering suppression mode", - json_start - ); - - // Return content before JSON that we haven't returned yet - let content_before_json = if json_start >= state.content_returned_up_to { - state.buffer[state.content_returned_up_to..json_start].to_string() - } else { - String::new() - }; - - state.content_returned_up_to = json_start; - - // Enter suppression mode - state.suppression_mode = true; - state.brace_depth = 0; - state.json_start_in_buffer = Some(json_start); - - // Count braces from the JSON start to see if it's complete - let buffer_clone = state.buffer.clone(); - for ch in buffer_clone[json_start..].chars() { - match ch { - '{' => state.brace_depth += 1, - '}' => { - state.brace_depth -= 1; - if state.brace_depth <= 0 { - // JSON is complete in this chunk - debug!("JSON tool call completed in same chunk"); - let result = extract_content_without_json(&buffer_clone, json_start); - - // Return content before JSON plus content after JSON - let content_after_json = if result.len() > json_start { - &result[json_start..] - } else { - "" - }; - - let final_result = - format!("{}{}", content_before_json, content_after_json); - state.reset(); - return final_result; - } - } - _ => {} - } + State::Buffering => { + handle_buffering_char(&mut state, ch, &mut output); + } + State::Suppressing => { + handle_suppressing_char(&mut state, ch, &mut output); } - - // JSON is incomplete, return only the content before JSON - return content_before_json; } } - - // No JSON tool call detected, return only the new content we haven't returned yet - if state.buffer.len() > state.content_returned_up_to { - let result = state.buffer[state.content_returned_up_to..].to_string(); - state.content_returned_up_to = state.buffer.len(); - result - } else { - String::new() - } + + output }) } -/// Extracts content from buffer, removing the JSON tool call. -/// -/// Given a buffer and the start position of a JSON tool call, this function: -/// 1. Extracts all content before the JSON -/// 2. Finds the end of the JSON (matching closing brace) -/// 3. Extracts all content after the JSON -/// 4. Returns the concatenation of before + after (JSON removed) -/// -/// # Arguments -/// * `full_content` - The full content buffer -/// * `json_start` - Position where the JSON tool call begins -fn extract_content_without_json(full_content: &str, json_start: usize) -> String { - // Find the end of the JSON using proper brace counting with string handling - let mut brace_depth = 0; - let mut json_end = json_start; - let mut in_string = false; - let mut escape_next = false; - - for (i, ch) in full_content[json_start..].char_indices() { - if escape_next { - escape_next = false; - continue; +/// Handle a character in Streaming state +fn handle_streaming_char(state: &mut FilterState, ch: char, output: &mut String) { + match ch { + '\n' => { + // Output the newline and any pending whitespace + output.push_str(&state.pending_whitespace); + output.push(ch); + state.pending_whitespace.clear(); + state.at_line_start = true; } - - match ch { - '\\' if in_string => escape_next = true, - '"' if !escape_next => in_string = !in_string, - '{' if !in_string => { - brace_depth += 1; - } - '}' if !in_string => { - brace_depth -= 1; - if brace_depth == 0 { - json_end = json_start + i + 1; // +1 to include the closing brace - break; - } - } - _ => {} + ' ' | '\t' if state.at_line_start => { + // Accumulate whitespace at line start + state.pending_whitespace.push(ch); + } + '{' if state.at_line_start => { + // Potential tool call! Enter buffering mode + debug!("Potential tool call detected - entering Buffering state"); + state.state = State::Buffering; + state.buffer.clear(); + state.buffer.push(ch); + // Don't output pending_whitespace yet - we might need to suppress it + } + _ => { + // Regular character - output any pending whitespace first + output.push_str(&state.pending_whitespace); + state.pending_whitespace.clear(); + output.push(ch); + state.at_line_start = false; } } +} - // Return content before and after the JSON (excluding the JSON itself) - let before = &full_content[..json_start]; - let after = if json_end < full_content.len() { - &full_content[json_end..] - } else { - "" - }; +/// Handle a character in Buffering state +fn handle_buffering_char(state: &mut FilterState, ch: char, output: &mut String) { + state.buffer.push(ch); + + // Check if we can determine tool call status + match check_tool_pattern(&state.buffer) { + Some(true) => { + // Confirmed tool call! Enter suppression mode + debug!("Confirmed tool call - entering Suppressing state"); + state.state = State::Suppressing; + state.brace_depth = 1; // We already have the opening { + state.in_string = true; // We're inside the "tool" value string + state.escape_next = false; + // Discard pending_whitespace (it's part of the tool call line) + state.pending_whitespace.clear(); + state.buffer.clear(); + } + Some(false) => { + // Not a tool call - release buffered content + debug!("Not a tool call - releasing buffer"); + output.push_str(&state.pending_whitespace); + output.push_str(&state.buffer); + state.pending_whitespace.clear(); + state.buffer.clear(); + state.state = State::Streaming; + state.at_line_start = ch == '\n'; + } + None => { + // Need more data - check if buffer is getting too long + if state.buffer.len() > MAX_BUFFER_FOR_DETECTION { + // Too long without confirmation - not a tool call + debug!("Buffer exceeded max length - not a tool call"); + output.push_str(&state.pending_whitespace); + output.push_str(&state.buffer); + state.pending_whitespace.clear(); + state.buffer.clear(); + state.state = State::Streaming; + state.at_line_start = false; + } + // Otherwise keep buffering + } + } +} - format!("{}{}", before, after) +/// Handle a character in Suppressing state (string-aware brace counting) +fn handle_suppressing_char(state: &mut FilterState, ch: char, _output: &mut String) { + // Track chars to detect if we see a new tool call pattern while suppressing + // This handles truncated JSON followed by complete JSON + state.buffer.push(ch); + + // Handle escape sequences + if state.escape_next { + state.escape_next = false; + return; + } + + match ch { + '\\' if state.in_string => { + state.escape_next = true; + } + '"' => { + state.in_string = !state.in_string; + } + '{' if !state.in_string => { + state.brace_depth += 1; + } + '}' if !state.in_string => { + state.brace_depth -= 1; + if state.brace_depth <= 0 { + // JSON complete! Return to streaming + debug!("Tool call complete - returning to Streaming state"); + state.state = State::Streaming; + state.at_line_start = false; // We're right after the } + state.in_string = false; + state.escape_next = false; + state.buffer.clear(); + } + } + _ => {} + } + + // Check if we're seeing a new tool call pattern (truncated JSON case) + // This can happen with or without a newline before the new { + // Look for { followed by tool pattern in the buffer + if state.buffer.len() >= 10 { + // Find the last { that could start a new tool call + for (i, c) in state.buffer.char_indices().rev() { + if c == '{' && i > 0 { + let potential_tool = &state.buffer[i..]; + if let Some(true) = check_tool_pattern(potential_tool) { + // New tool call detected! Restart suppression from here + debug!("New tool call detected while suppressing - restarting"); + state.brace_depth = 1; + state.in_string = true; + // Keep only the part after the new { for continued tracking + state.buffer = potential_tool.to_string(); + return; + } + } + } + + // Limit buffer size to prevent unbounded growth + if state.buffer.len() > 200 { + let keep_from = state.buffer.len() - 100; + state.buffer = state.buffer[keep_from..].to_string(); + } + } } /// Resets the global JSON filtering state. @@ -475,3 +339,78 @@ pub fn reset_json_tool_state() { state.reset(); }); } + + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_check_tool_pattern_confirmed() { + assert_eq!(check_tool_pattern(r#"{"tool":"" +"#), Some(true)); + assert_eq!(check_tool_pattern(r#"{"tool": "shell""#), Some(true)); + assert_eq!(check_tool_pattern(r#"{ "tool" : "test""#), Some(true)); + } + + #[test] + fn test_check_tool_pattern_rejected() { + assert_eq!(check_tool_pattern(r#"{"other": "value"}"#), Some(false)); + assert_eq!(check_tool_pattern(r#"{"tools": "value"}"#), Some(false)); + assert_eq!(check_tool_pattern(r#"{"tool": 123}"#), Some(false)); // number not string + } + + #[test] + fn test_check_tool_pattern_need_more() { + assert_eq!(check_tool_pattern(r#"{"#), None); + assert_eq!(check_tool_pattern(r#"{"tool"#), None); + assert_eq!(check_tool_pattern(r#"{"tool":"#), None); + } + + #[test] + fn test_passthrough_no_tool() { + reset_json_tool_state(); + let input = "Hello world"; + assert_eq!(filter_json_tool_calls(input), input); + } + + #[test] + fn test_simple_tool_filtered() { + reset_json_tool_state(); + let input = "Before\n{\"tool\": \"shell\", \"args\": {}}\nAfter"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Before\n\nAfter"); + } + + #[test] + fn test_tool_with_braces_in_string() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": \"shell\", \"args\": {\"cmd\": \"echo }\"}}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); + } + + #[test] + fn test_non_tool_json_passes_through() { + reset_json_tool_state(); + let input = "Text\n{\"other\": \"value\"}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, input); + } + + #[test] + fn test_streaming_chunks() { + reset_json_tool_state(); + let chunks = vec![ + "Before\n", + "{\"tool\": \"", + "shell\", \"args\": {}", + "}\nAfter", + ]; + let mut result = String::new(); + for chunk in chunks { + result.push_str(&filter_json_tool_calls(chunk)); + } + assert_eq!(result, "Before\n\nAfter"); + } +} diff --git a/crates/g3-cli/tests/filter_json_stress_test.rs b/crates/g3-cli/tests/filter_json_stress_test.rs new file mode 100644 index 0000000..2fc8d7a --- /dev/null +++ b/crates/g3-cli/tests/filter_json_stress_test.rs @@ -0,0 +1,639 @@ +//! Stress tests for JSON tool call filtering. +//! +//! These tests hammer the filter with malformed JSON, partial tool calls, +//! edge cases, and adversarial inputs to ensure robustness. + +use g3_cli::filter_json::{filter_json_tool_calls, reset_json_tool_state}; + +// ============================================================================ +// Malformed JSON Tests +// ============================================================================ + +#[test] +fn test_unclosed_brace_at_end() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": \"shell\", \"args\": {\"cmd\": \"ls\""; + let result = filter_json_tool_calls(input); + // Should suppress the incomplete tool call + assert_eq!(result, "Text\n"); +} + +#[test] +fn test_missing_closing_quote() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": \"shell\", \"args\": {\"cmd\": \"ls}}\nMore"; + let result = filter_json_tool_calls(input); + // The unbalanced quote makes brace counting tricky + // Should still filter the tool call attempt + assert_eq!(result, "Text\n"); +} + +#[test] +fn test_extra_closing_braces() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": \"shell\", \"args\": {}}}}}\nMore"; + let result = filter_json_tool_calls(input); + // Extra braces after valid JSON should pass through + assert_eq!(result, "Text\n}}}\nMore"); +} + +#[test] +fn test_deeply_nested_malformed() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": \"x\", \"args\": {{{{{{}}}}}}}\nMore"; + let result = filter_json_tool_calls(input); + // Should handle deep nesting - extra braces get consumed as part of the tool call + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_null_bytes_in_json() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": \"shell\0\", \"args\": {}}\nMore"; + let result = filter_json_tool_calls(input); + // Should handle null bytes gracefully + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_unicode_in_tool_name() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": \"shëll\", \"args\": {}}\nMore"; + let result = filter_json_tool_calls(input); + // Unicode in tool name - still a valid tool call pattern + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_emoji_in_args() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": \"shell\", \"args\": {\"msg\": \"Hello 🎉\"}}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_very_long_string_value() { + reset_json_tool_state(); + let long_string = "x".repeat(10000); + let input = format!("Text\n{{\"tool\": \"shell\", \"args\": {{\"data\": \"{}\"}}}}\nMore", long_string); + let result = filter_json_tool_calls(&input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_many_escaped_quotes() { + reset_json_tool_state(); + let input = r#"Text +{"tool": "shell", "args": {"cmd": "echo \"a\" \"b\" \"c\" \"d\" \"e\""}} +More"#; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_escaped_backslash_before_quote() { + reset_json_tool_state(); + // This is: {"tool": "shell", "args": {"path": "C:\\"}} + let input = "Text\n{\"tool\": \"shell\", \"args\": {\"path\": \"C:\\\\\"}}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_newlines_inside_string() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": \"shell\", \"args\": {\"cmd\": \"echo\\nworld\"}}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +// ============================================================================ +// Partial Tool Call Pattern Tests +// ============================================================================ + +#[test] +fn test_just_opening_brace() { + reset_json_tool_state(); + let result = filter_json_tool_calls("Text\n{"); + // Should buffer, waiting for more + assert_eq!(result, "Text\n"); + + // Now send something that's not a tool call + let result2 = filter_json_tool_calls("\"other\": 1}\nMore"); + assert_eq!(result2, "{\"other\": 1}\nMore"); +} + +#[test] +fn test_partial_tool_keyword() { + reset_json_tool_state(); + let chunks = vec!["Text\n{", "\"to", "ol", "\": ", "\"sh", "ell\"", ", \"args\": {}", "}\nMore"]; + let mut result = String::new(); + for chunk in chunks { + result.push_str(&filter_json_tool_calls(chunk)); + } + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_tool_then_not_colon() { + reset_json_tool_state(); + let input = "Text\n{\"tool\" \"shell\"}\nMore"; // Missing colon + let result = filter_json_tool_calls(input); + // Not a valid tool call pattern - should pass through + assert_eq!(result, input); +} + +#[test] +fn test_tool_colon_then_number() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": 123}\nMore"; // Number instead of string + let result = filter_json_tool_calls(input); + // Not a valid tool call pattern - should pass through + assert_eq!(result, input); +} + +#[test] +fn test_tool_colon_then_null() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": null}\nMore"; + let result = filter_json_tool_calls(input); + // Not a valid tool call pattern - should pass through + assert_eq!(result, input); +} + +#[test] +fn test_tool_colon_then_array() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": []}\nMore"; + let result = filter_json_tool_calls(input); + // Not a valid tool call pattern - should pass through + assert_eq!(result, input); +} + +#[test] +fn test_tool_colon_then_object() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": {}}\nMore"; + let result = filter_json_tool_calls(input); + // Not a valid tool call pattern - should pass through + assert_eq!(result, input); +} + +#[test] +fn test_tools_plural() { + reset_json_tool_state(); + let input = "Text\n{\"tools\": \"shell\"}\nMore"; + let result = filter_json_tool_calls(input); + // "tools" is not "tool" - should pass through + assert_eq!(result, input); +} + +#[test] +fn test_tool_with_prefix() { + reset_json_tool_state(); + let input = "Text\n{\"mytool\": \"shell\"}\nMore"; + let result = filter_json_tool_calls(input); + // "mytool" is not "tool" - should pass through + assert_eq!(result, input); +} + +#[test] +fn test_tool_uppercase() { + reset_json_tool_state(); + let input = "Text\n{\"TOOL\": \"shell\"}\nMore"; + let result = filter_json_tool_calls(input); + // "TOOL" is not "tool" - should pass through + assert_eq!(result, input); +} + +// ============================================================================ +// Streaming Edge Cases +// ============================================================================ + +#[test] +fn test_single_char_streaming() { + reset_json_tool_state(); + let input = "Hi\n{\"tool\": \"x\", \"args\": {}}\nBye"; + let mut result = String::new(); + for ch in input.chars() { + result.push_str(&filter_json_tool_calls(&ch.to_string())); + } + assert_eq!(result, "Hi\n\nBye"); +} + +#[test] +fn test_two_char_streaming() { + reset_json_tool_state(); + let input = "Hi\n{\"tool\": \"x\", \"args\": {}}\nBye"; + let mut result = String::new(); + let chars: Vec = input.chars().collect(); + for chunk in chars.chunks(2) { + let s: String = chunk.iter().collect(); + result.push_str(&filter_json_tool_calls(&s)); + } + assert_eq!(result, "Hi\n\nBye"); +} + +#[test] +fn test_random_chunk_sizes() { + reset_json_tool_state(); + let input = "Before\n{\"tool\": \"shell\", \"args\": {\"cmd\": \"ls -la\"}}\nAfter"; + + // Chunk at various sizes + let chunk_sizes = [1, 3, 7, 11, 13, 17]; + + for &size in &chunk_sizes { + reset_json_tool_state(); + let mut result = String::new(); + let mut pos = 0; + while pos < input.len() { + let end = (pos + size).min(input.len()); + let chunk = &input[pos..end]; + result.push_str(&filter_json_tool_calls(chunk)); + pos = end; + } + assert_eq!(result, "Before\n\nAfter", "Failed with chunk size {}", size); + } +} + +#[test] +fn test_chunk_boundary_at_brace() { + reset_json_tool_state(); + let chunks = vec!["Text\n", "{", "\"tool\": \"x\", \"args\": {}", "}", "\nMore"]; + let mut result = String::new(); + for chunk in chunks { + result.push_str(&filter_json_tool_calls(chunk)); + } + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_chunk_boundary_at_quote() { + reset_json_tool_state(); + let chunks = vec!["Text\n{\"tool\": \"", "shell", "\", \"args\": {}}", "\nMore"]; + let mut result = String::new(); + for chunk in chunks { + result.push_str(&filter_json_tool_calls(chunk)); + } + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_chunk_boundary_at_colon() { + reset_json_tool_state(); + let chunks = vec!["Text\n{\"tool\"", ":", " \"shell\", \"args\": {}}\nMore"]; + let mut result = String::new(); + for chunk in chunks { + result.push_str(&filter_json_tool_calls(chunk)); + } + assert_eq!(result, "Text\n\nMore"); +} + +// ============================================================================ +// Multiple Tool Calls +// ============================================================================ + +#[test] +fn test_two_tool_calls_same_line() { + reset_json_tool_state(); + // Two tool calls on same line (no newline between) + let input = "Text\n{\"tool\": \"a\", \"args\": {}}{\"tool\": \"b\", \"args\": {}}\nMore"; + let result = filter_json_tool_calls(input); + // First is filtered (starts at line beginning) + // Second starts immediately after first's }, not at line start, so passes through + // This is acceptable - LLMs typically put tool calls on separate lines + assert_eq!(result, "Text\n{\"tool\": \"b\", \"args\": {}}\nMore"); +} + +#[test] +fn test_three_tool_calls_separate_lines() { + reset_json_tool_state(); + let input = "A\n{\"tool\": \"x\", \"args\": {}}\nB\n{\"tool\": \"y\", \"args\": {}}\nC\n{\"tool\": \"z\", \"args\": {}}\nD"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "A\n\nB\n\nC\n\nD"); +} + +#[test] +fn test_tool_call_then_regular_json() { + reset_json_tool_state(); + let input = "A\n{\"tool\": \"x\", \"args\": {}}\nB\n{\"data\": 123}\nC"; + let result = filter_json_tool_calls(input); + // First is tool call (filtered), second is regular JSON (kept) + assert_eq!(result, "A\n\nB\n{\"data\": 123}\nC"); +} + +#[test] +fn test_regular_json_then_tool_call() { + reset_json_tool_state(); + let input = "A\n{\"data\": 123}\nB\n{\"tool\": \"x\", \"args\": {}}\nC"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "A\n{\"data\": 123}\nB\n\nC"); +} + +// ============================================================================ +// Adversarial Inputs +// ============================================================================ + +#[test] +fn test_fake_tool_in_string() { + reset_json_tool_state(); + // The tool pattern appears inside a string value + let input = r#"Text +{"message": "{\"tool\": \"shell\"}"} +More"#; + let result = filter_json_tool_calls(input); + // Should pass through - the pattern is inside a string + assert_eq!(result, input); +} + +#[test] +fn test_nested_json_with_tool_key() { + reset_json_tool_state(); + // Nested object has "tool" key but outer doesn't match pattern + let input = "Text\n{\"outer\": {\"tool\": \"inner\"}}\nMore"; + let result = filter_json_tool_calls(input); + // Should pass through - outer object doesn't start with "tool" + assert_eq!(result, input); +} + +#[test] +fn test_brace_bomb() { + reset_json_tool_state(); + // Many braces to stress the counter + let input = "Text\n{\"tool\": \"x\", \"args\": {\"a\": {\"b\": {\"c\": {\"d\": {\"e\": {}}}}}}}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_string_with_many_braces() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": \"x\", \"args\": {\"code\": \"{{{{}}}}\"}}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_alternating_braces_in_string() { + reset_json_tool_state(); + let input = "Text\n{\"tool\": \"x\", \"args\": {\"pat\": \"}{}{}{\"}}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_quote_after_backslash_in_string() { + reset_json_tool_state(); + // Tricky: \" inside string should not end the string + let input = r#"Text +{"tool": "x", "args": {"msg": "say \"hi\""}} +More"#; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_double_backslash_then_quote() { + reset_json_tool_state(); + // \\ followed by " - the quote DOES end the string + let input = "Text\n{\"tool\": \"x\", \"args\": {\"path\": \"C:\\\\\"}}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_triple_backslash_then_quote() { + reset_json_tool_state(); + // \\\" - escaped backslash followed by escaped quote + let input = "Text\n{\"tool\": \"x\", \"args\": {\"s\": \"a\\\\\\\"b\"}}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +// ============================================================================ +// Whitespace Variations +// ============================================================================ + +#[test] +fn test_tabs_before_brace() { + reset_json_tool_state(); + let input = "Text\n\t\t{\"tool\": \"x\", \"args\": {}}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_spaces_before_brace() { + reset_json_tool_state(); + let input = "Text\n {\"tool\": \"x\", \"args\": {}}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_mixed_whitespace_before_brace() { + reset_json_tool_state(); + let input = "Text\n \t \t {\"tool\": \"x\", \"args\": {}}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_space_after_opening_brace() { + reset_json_tool_state(); + let input = "Text\n{ \"tool\": \"x\", \"args\": {}}\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_lots_of_space_in_json() { + reset_json_tool_state(); + let input = "Text\n{ \"tool\" : \"x\" , \"args\" : { } }\nMore"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Text\n\nMore"); +} + +#[test] +fn test_crlf_line_endings() { + reset_json_tool_state(); + let input = "Text\r\n{\"tool\": \"x\", \"args\": {}}\r\nMore"; + let result = filter_json_tool_calls(input); + // \r is not treated as line start, so { after \r\n should work + // Actually \n triggers line start, \r is just a regular char + assert_eq!(result, "Text\r\n\r\nMore"); +} + +// ============================================================================ +// Empty and Minimal Cases +// ============================================================================ + +#[test] +fn test_empty_input() { + reset_json_tool_state(); + assert_eq!(filter_json_tool_calls(""), ""); +} + +#[test] +fn test_just_newline() { + reset_json_tool_state(); + assert_eq!(filter_json_tool_calls("\n"), "\n"); +} + +#[test] +fn test_just_brace() { + reset_json_tool_state(); + let r1 = filter_json_tool_calls("{"); + // At start of input (line start), { triggers buffering + assert_eq!(r1, ""); + + // Send non-tool content - the newline comes through + let r2 = filter_json_tool_calls("}\n"); + assert_eq!(r2, "{}\n"); +} + +#[test] +fn test_minimal_tool_call() { + reset_json_tool_state(); + let input = "{\"tool\":\"x\",\"args\":{}}"; + let result = filter_json_tool_calls(input); + assert_eq!(result, ""); +} + +#[test] +fn test_tool_call_at_very_start() { + reset_json_tool_state(); + let input = "{\"tool\": \"x\", \"args\": {}}\nAfter"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "\nAfter"); +} + +// ============================================================================ +// State Reset Tests +// ============================================================================ + +#[test] +fn test_reset_clears_buffering_state() { + reset_json_tool_state(); + + // Start a potential tool call + let _ = filter_json_tool_calls("Text\n{"); + + // Reset + reset_json_tool_state(); + + // New input should work fresh + let result = filter_json_tool_calls("Fresh start"); + assert_eq!(result, "Fresh start"); +} + +#[test] +fn test_reset_clears_suppressing_state() { + reset_json_tool_state(); + + // Start suppressing a tool call + let _ = filter_json_tool_calls("Text\n{\"tool\": \"x\", \"args\": {"); + + // Reset + reset_json_tool_state(); + + // New input should work fresh + let result = filter_json_tool_calls("Fresh start"); + assert_eq!(result, "Fresh start"); +} + +// ============================================================================ +// Real-World Patterns from Bug Reports +// ============================================================================ + +#[test] +fn test_str_replace_with_diff() { + reset_json_tool_state(); + let input = r#"I'll update the file: +{"tool": "str_replace", "args": {"file_path": "src/main.rs", "diff": "@@ -1,3 +1,4 @@\n fn main() {\n+ println!(\"Hello\");\n }"}} +Done!"#; + let result = filter_json_tool_calls(input); + assert_eq!(result, "I'll update the file:\n\nDone!"); +} + +#[test] +fn test_shell_with_complex_command() { + reset_json_tool_state(); + let input = r#"Running command: +{"tool": "shell", "args": {"command": "find . -name '*.rs' -exec grep -l 'TODO' {} \;"}} +Results above."#; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Running command:\n\nResults above."); +} + +#[test] +fn test_write_file_with_json_content() { + reset_json_tool_state(); + let input = r#"Creating config: +{"tool": "write_file", "args": {"file_path": "config.json", "content": "{\"key\": \"value\"}"}} +File created."#; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Creating config:\n\nFile created."); +} + +#[test] +fn test_read_file_simple() { + reset_json_tool_state(); + let input = "Let me check:\n{\"tool\": \"read_file\", \"args\": {\"file_path\": \"README.md\"}}\nHere's what I found:"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Let me check:\n\nHere's what I found:"); +} + +#[test] +fn test_final_output() { + reset_json_tool_state(); + let input = "Task complete.\n{\"tool\": \"final_output\", \"args\": {\"summary\": \"# Summary\\n\\nI completed the task.\\n\\n## Details\\n- Item 1\\n- Item 2\"}}\n"; + let result = filter_json_tool_calls(input); + assert_eq!(result, "Task complete.\n\n"); +} + +// ============================================================================ +// Truncated JSON followed by Complete JSON (the original bug) +// ============================================================================ + +#[test] +fn test_truncated_then_complete_streaming() { + reset_json_tool_state(); + + // Chunk 1: text + let r1 = filter_json_tool_calls("Some text\n"); + assert_eq!(r1, "Some text\n"); + + // Chunk 2: truncated tool call + let r2 = filter_json_tool_calls(r#"{"tool": "str_replace", "args": {"diff":"partial"#); + assert_eq!(r2, ""); + + // Chunk 3: new complete tool call (LLM retry) + let r3 = filter_json_tool_calls(r#"{"tool": "str_replace", "args": {"diff":"complete", "file_path":"x.rs"}}"#); + assert_eq!(r3, ""); + + // Chunk 4: text after + let r4 = filter_json_tool_calls("\nMore text"); + assert_eq!(r4, "\nMore text"); +} + +#[test] +fn test_multiple_truncated_then_complete() { + reset_json_tool_state(); + + let chunks = vec![ + "Start\n", + r#"{"tool": "a", "args": {"x": "trunc"#, // truncated + r#"{"tool": "b", "args": {"y": "also_trunc"#, // another truncated + r#"{"tool": "c", "args": {"z": "complete"}}"#, // finally complete + "\nEnd", + ]; + + let mut result = String::new(); + for chunk in chunks { + result.push_str(&filter_json_tool_calls(chunk)); + } + + assert_eq!(result, "Start\n\nEnd"); +}