large tool response handling (#2629)

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(

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<Vec<Content>, ToolError>,
+) -> Result<Vec<Content>, 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<String, std::io::Error> {
+    // 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<Vec<Content>, 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"),
+        }
+    }
+}

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;