Bound agent context and preprocess oversized image attachments

This commit is contained in:
Ryan Ward
2026-09-24 21:54:20 -05:00
parent 6bf81519dc
commit 2a89d03866
15 changed files with 1606 additions and 283 deletions
+187 -37
View File
@@ -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),
}
}
+82 -2
View File
@@ -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)