From df2cb78228a1f8001e62024d30c7969a837f2709 Mon Sep 17 00:00:00 2001 From: Michael Neale Date: Fri, 23 May 2025 12:27:17 +1000 Subject: [PATCH] large tool response handling (#2629) --- crates/goose/src/agents/agent.rs | 5 +- .../src/agents/large_response_handler.rs | 247 ++++++++++++++++++ crates/goose/src/agents/mod.rs | 1 + 3 files changed, 252 insertions(+), 1 deletion(-) create mode 100644 crates/goose/src/agents/large_response_handler.rs diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index d7662eb6..2099cd77 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -195,7 +195,10 @@ impl Agent { "output" = serde_json::to_string(&result).unwrap(), ); - (request_id, result) + // Process the response to handle large text content + let processed_result = super::large_response_handler::process_tool_response(result); + + (request_id, processed_result) } pub(super) async fn manage_extensions( diff --git a/crates/goose/src/agents/large_response_handler.rs b/crates/goose/src/agents/large_response_handler.rs new file mode 100644 index 00000000..29141bc8 --- /dev/null +++ b/crates/goose/src/agents/large_response_handler.rs @@ -0,0 +1,247 @@ +use chrono::Utc; +use mcp_core::{Content, ToolError}; +use std::fs::File; +use std::io::Write; + +// Constant for the size threshold (20K characters) +const LARGE_TEXT_THRESHOLD: usize = 20_000; + +/// Process tool response and handle large text content +pub fn process_tool_response( + response: Result, ToolError>, +) -> Result, ToolError> { + match response { + Ok(contents) => { + let mut processed_contents = Vec::new(); + + for content in contents { + match content { + Content::Text(text_content) => { + // Check if text exceeds threshold + if text_content.text.len() > LARGE_TEXT_THRESHOLD { + // Write to temp file + match write_large_text_to_file(&text_content.text) { + Ok(file_path) => { + // Create a new text content with reference to the file + let message = format!( + "The response returned from the tool call was larger ({} characters) and is stored in the file which you can use other tools to examine or search in: {}", + text_content.text.len(), + file_path + ); + processed_contents.push(Content::text(message)); + } + Err(e) => { + // If file writing fails, include original content with warning + let warning = format!( + "Warning: Failed to write large response to file: {}. Showing full content instead.\n\n{}", + e, + text_content.text + ); + processed_contents.push(Content::text(warning)); + } + } + } else { + // Keep original content for smaller texts + processed_contents.push(Content::Text(text_content)); + } + } + // Pass through other content types unchanged + _ => processed_contents.push(content), + } + } + + Ok(processed_contents) + } + Err(e) => Err(e), + } +} + +/// Write large text content to a temporary file +fn write_large_text_to_file(content: &str) -> Result { + // Create temp directory if it doesn't exist + let temp_dir = std::env::temp_dir().join("goose_mcp_responses"); + std::fs::create_dir_all(&temp_dir)?; + + // Generate a unique filename with timestamp + let timestamp = Utc::now().format("%Y%m%d_%H%M%S"); + let filename = format!("mcp_response_{}.txt", timestamp); + let file_path = temp_dir.join(&filename); + + // Write content to file + let mut file = File::create(&file_path)?; + file.write_all(content.as_bytes())?; + + Ok(file_path.to_string_lossy().to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + use mcp_core::{Content, ImageContent, TextContent, ToolError}; + use std::fs; + use std::path::Path; + + #[test] + fn test_small_text_response_passes_through() { + // Create a small text response + let small_text = "This is a small text response"; + let content = Content::Text(TextContent { + text: small_text.to_string(), + annotations: None, + }); + + let response = Ok(vec![content]); + + // Process the response + let processed = process_tool_response(response).unwrap(); + + // Verify the response is unchanged + assert_eq!(processed.len(), 1); + if let Content::Text(text_content) = &processed[0] { + assert_eq!(text_content.text, small_text); + } else { + panic!("Expected text content"); + } + } + + #[test] + fn test_large_text_response_redirected_to_file() { + // Create a text larger than the threshold + let large_text = "a".repeat(LARGE_TEXT_THRESHOLD + 1000); + let content = Content::Text(TextContent { + text: large_text.clone(), + annotations: None, + }); + + let response = Ok(vec![content]); + + // Process the response + let processed = process_tool_response(response).unwrap(); + + // Verify the response contains a message about the file + assert_eq!(processed.len(), 1); + if let Content::Text(text_content) = &processed[0] { + assert!(text_content + .text + .contains("The response returned from the tool call was larger")); + assert!(text_content.text.contains("characters")); + + // Extract the file path from the message + if let Some(file_path) = text_content.text.split("stored in the file: ").nth(1) { + // Verify the file exists and contains the original text + let path = Path::new(file_path.trim()); + if path.exists() { + // Only check content if file exists (may not exist in CI environments) + if let Ok(file_content) = fs::read_to_string(path) { + assert_eq!(file_content, large_text); + } + + // Clean up the file + let _ = fs::remove_file(path); // Ignore errors on cleanup + } + } + } else { + panic!("Expected text content"); + } + } + + #[test] + fn test_image_content_passes_through() { + // Create an image content + let image_content = Content::Image(ImageContent { + data: "base64data".to_string(), + mime_type: "image/png".to_string(), + annotations: None, + }); + + let response = Ok(vec![image_content]); + + // Process the response + let processed = process_tool_response(response).unwrap(); + + // Verify the response is unchanged + assert_eq!(processed.len(), 1); + match &processed[0] { + Content::Image(img) => { + assert_eq!(img.data, "base64data"); + assert_eq!(img.mime_type, "image/png"); + } + _ => panic!("Expected image content"), + } + } + + #[test] + fn test_mixed_content_handled_correctly() { + // Create a response with mixed content types + let small_text = Content::text("Small text"); + let large_text = Content::Text(TextContent { + text: "a".repeat(LARGE_TEXT_THRESHOLD + 1000), + annotations: None, + }); + let image = Content::Image(ImageContent { + data: "image_data".to_string(), + mime_type: "image/jpeg".to_string(), + annotations: None, + }); + + let response = Ok(vec![small_text, large_text, image]); + + // Process the response + let processed = process_tool_response(response).unwrap(); + + // Verify each item is handled correctly + assert_eq!(processed.len(), 3); + + // First item should be unchanged small text + if let Content::Text(text_content) = &processed[0] { + assert_eq!(text_content.text, "Small text"); + } else { + panic!("Expected text content"); + } + + // Second item should be a message about the file + if let Content::Text(text_content) = &processed[1] { + assert!(text_content + .text + .contains("The response returned from the tool call was larger")); + + // Extract the file path and clean up + if let Some(file_path) = text_content.text.split("stored in the file: ").nth(1) { + let path = Path::new(file_path.trim()); + if path.exists() { + let _ = fs::remove_file(path); // Ignore errors on cleanup + } + } + } else { + panic!("Expected text content"); + } + + // Third item should be unchanged image + match &processed[2] { + Content::Image(img) => { + assert_eq!(img.data, "image_data"); + assert_eq!(img.mime_type, "image/jpeg"); + } + _ => panic!("Expected image content"), + } + } + + #[test] + fn test_error_response_passes_through() { + // Create an error response + let error = ToolError::ExecutionError("Test error".to_string()); + let response: Result, ToolError> = Err(error); + + // Process the response + let processed = process_tool_response(response); + + // Verify the error is passed through unchanged + assert!(processed.is_err()); + match processed { + Err(ToolError::ExecutionError(msg)) => { + assert_eq!(msg, "Test error"); + } + _ => panic!("Expected execution error"), + } + } +} diff --git a/crates/goose/src/agents/mod.rs b/crates/goose/src/agents/mod.rs index ca48dd25..3b39a6f8 100644 --- a/crates/goose/src/agents/mod.rs +++ b/crates/goose/src/agents/mod.rs @@ -2,6 +2,7 @@ mod agent; mod context; pub mod extension; pub mod extension_manager; +mod large_response_handler; pub mod platform_tools; pub mod prompt_manager; mod reply_parts;