[fix] Fixes issues with Tool Input parsing.

* Ocassionally the model will generate a tool use which parameters are not a valid json. When this happens it corrupts the conversation history.
* Here we first avoid storing the tool use and add the propert validation logic to the conversation history.

Author: kennvene-aws

crates/chat-cli/src/cli/chat/consts.rs

@@ -13,6 +13,9 @@ pub const MAX_USER_MESSAGE_SIZE: usize = 400_000;
 
 pub const DUMMY_TOOL_NAME: &str = "dummy";
 
+/// Marker key used to identify invalid tool arguments (non-JSON objects)
+pub const INVALID_TOOL_ARGS_MARKER: &str = "__error__invalid_args_json__";
+
 pub const MAX_NUMBER_OF_IMAGES_PER_REQUEST: usize = 10;
 
 /// In bytes - 10 MB

crates/chat-cli/src/cli/chat/parser.rs

@@ -25,6 +25,7 @@ use tracing::{
     warn,
 };
 
+use super::consts::INVALID_TOOL_ARGS_MARKER;
 use super::message::{
     AssistantMessage,
     AssistantToolUse,
@@ -472,7 +473,18 @@ impl ResponseParser {
         }
 
         let args = match serde_json::from_str(&tool_string) {
-            Ok(args) => args,
+            Ok(args) => {
+                // Ensure we have a valid JSON object
+                match args {
+                    serde_json::Value::Object(_) => args,
+                    _ => {
+                        error!("Received non-object JSON for tool arguments: {:?}", args);
+                        serde_json::json!({
+                            INVALID_TOOL_ARGS_MARKER: format!("Expected JSON object, got: {:?}", args)
+                        })
+                    },
+                }
+            },
             Err(err) if !tool_string.is_empty() => {
                 // If we failed deserializing after waiting for a long time, then this is most
                 // likely bedrock responding with a stop event for some reason without actually
@@ -753,4 +765,73 @@ mod tests {
             "assistant text preceding a code reference should be ignored as this indicates licensed code is being returned"
         );
     }
+
+    #[tokio::test]
+    async fn test_response_parser_avoid_invalid_json() {
+        let content_to_ignore = "IGNORE ME PLEASE";
+        let tool_use_id = "TEST_ID".to_string();
+        let tool_name = "execute_bash".to_string();
+        let tool_args = serde_json::json!("invalid json").to_string();
+        let mut events = vec![
+            ChatResponseStream::AssistantResponseEvent {
+                content: "hi".to_string(),
+            },
+            ChatResponseStream::AssistantResponseEvent {
+                content: " there".to_string(),
+            },
+            ChatResponseStream::AssistantResponseEvent {
+                content: content_to_ignore.to_string(),
+            },
+            ChatResponseStream::CodeReferenceEvent(()),
+            ChatResponseStream::ToolUseEvent {
+                tool_use_id: tool_use_id.clone(),
+                name: tool_name.clone(),
+                input: None,
+                stop: None,
+            },
+            ChatResponseStream::ToolUseEvent {
+                tool_use_id: tool_use_id.clone(),
+                name: tool_name.clone(),
+                input: Some(tool_args),
+                stop: None,
+            },
+        ];
+        events.reverse();
+        let mock = SendMessageOutput::Mock(events);
+        let mut parser = ResponseParser::new(
+            mock,
+            "".to_string(),
+            None,
+            1,
+            vec![],
+            mpsc::channel(32).0,
+            Instant::now(),
+            SystemTime::now(),
+            CancellationToken::new(),
+            Arc::new(Mutex::new(None)),
+        );
+
+        let mut output = String::new();
+        let mut found_invalid_marker = false;
+        for _ in 0..5 {
+            let event = parser.recv().await.unwrap();
+            output.push_str(&format!("{:?}", event));
+
+            // Check for invalid args marker in ToolUse events
+            if let ResponseEvent::ToolUse(tool_use) = event {
+                if tool_use.args.get(INVALID_TOOL_ARGS_MARKER).is_some() {
+                    found_invalid_marker = true;
+                }
+            }
+        }
+
+        assert!(
+            !output.contains(content_to_ignore),
+            "assistant text preceding a code reference should be ignored as this indicates licensed code is being returned"
+        );
+        assert!(
+            found_invalid_marker,
+            "Expected to find invalid args marker for non-object JSON"
+        );
+    }
 }

crates/chat-cli/src/cli/chat/tool_manager.rs

@@ -64,7 +64,10 @@ use crate::cli::agent::{
     McpServerConfig,
 };
 use crate::cli::chat::cli::prompts::GetPromptError;
-use crate::cli::chat::consts::DUMMY_TOOL_NAME;
+use crate::cli::chat::consts::{
+    DUMMY_TOOL_NAME,
+    INVALID_TOOL_ARGS_MARKER,
+};
 use crate::cli::chat::message::AssistantToolUse;
 use crate::cli::chat::server_messenger::{
     ServerMessengerBuilder,
@@ -847,6 +850,20 @@ impl ToolManager {
     }
 
     pub async fn get_tool_from_tool_use(&mut self, value: AssistantToolUse) -> Result<Tool, ToolResult> {
+        // Check for invalid args marker
+        // in case parser, identified some
+        // fundamental error in the inputs.
+        if let Some(error_msg) = value.args.get(INVALID_TOOL_ARGS_MARKER).and_then(|v| v.as_str()) {
+            return Err(ToolResult {
+                tool_use_id: value.id.clone(),
+                content: vec![ToolResultContentBlock::Text(format!(
+                    "The tool \"{}\" is supplied with invalid input format. {}",
+                    value.name, error_msg
+                ))],
+                status: ToolResultStatus::Error,
+            });
+        }
+
         let map_err = |parse_error| ToolResult {
             tool_use_id: value.id.clone(),
             content: vec![ToolResultContentBlock::Text(format!(