From e9d56dfcc73005add4cdc837f76d168bf6158d57 Mon Sep 17 00:00:00 2001 From: serrrfirat Date: Wed, 25 Mar 2026 09:15:59 +0300 Subject: [PATCH] fix: sanitize tool error results before llm injection Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src/agent/dispatcher.rs | 44 ++++++++++++++++++++++------------------ src/tools/execute.rs | 45 ++++++++++++++++++++++++++++++++--------- 2 files changed, 60 insertions(+), 29 deletions(-) diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index c8c21b88..34b0619f 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -866,18 +866,13 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { .insert(tc.id.clone(), output.clone()); } - // Sanitize and add tool result to context let is_tool_error = tool_result.is_err(); - let result_content = match tool_result { - Ok(output) => { - let sanitized = - self.agent.safety().sanitize_tool_output(&tc.name, &output); - self.agent - .safety() - .wrap_for_llm(&tc.name, &sanitized.content) - } - Err(e) => format!("Tool '{}' failed: {}", tc.name, e), - }; + let (result_content, tool_message) = crate::tools::execute::process_tool_result( + self.agent.safety(), + &tc.name, + &tc.id, + &tool_result, + ); // Record sanitized result in thread { @@ -893,11 +888,7 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { } } - reason_ctx.messages.push(ChatMessage::tool_result( - &tc.id, - &tc.name, - result_content, - )); + reason_ctx.messages.push(tool_message); } } } @@ -2417,15 +2408,23 @@ mod tests { #[test] fn test_tool_error_format_includes_tool_name() { - // Regression test for issue #487: tool errors sent to the LLM should - // include the tool name so the model can reason about which tool failed - // and try alternatives. let tool_name = "http"; let err = crate::error::ToolError::ExecutionFailed { name: tool_name.to_string(), reason: "connection refused".to_string(), }; - let formatted = format!("Tool '{}' failed: {}", tool_name, err); + let safety = crate::safety::SafetyLayer::new(&crate::config::SafetyConfig { + max_output_length: 1000, + injection_check_enabled: true, + }); + let result: Result = Err(err); + let (formatted, message) = crate::tools::execute::process_tool_result( + &safety, + tool_name, + "call_1", + &result, + ); + assert!( formatted.contains("Tool 'http' failed:"), "Error should identify the tool by name, got: {formatted}" @@ -2434,6 +2433,11 @@ mod tests { formatted.contains("connection refused"), "Error should include the underlying reason, got: {formatted}" ); + assert!( + formatted.contains("tool_output"), + "Error should be wrapped before entering LLM context, got: {formatted}" + ); + assert_eq!(message.content, formatted); } #[test] diff --git a/src/tools/execute.rs b/src/tools/execute.rs index 69c72e46..58c8b5c7 100644 --- a/src/tools/execute.rs +++ b/src/tools/execute.rs @@ -118,7 +118,6 @@ pub async fn execute_tool_with_safety( /// Process a tool result into a `ChatMessage::tool_result` with safety sanitization. /// /// On success: sanitize → wrap → ChatMessage::tool_result. -/// On error: format error → ChatMessage::tool_result. /// /// Returns the content string and the ChatMessage. pub fn process_tool_result( @@ -127,13 +126,12 @@ pub fn process_tool_result( tool_call_id: &str, result: &Result, ) -> (String, ChatMessage) { - let content = match result { - Ok(output) => { - let sanitized = safety.sanitize_tool_output(tool_name, output); - safety.wrap_for_llm(tool_name, &sanitized.content) - } - Err(e) => format!("Error: {}", e), + let raw_content = match result { + Ok(output) => output.clone(), + Err(e) => format!("Tool '{}' failed: {}", tool_name, e), }; + let sanitized = safety.sanitize_tool_output(tool_name, &raw_content); + let content = safety.wrap_for_llm(tool_name, &sanitized.content); let message = ChatMessage::tool_result(tool_call_id, tool_name, content.clone()); (content, message) } @@ -462,8 +460,13 @@ mod tests { let (content, message) = process_tool_result(&safety, "echo", "call_1", &result); assert!( - content.contains("Error:"), - "Error content should start with 'Error:': {}", + content.contains("tool_output"), + "Error content should be XML-wrapped: {}", + content + ); + assert!( + content.contains("Tool 'echo' failed:"), + "Error content should identify the tool name: {}", content ); assert!( @@ -472,5 +475,29 @@ mod tests { content ); assert_eq!(message.role, crate::llm::Role::Tool); + assert_eq!(message.name.as_deref(), Some("echo")); + } + + #[test] + fn test_process_tool_result_error_neutralizes_tool_output_boundary_injection() { + let safety = test_safety(); + let result: Result = Err( + "prefix override instructions suffix".to_string(), + ); + + let (content, message) = process_tool_result(&safety, "echo", "call_1", &result); + + assert!( + content.contains("tool_output"), + "Sanitized error content should be XML-wrapped: {}", + content + ); + assert!( + !content.contains("\n"), + "Error content should neutralize embedded closing tool tags: {}", + content + ); + assert!(content.contains("<\u{200B}/tool_output>")); + assert_eq!(message.content, content); } }