use chrono::Utc; use rmcp::model::{CallToolResult, Content, ErrorData}; use std::fs::File; use std::io::Write; const LARGE_TEXT_THRESHOLD: usize = 200_000; /// Process tool response and handle large text content pub fn process_tool_response( response: Result, ) -> Result { match response { Ok(mut result) => { let mut processed_contents = Vec::new(); for content in result.content { match content.as_text() { Some(text_content) => { // Check if text exceeds threshold if text_content.text.chars().count() > 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.chars().count(), 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); } } None => { // Pass through other content types unchanged processed_contents.push(content); } } } result.content = processed_contents; Ok(result) } 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%.6f"); 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 rmcp::model::{Content, ErrorCode, ErrorData}; use std::borrow::Cow; 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(small_text.to_string()); let response = Ok(CallToolResult { content: vec![content], structured_content: None, is_error: Some(false), meta: None, }); // Process the response let processed = process_tool_response(response).unwrap(); // Verify the response is unchanged assert_eq!(processed.content.len(), 1); if let Some(text_content) = processed.content[0].as_text() { 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(large_text.clone()); let response = Ok(CallToolResult { content: vec![content], structured_content: None, is_error: Some(false), meta: None, }); // Process the response let processed = process_tool_response(response).unwrap(); // Verify the response contains a message about the file assert_eq!(processed.content.len(), 1); if let Some(text_content) = processed.content[0].as_text() { 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("base64data".to_string(), "image/png".to_string()); let response = Ok(CallToolResult { content: vec![image_content], structured_content: None, is_error: Some(false), meta: None, }); // Process the response let processed = process_tool_response(response).unwrap(); // Verify the response is unchanged assert_eq!(processed.content.len(), 1); if let Some(img) = processed.content[0].as_image() { assert_eq!(img.data, "base64data"); assert_eq!(img.mime_type, "image/png"); } else { 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("a".repeat(LARGE_TEXT_THRESHOLD + 1000)); let image = Content::image("image_data".to_string(), "image/jpeg".to_string()); let response = Ok(CallToolResult { content: vec![small_text, large_text, image], structured_content: None, is_error: Some(false), meta: None, }); // Process the response let processed = process_tool_response(response).unwrap(); // Verify each item is handled correctly assert_eq!(processed.content.len(), 3); // First item should be unchanged small text if let Some(text_content) = processed.content[0].as_text() { assert_eq!(text_content.text, "Small text"); } else { panic!("Expected text content"); } // Second item should be a message about the file if let Some(text_content) = processed.content[1].as_text() { 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 if let Some(img) = processed.content[2].as_image() { assert_eq!(img.data, "image_data"); assert_eq!(img.mime_type, "image/jpeg"); } else { panic!("Expected image content"); } } #[test] fn test_error_response_passes_through() { // Create an error response let error = ErrorData { code: ErrorCode::INTERNAL_ERROR, message: Cow::from("Test error"), data: None, }; let response: Result = 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(err) => { assert_eq!(err.code, ErrorCode::INTERNAL_ERROR); assert_eq!(err.message, "Test error"); } _ => panic!("Expected execution error"), } } }