[fix] Update to use a new RecvErrorKind instead of custom error handling.

Author: kennvene-aws

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

@@ -13,9 +13,6 @@ 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/mod.rs

@@ -2694,6 +2694,53 @@ impl ChatSession {
                                     .await?,
                             ));
                         },
+                        RecvErrorKind::ToolValidationError {
+                            tool_use_id,
+                            name,
+                            message,
+                            error_message,
+                        } => {
+                            self.send_chat_telemetry(
+                                os,
+                                TelemetryResult::Failed,
+                                Some(reason),
+                                Some(reason_desc),
+                                status_code,
+                                false, // We retry the request, so don't end the current turn yet.
+                            )
+                            .await;
+
+                            error!(
+                                recv_error.request_metadata.request_id,
+                                tool_use_id, name, error_message, "Tool validation failed"
+                            );
+                            self.conversation
+                                .push_assistant_message(os, *message, Some(recv_error.request_metadata));
+                            let tool_results = vec![ToolUseResult {
+                                tool_use_id,
+                                content: vec![ToolUseResultBlock::Text(format!(
+                                    "Tool validation failed: {}. Please ensure tool arguments are provided as a valid JSON object.",
+                                    error_message
+                                ))],
+                                status: ToolResultStatus::Error,
+                            }];
+                            // User hint of what happened
+                            let _ = queue!(
+                                self.stdout,
+                                style::Print("\n\n"),
+                                style::SetForegroundColor(Color::Yellow),
+                                style::Print(format!("Tool validation failed: {}", error_message)),
+                                style::ResetColor,
+                                style::Print("\n"),
+                            );
+                            self.conversation.add_tool_results(tool_results);
+                            self.send_tool_use_telemetry(os).await;
+                            return Ok(ChatState::HandleResponseStream(
+                                self.conversation
+                                    .as_sendable_conversation_state(os, &mut self.stderr, false)
+                                    .await?,
+                            ));
+                        },
                         _ => {
                             self.send_chat_telemetry(
                                 os,

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

@@ -25,7 +25,6 @@ use tracing::{
     warn,
 };
 
-use super::consts::INVALID_TOOL_ARGS_MARKER;
 use super::message::{
     AssistantMessage,
     AssistantToolUse,
@@ -92,6 +91,7 @@ impl RecvError {
             RecvErrorKind::StreamTimeout { .. } => None,
             RecvErrorKind::UnexpectedToolUseEos { .. } => None,
             RecvErrorKind::Cancelled => None,
+            RecvErrorKind::ToolValidationError { .. } => None,
         }
     }
 }
@@ -104,6 +104,7 @@ impl ReasonCode for RecvError {
             RecvErrorKind::StreamTimeout { .. } => "RecvErrorStreamTimeout".to_string(),
             RecvErrorKind::UnexpectedToolUseEos { .. } => "RecvErrorUnexpectedToolUseEos".to_string(),
             RecvErrorKind::Cancelled => "Interrupted".to_string(),
+            RecvErrorKind::ToolValidationError { .. } => "RecvErrorToolValidation".to_string(),
         }
     }
 }
@@ -152,6 +153,14 @@ pub enum RecvErrorKind {
     /// The stream processing task was cancelled
     #[error("Stream handling was cancelled")]
     Cancelled,
+    /// Tool validation failed due to invalid arguments
+    #[error("Tool validation failed for tool: {} with id: {}", .name, .tool_use_id)]
+    ToolValidationError {
+        tool_use_id: String,
+        name: String,
+        message: Box<AssistantMessage>,
+        error_message: String,
+    },
 }
 
 /// Represents a response stream from a call to the SendMessage API.
@@ -479,9 +488,34 @@ impl ResponseParser {
                     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)
-                        })
+                        let warning_args = serde_json::Value::Object(
+                            [(
+                                "key".to_string(),
+                                serde_json::Value::String(
+                                    "WARNING: the actual tool use arguments were not a valid JSON object".to_string(),
+                                ),
+                            )]
+                            .into_iter()
+                            .collect(),
+                        );
+                        self.tool_uses.push(AssistantToolUse {
+                            id: id.clone(),
+                            name: name.clone(),
+                            orig_name: name.clone(),
+                            args: warning_args.clone(),
+                            orig_args: warning_args.clone(),
+                        });
+                        let message = Box::new(AssistantMessage::new_tool_use(
+                            Some(self.message_id.clone()),
+                            std::mem::take(&mut self.assistant_text),
+                            self.tool_uses.clone().into_iter().collect(),
+                        ));
+                        return Err(self.error(RecvErrorKind::ToolValidationError {
+                            tool_use_id: id,
+                            name,
+                            message,
+                            error_message: format!("Expected JSON object, got: {:?}", args),
+                        }));
                     },
                 }
             },
@@ -812,16 +846,18 @@ mod tests {
         );
 
         let mut output = String::new();
-        let mut found_invalid_marker = false;
+        let mut found_validation_error = 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;
-                }
+            match parser.recv().await {
+                Ok(event) => {
+                    output.push_str(&format!("{:?}", event));
+                },
+                Err(recv_error) => {
+                    if matches!(recv_error.source, RecvErrorKind::ToolValidationError { .. }) {
+                        found_validation_error = true;
+                    }
+                    break;
+                },
             }
         }
 
@@ -830,8 +866,8 @@ mod tests {
             "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"
+            found_validation_error,
+            "Expected to find tool validation error for non-object JSON"
         );
     }
 }

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

@@ -64,10 +64,7 @@ use crate::cli::agent::{
     McpServerConfig,
 };
 use crate::cli::chat::cli::prompts::GetPromptError;
-use crate::cli::chat::consts::{
-    DUMMY_TOOL_NAME,
-    INVALID_TOOL_ARGS_MARKER,
-};
+use crate::cli::chat::consts::DUMMY_TOOL_NAME;
 use crate::cli::chat::message::AssistantToolUse;
 use crate::cli::chat::server_messenger::{
     ServerMessengerBuilder,
@@ -849,54 +846,7 @@ impl ToolManager {
         Ok(self.schema.clone())
     }
 
-    /// Validates that the tool arguments are a valid JSON object and checks for invalid args
-    /// marker.
-    ///
-    /// # Arguments
-    /// * `args` - The JSON value containing tool arguments
-    /// * `tool_name` - Name of the tool for error reporting
-    /// * `tool_use_id` - ID of the tool use for error reporting
-    ///
-    /// # Returns
-    /// * `Ok(())` if validation passes
-    /// * `Err(ToolResult)` if validation fails with appropriate error message
-    fn validate_tool_args(
-        args: &serde_json::Value,
-        tool_name: &str,
-        tool_use_id: &str,
-    ) -> Result<(), ToolResult> {
-        // Ensure args is an object for safe access
-        if !args.is_object() {
-            return Err(ToolResult {
-                tool_use_id: tool_use_id.to_string(),
-                content: vec![ToolResultContentBlock::Text(format!(
-                    "The tool \"{}\" requires arguments to be a JSON object, but received: {}",
-                    tool_name,
-                    args
-                ))],
-                status: ToolResultStatus::Error,
-            });
-        }
-
-        // Now safely check for invalid args marker (args is guaranteed to be an object)
-        if let Some(error_msg) = args.get(INVALID_TOOL_ARGS_MARKER).and_then(|v| v.as_str()) {
-            return Err(ToolResult {
-                tool_use_id: tool_use_id.to_string(),
-                content: vec![ToolResultContentBlock::Text(format!(
-                    "The tool \"{}\" is supplied with invalid input format. {}",
-                    tool_name, error_msg
-                ))],
-                status: ToolResultStatus::Error,
-            });
-        }
-
-        Ok(())
-    }
-
     pub async fn get_tool_from_tool_use(&mut self, value: AssistantToolUse) -> Result<Tool, ToolResult> {
-        // Validate tool arguments
-        Self::validate_tool_args(&value.args, &value.name, &value.id)?;
-
         let map_err = |parse_error| ToolResult {
             tool_use_id: value.id.clone(),
             content: vec![ToolResultContentBlock::Text(format!(