Initial public release of Warp.

Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b

app/src/util/image.rs

@@ -0,0 +1,113 @@
+//! Shared image processing utilities for agent mode.
+//!
+//! This module provides common functionality for processing images before they are
+//! sent to the AI agent, whether attached by the user or read via the read_files tool.
+
+use image::{GenericImageView, ImageError};
+
+/// Max image size is 3.75 MB.
+/// The max size of an image we will send is 5MB. However, due to the 33% inflation of Base64, this means
+/// the largest size a user can attach is actually ~3.75MB.
+pub const MAX_IMAGE_SIZE_BYTES: usize = 3750 * 1000;
+
+/// 1.15 Megapixels
+pub const MAX_IMAGE_PIXELS: f64 = 1150. * 1000.;
+
+/// Maximum dimension (width or height) for images.
+pub const MAX_IMAGE_DIMENSION: f64 = 2000.;
+
+/// Maximum number of images that can be attached per query/task.
+pub const MAX_IMAGE_COUNT_FOR_QUERY: usize = 20;
+
+/// Minimum bytes needed for image format detection using magic number signatures.
+pub const MIN_IMAGE_HEADER_SIZE: usize = 8;
+
+/// Supported image MIME types for agent mode.
+pub const SUPPORTED_IMAGE_MIME_TYPES: &[&str] = &[
+    "image/png",
+    "image/jpeg",
+    "image/jpg",
+    "image/gif",
+    "image/webp",
+];
+
+/// Checks if the given MIME type is a supported image type.
+pub fn is_supported_image_mime_type(mime_type: &str) -> bool {
+    SUPPORTED_IMAGE_MIME_TYPES.contains(&mime_type)
+}
+
+/// Resizes an image if it exceeds the maximum pixel count, and ensures
+/// resized outputs also respect the maximum dimension (width or height).
+///
+/// Returns the original image bytes if the image is already within the
+/// pixel limit; otherwise returns the resized image bytes in the original
+/// format.
+pub fn resize_image(image: &[u8]) -> Result<Vec<u8>, ImageError> {
+    let img = image::load_from_memory(image)?;
+
+    let (current_width, current_height) = img.dimensions();
+    let current_pixels = (current_width * current_height) as f64;
+
+    if current_pixels <= MAX_IMAGE_PIXELS {
+        return Ok(image.to_vec());
+    }
+
+    let original_format = image::guess_format(image)?;
+
+    let scale = (MAX_IMAGE_PIXELS / current_pixels).sqrt();
+
+    let mut new_width = current_width as f64 * scale;
+    let mut new_height = current_height as f64 * scale;
+
+    let scale_by_width = MAX_IMAGE_DIMENSION / new_width;
+    let scale_by_height = MAX_IMAGE_DIMENSION / new_height;
+    let scale = scale_by_width.min(scale_by_height).min(1.0);
+
+    new_width *= scale;
+    new_height *= scale;
+
+    let resized_img = img.thumbnail(new_width.round() as u32, new_height.round() as u32);
+
+    let mut output_bytes: Vec<u8> = Vec::new();
+    let mut writer = std::io::Cursor::new(&mut output_bytes);
+
+    resized_img.write_to(&mut writer, original_format)?;
+
+    Ok(output_bytes)
+}
+
+/// Result of processing an image for agent mode.
+#[derive(Debug)]
+pub enum ProcessImageResult {
+    /// Image was successfully processed and is within size limits.
+    Success {
+        /// The processed image bytes (resized if needed).
+        data: Vec<u8>,
+    },
+    /// Image is too large even after resizing.
+    TooLarge,
+    /// Error processing the image.
+    Error(ImageError),
+}
+
+/// Processes an image for agent mode: resizes if needed and checks size limits.
+///
+/// This applies the same processing that user-attached images go through.
+pub fn process_image_for_agent(image_data: &[u8]) -> ProcessImageResult {
+    match resize_image(image_data) {
+        Ok(resized_bytes) => {
+            if resized_bytes.len() > MAX_IMAGE_SIZE_BYTES {
+                ProcessImageResult::TooLarge
+            } else {
+                ProcessImageResult::Success {
+                    data: resized_bytes,
+                }
+            }
+        }
+        Err(err) => ProcessImageResult::Error(err),
+    }
+}
+
+#[cfg(test)]
+#[path = "image_tests.rs"]
+mod tests;