diff --git a/crates/goose-mcp/src/google_drive/mod.rs b/crates/goose-mcp/src/google_drive/mod.rs index 02b5b9753ccc..2558751233ff 100644 --- a/crates/goose-mcp/src/google_drive/mod.rs +++ b/crates/goose-mcp/src/google_drive/mod.rs @@ -1,6 +1,8 @@ mod oauth_pkce; pub mod storage; +use anyhow::{Context, Error}; +use base64::Engine; use indoc::indoc; use oauth_pkce::PkceOAuth2Client; use regex::Regex; @@ -632,6 +634,67 @@ impl GoogleDriveRouter { image_regex.replace_all(input, "").to_string() } + // Helper function that processes one captured image. + // It decodes the base64 data, resizes the image if its width exceeds `max_width`, + // and then returns a new image tag (always output as PNG). + // logic copied from developer/mod.rs + fn process_image(&self, caps: ®ex::Captures, max_width: u32) -> Result { + let base64_data = &caps["data"]; + + // Decode the Base64 data. + let image_bytes = base64::prelude::BASE64_STANDARD + .decode(base64_data) + .context("Failed to decode base64 image data")?; + + // Load the image from the decoded bytes. + let img = xcap::image::load_from_memory(&image_bytes) + .context("Failed to load image from memory")?; + + // Resize the image if necessary. + let mut processed_image = img; + if processed_image.width() > max_width { + let scale = max_width as f32 / processed_image.width() as f32; + let new_height = (processed_image.height() as f32 * scale) as u32; + processed_image = xcap::image::DynamicImage::ImageRgba8(xcap::image::imageops::resize( + &processed_image, + max_width, + new_height, + xcap::image::imageops::FilterType::Lanczos3, + )); + } + + // Write the processed image to an in-memory buffer in PNG format. + let mut buffer: Vec = Vec::new(); + processed_image + .write_to(&mut Cursor::new(&mut buffer), xcap::image::ImageFormat::Png) + .context("Failed to write processed image to buffer")?; + + // Re-encode the buffer back into a Base64 string. + let data = base64::prelude::BASE64_STANDARD.encode(&buffer); + Ok(Content::image(data, "image/png")) + } + + /// Resizes all base64-encoded images found in the input string. + /// If any image fails to process, an error is returned. + fn resize_images(&self, input: &str) -> Result, Error> { + // Regex to match and capture the MIME type and Base64 data. + let image_regex = + Regex::new(r"[a-zA-Z0-9.+-]+);base64,(?P[^>]+)>") + .context("Failed to compile regex")?; + + let mut result: Vec = Vec::new(); + + // Iterate over all matches, process them, and rebuild the output string. + for caps in image_regex.captures_iter(input) { + let processed_tag = self + .process_image(&caps, 768) + .context("Failed to process one of the images")?; + result.push(processed_tag); + } + + Ok(result) + } + // Downloading content with alt=media only works if the file is stored in Drive. // To download Google Docs, Sheets, and Slides use files.export instead. async fn export_google_file( @@ -665,13 +728,22 @@ impl GoogleDriveRouter { Ok(r) => { if let Ok(body) = r.into_body().collect().await { if let Ok(response) = String::from_utf8(body.to_bytes().to_vec()) { - let content = if !include_images { - self.strip_image_body(&response) + if !include_images { + let content = self.strip_image_body(&response); + Ok(vec![Content::text(content).with_priority(0.1)]) } else { - response - }; + let images = self.resize_images(&response).map_err(|e| { + ToolError::ExecutionError(format!( + "Failed to resize image(s): {}", + e + )) + })?; - Ok(vec![Content::text(content).with_priority(0.1)]) + let content = self.strip_image_body(&response); + Ok(std::iter::once(Content::text(content).with_priority(0.1)) + .chain(images.iter().cloned()) + .collect::>()) + } } else { Err(ToolError::ExecutionError(format!( "Failed to export google drive to string, {}.", @@ -717,13 +789,22 @@ impl GoogleDriveRouter { if mime_type.starts_with("text/") || mime_type == "application/json" { if let Ok(body) = r.0.into_body().collect().await { if let Ok(response) = String::from_utf8(body.to_bytes().to_vec()) { - let content = if !include_images { - self.strip_image_body(&response) + if !include_images { + let content = self.strip_image_body(&response); + Ok(vec![Content::text(content).with_priority(0.1)]) } else { - response - }; - - Ok(vec![Content::text(content).with_priority(0.1)]) + let images = self.resize_images(&response).map_err(|e| { + ToolError::ExecutionError(format!( + "Failed to resize image(s): {}", + e + )) + })?; + + let content = self.strip_image_body(&response); + Ok(std::iter::once(Content::text(content).with_priority(0.1)) + .chain(images.iter().cloned()) + .collect::>()) + } } else { Err(ToolError::ExecutionError(format!( "Failed to convert google drive to string, {}.",