Bound agent context and preprocess oversized image attachments
This commit is contained in:
+187
-37
@@ -3,15 +3,24 @@
|
||||
//! 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 std::io::Cursor;
|
||||
use std::path::Path;
|
||||
|
||||
use image::{GenericImageView, ImageError};
|
||||
use image::codecs::jpeg::JpegEncoder;
|
||||
use image::{
|
||||
DynamicImage, ExtendedColorType, GenericImageView, ImageEncoder, ImageError, ImageFormat,
|
||||
};
|
||||
use mime_guess::from_path;
|
||||
|
||||
/// 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;
|
||||
const JPEG_MIME_TYPE: &str = "image/jpeg";
|
||||
const PNG_MIME_TYPE: &str = "image/png";
|
||||
const MIN_READABLE_IMAGE_DIMENSION: u32 = 512;
|
||||
const JPEG_QUALITIES: &[u8] = &[85, 75, 65, 55];
|
||||
const MAX_IMAGE_ENCODE_ATTEMPTS: usize = 8;
|
||||
|
||||
/// Cap on individual image bytes when relaying a dropped/pasted image to a
|
||||
/// CLI agent (Claude Code etc.) via the system clipboard. CLI agents handle
|
||||
@@ -62,44 +71,187 @@ 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).
|
||||
#[derive(Debug)]
|
||||
pub struct ProcessedImage {
|
||||
pub data: Vec<u8>,
|
||||
pub mime_type: String,
|
||||
}
|
||||
|
||||
/// Resizes and re-encodes an image until it fits the model-facing byte budget.
|
||||
///
|
||||
/// 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> {
|
||||
/// The source bytes are never modified. Opaque images may be converted to JPEG
|
||||
/// because that is the most reliable way to fit photo-like uploads under the
|
||||
/// provider cap; images with transparency stay PNG so alpha is preserved.
|
||||
pub fn prepare_image_for_agent(image: &[u8]) -> Result<Option<ProcessedImage>, ImageError> {
|
||||
let img = image::load_from_memory(image)?;
|
||||
let original_format = image::guess_format(image)?;
|
||||
let original_mime_type = mime_type_for_format(original_format);
|
||||
|
||||
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());
|
||||
if image.len() <= MAX_IMAGE_SIZE_BYTES && image_within_dimension_limits(&img) {
|
||||
return Ok(Some(ProcessedImage {
|
||||
data: image.to_vec(),
|
||||
mime_type: original_mime_type.to_string(),
|
||||
}));
|
||||
}
|
||||
|
||||
let original_format = image::guess_format(image)?;
|
||||
let has_transparency = image_has_transparency(&img);
|
||||
let mut width = target_dimension(img.width(), img.height()).0;
|
||||
let mut height = target_dimension(img.width(), img.height()).1;
|
||||
let mut encode_attempts = 0;
|
||||
|
||||
let scale = (MAX_IMAGE_PIXELS / current_pixels).sqrt();
|
||||
while encode_attempts < MAX_IMAGE_ENCODE_ATTEMPTS {
|
||||
let candidate = if width == img.width() && height == img.height() {
|
||||
img.clone()
|
||||
} else {
|
||||
img.thumbnail(width, height)
|
||||
};
|
||||
|
||||
let mut new_width = current_width as f64 * scale;
|
||||
let mut new_height = current_height as f64 * scale;
|
||||
if has_transparency {
|
||||
encode_attempts += 1;
|
||||
let data = encode_png(&candidate)?;
|
||||
if data.len() <= MAX_IMAGE_SIZE_BYTES {
|
||||
return Ok(Some(ProcessedImage {
|
||||
data,
|
||||
mime_type: PNG_MIME_TYPE.to_string(),
|
||||
}));
|
||||
}
|
||||
} else {
|
||||
encode_attempts += 1;
|
||||
let data = encode_png(&candidate)?;
|
||||
if data.len() <= MAX_IMAGE_SIZE_BYTES {
|
||||
return Ok(Some(ProcessedImage {
|
||||
data,
|
||||
mime_type: PNG_MIME_TYPE.to_string(),
|
||||
}));
|
||||
}
|
||||
|
||||
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);
|
||||
for quality in JPEG_QUALITIES {
|
||||
if encode_attempts >= MAX_IMAGE_ENCODE_ATTEMPTS {
|
||||
break;
|
||||
}
|
||||
encode_attempts += 1;
|
||||
let data = encode_jpeg(&candidate, *quality)?;
|
||||
if data.len() <= MAX_IMAGE_SIZE_BYTES {
|
||||
return Ok(Some(ProcessedImage {
|
||||
data,
|
||||
mime_type: JPEG_MIME_TYPE.to_string(),
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
new_width *= scale;
|
||||
new_height *= scale;
|
||||
let next = next_smaller_dimensions(width, height);
|
||||
if next == (width, height) {
|
||||
break;
|
||||
}
|
||||
(width, height) = next;
|
||||
}
|
||||
|
||||
let resized_img = img.thumbnail(new_width.round() as u32, new_height.round() as u32);
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
let mut output_bytes: Vec<u8> = Vec::new();
|
||||
let mut writer = std::io::Cursor::new(&mut output_bytes);
|
||||
/// Backwards-compatible helper for callers that only need processed bytes.
|
||||
pub fn resize_image(image: &[u8]) -> Result<Vec<u8>, ImageError> {
|
||||
prepare_image_for_agent(image)?.map_or_else(
|
||||
|| {
|
||||
Err(ImageError::Limits(image::error::LimitError::from_kind(
|
||||
image::error::LimitErrorKind::InsufficientMemory,
|
||||
)))
|
||||
},
|
||||
|processed| Ok(processed.data),
|
||||
)
|
||||
}
|
||||
|
||||
resized_img.write_to(&mut writer, original_format)?;
|
||||
fn image_within_dimension_limits(img: &DynamicImage) -> bool {
|
||||
let (width, height) = img.dimensions();
|
||||
let pixels = (width as f64) * (height as f64);
|
||||
pixels <= MAX_IMAGE_PIXELS
|
||||
&& (width as f64) <= MAX_IMAGE_DIMENSION
|
||||
&& (height as f64) <= MAX_IMAGE_DIMENSION
|
||||
}
|
||||
|
||||
Ok(output_bytes)
|
||||
fn target_dimension(width: u32, height: u32) -> (u32, u32) {
|
||||
let pixels = (width as f64) * (height as f64);
|
||||
let pixel_scale = if pixels > MAX_IMAGE_PIXELS {
|
||||
(MAX_IMAGE_PIXELS / pixels).sqrt()
|
||||
} else {
|
||||
1.0
|
||||
};
|
||||
let dimension_scale = (MAX_IMAGE_DIMENSION / width as f64)
|
||||
.min(MAX_IMAGE_DIMENSION / height as f64)
|
||||
.min(1.0);
|
||||
let scale = pixel_scale.min(dimension_scale).min(1.0);
|
||||
scaled_dimensions(width, height, scale)
|
||||
}
|
||||
|
||||
fn next_smaller_dimensions(width: u32, height: u32) -> (u32, u32) {
|
||||
if width <= MIN_READABLE_IMAGE_DIMENSION && height <= MIN_READABLE_IMAGE_DIMENSION {
|
||||
return (width, height);
|
||||
}
|
||||
|
||||
let next_width = ((width as f64) * 0.85).round() as u32;
|
||||
let next_height = ((height as f64) * 0.85).round() as u32;
|
||||
(
|
||||
next_width.max(MIN_READABLE_IMAGE_DIMENSION.min(width)),
|
||||
next_height.max(MIN_READABLE_IMAGE_DIMENSION.min(height)),
|
||||
)
|
||||
}
|
||||
|
||||
fn scaled_dimensions(width: u32, height: u32, scale: f64) -> (u32, u32) {
|
||||
(
|
||||
((width as f64) * scale).round().max(1.0) as u32,
|
||||
((height as f64) * scale).round().max(1.0) as u32,
|
||||
)
|
||||
}
|
||||
|
||||
fn image_has_transparency(img: &DynamicImage) -> bool {
|
||||
if !img.color().has_alpha() {
|
||||
return false;
|
||||
}
|
||||
img.to_rgba8().pixels().any(|pixel| pixel.0[3] < 255)
|
||||
}
|
||||
|
||||
fn encode_image_format(img: &DynamicImage, format: ImageFormat) -> Result<Vec<u8>, ImageError> {
|
||||
let mut data = Vec::new();
|
||||
img.write_to(&mut Cursor::new(&mut data), format)?;
|
||||
Ok(data)
|
||||
}
|
||||
|
||||
fn encode_png(img: &DynamicImage) -> Result<Vec<u8>, ImageError> {
|
||||
encode_image_format(img, ImageFormat::Png)
|
||||
}
|
||||
|
||||
fn encode_jpeg(img: &DynamicImage, quality: u8) -> Result<Vec<u8>, ImageError> {
|
||||
let rgb = img.to_rgb8();
|
||||
let mut data = Vec::new();
|
||||
JpegEncoder::new_with_quality(&mut data, quality).write_image(
|
||||
rgb.as_raw(),
|
||||
rgb.width(),
|
||||
rgb.height(),
|
||||
ExtendedColorType::Rgb8,
|
||||
)?;
|
||||
Ok(data)
|
||||
}
|
||||
|
||||
fn mime_type_for_format(format: ImageFormat) -> &'static str {
|
||||
match format {
|
||||
ImageFormat::Png => PNG_MIME_TYPE,
|
||||
ImageFormat::Jpeg => JPEG_MIME_TYPE,
|
||||
ImageFormat::Gif => "image/gif",
|
||||
ImageFormat::WebP => "image/webp",
|
||||
ImageFormat::Pnm => "image/x-portable-anymap",
|
||||
ImageFormat::Tiff => "image/tiff",
|
||||
ImageFormat::Tga => "image/x-tga",
|
||||
ImageFormat::Dds => "image/vnd-ms.dds",
|
||||
ImageFormat::Bmp => "image/bmp",
|
||||
ImageFormat::Ico => "image/x-icon",
|
||||
ImageFormat::Hdr => "image/vnd.radiance",
|
||||
ImageFormat::OpenExr => "image/x-exr",
|
||||
ImageFormat::Farbfeld => "image/farbfeld",
|
||||
ImageFormat::Avif => "image/avif",
|
||||
ImageFormat::Qoi => "image/qoi",
|
||||
_ => "application/octet-stream",
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of processing an image for agent mode.
|
||||
@@ -107,8 +259,10 @@ pub fn resize_image(image: &[u8]) -> Result<Vec<u8>, ImageError> {
|
||||
pub enum ProcessImageResult {
|
||||
/// Image was successfully processed and is within size limits.
|
||||
Success {
|
||||
/// The processed image bytes (resized if needed).
|
||||
/// The processed image bytes (resized/re-encoded if needed).
|
||||
data: Vec<u8>,
|
||||
/// MIME type matching the processed image bytes.
|
||||
mime_type: String,
|
||||
},
|
||||
/// Image is too large even after resizing.
|
||||
TooLarge,
|
||||
@@ -120,16 +274,12 @@ pub enum ProcessImageResult {
|
||||
///
|
||||
/// 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,
|
||||
}
|
||||
}
|
||||
}
|
||||
match prepare_image_for_agent(image_data) {
|
||||
Ok(Some(processed)) => ProcessImageResult::Success {
|
||||
data: processed.data,
|
||||
mime_type: processed.mime_type,
|
||||
},
|
||||
Ok(None) => ProcessImageResult::TooLarge,
|
||||
Err(err) => ProcessImageResult::Error(err),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use image::{ImageBuffer, Rgba};
|
||||
use image::{GenericImageView, ImageBuffer, Rgba};
|
||||
|
||||
use super::*;
|
||||
|
||||
@@ -32,11 +32,12 @@ fn test_process_image_for_agent_small_image() {
|
||||
let result = process_image_for_agent(&small_png);
|
||||
|
||||
match result {
|
||||
ProcessImageResult::Success { data } => {
|
||||
ProcessImageResult::Success { data, mime_type } => {
|
||||
// The processed image should be non-empty
|
||||
assert!(!data.is_empty());
|
||||
// For a small image, it should not be resized, so data should be similar
|
||||
assert!(data.len() <= MAX_IMAGE_SIZE_BYTES);
|
||||
assert_eq!(mime_type, "image/png");
|
||||
}
|
||||
other => panic!("Expected Success, got {:?}", other),
|
||||
}
|
||||
@@ -52,6 +53,67 @@ fn test_resize_image_small_image_unchanged() {
|
||||
assert_eq!(result.len(), original_len);
|
||||
}
|
||||
|
||||
fn create_noisy_png(width: u32, height: u32, transparent: bool) -> Vec<u8> {
|
||||
let img: ImageBuffer<Rgba<u8>, Vec<u8>> = ImageBuffer::from_fn(width, height, |x, y| {
|
||||
let seed = x
|
||||
.wrapping_mul(1_103_515_245)
|
||||
.wrapping_add(y.wrapping_mul(12_345));
|
||||
let alpha = if transparent && (x + y) % 3 == 0 {
|
||||
128
|
||||
} else {
|
||||
255
|
||||
};
|
||||
Rgba([
|
||||
(seed & 0xff) as u8,
|
||||
((seed >> 8) & 0xff) as u8,
|
||||
((seed >> 16) & 0xff) as u8,
|
||||
alpha,
|
||||
])
|
||||
});
|
||||
|
||||
let mut bytes: Vec<u8> = Vec::new();
|
||||
let mut cursor = std::io::Cursor::new(&mut bytes);
|
||||
img.write_to(&mut cursor, image::ImageFormat::Png).unwrap();
|
||||
bytes
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_process_image_for_agent_compresses_byte_heavy_opaque_image() {
|
||||
let noisy_png = create_noisy_png(1500, 1500, false);
|
||||
assert!(noisy_png.len() > MAX_IMAGE_SIZE_BYTES);
|
||||
|
||||
let result = process_image_for_agent(&noisy_png);
|
||||
|
||||
match result {
|
||||
ProcessImageResult::Success { data, mime_type } => {
|
||||
assert!(data.len() <= MAX_IMAGE_SIZE_BYTES);
|
||||
assert_eq!(mime_type, "image/jpeg");
|
||||
let img = image::load_from_memory(&data).unwrap();
|
||||
let (width, height) = img.dimensions();
|
||||
assert!((width as f64 * height as f64) <= MAX_IMAGE_PIXELS);
|
||||
assert_eq!(width, height);
|
||||
}
|
||||
other => panic!("Expected Success, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_process_image_for_agent_preserves_transparency_as_png() {
|
||||
let transparent_png = create_noisy_png(1000, 1000, true);
|
||||
assert!(transparent_png.len() > MAX_IMAGE_SIZE_BYTES);
|
||||
|
||||
let result = process_image_for_agent(&transparent_png);
|
||||
|
||||
match result {
|
||||
ProcessImageResult::Success { data, mime_type } => {
|
||||
assert!(data.len() <= MAX_IMAGE_SIZE_BYTES);
|
||||
assert_eq!(mime_type, "image/png");
|
||||
assert!(image::load_from_memory(&data).unwrap().color().has_alpha());
|
||||
}
|
||||
other => panic!("Expected Success, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_process_image_for_agent_invalid_data() {
|
||||
let invalid_data = vec![0u8; 100];
|
||||
@@ -98,6 +160,24 @@ fn test_resize_image_large_image_gets_resized() {
|
||||
assert!(new_width < width || new_height < height);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_process_image_for_agent_keeps_resized_opaque_png_when_it_fits() {
|
||||
let large_png = create_large_test_png();
|
||||
|
||||
let result = process_image_for_agent(&large_png);
|
||||
|
||||
match result {
|
||||
ProcessImageResult::Success { data, mime_type } => {
|
||||
assert!(data.len() <= MAX_IMAGE_SIZE_BYTES);
|
||||
assert_eq!(mime_type, "image/png");
|
||||
let img = image::load_from_memory(&data).unwrap();
|
||||
let (width, height) = img.dimensions();
|
||||
assert!((width as f64 * height as f64) <= MAX_IMAGE_PIXELS);
|
||||
}
|
||||
other => panic!("Expected Success, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates a very tall/narrow image to test dimension clamping.
|
||||
fn create_tall_test_png() -> Vec<u8> {
|
||||
// Create a 100x20000 image (2M pixels, but very tall)
|
||||
|
||||
Reference in New Issue
Block a user