winx_code_agent/tools/
read_image.rs

1//! Bounded native-image delivery for multimodal MCP clients.
2//!
3//! Small valid images pass through unchanged. Large images are decoded with
4//! explicit resource limits, resized, and encoded as a bounded JPEG before they
5//! enter JSON/base64. A live session also remembers recent content fingerprints
6//! so an accidental repeated call does not resend the same multi-megabyte image.
7
8use std::io::Cursor;
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11
12use base64::{engine::general_purpose, Engine};
13use image::codecs::jpeg::JpegEncoder;
14use image::imageops::FilterType;
15use image::{DynamicImage, ExtendedColorType, ImageFormat, ImageReader, Limits};
16use serde::Serialize;
17use sha2::{Digest, Sha256};
18use tokio::sync::Mutex;
19use tracing::{debug, error, info, instrument};
20
21use crate::errors::{Result, WinxError};
22use crate::state::bash_state::BashState;
23use crate::types::ReadImage;
24use crate::utils::path::{expand_user, validate_path_in_workspace, validate_path_with_roots};
25
26/// Supported MCP image MIME types.
27pub const SUPPORTED_MIME_TYPES: [&str; 4] = ["image/jpeg", "image/png", "image/gif", "image/webp"];
28
29/// Maximum encoded source size accepted from disk. Decoded dimensions and
30/// allocations have independent limits below.
31const MAX_SOURCE_BYTES: u64 = 50 * 1024 * 1024;
32/// Maximum binary image payload before base64 expansion. This turns the 17 MiB
33/// production outlier into a response of at most ~2.7 MiB of base64.
34pub const MAX_DELIVERED_BYTES: usize = 2 * 1024 * 1024;
35/// Multimodal models downsample large inputs anyway; bounding the long edge
36/// preserves useful screenshot detail without shipping wallpaper-sized files.
37pub const MAX_DELIVERED_DIMENSION: u32 = 2_560;
38const MAX_SOURCE_DIMENSION: u32 = 16_384;
39const MAX_SOURCE_PIXELS: u64 = 64 * 1024 * 1024;
40const MAX_DECODE_ALLOC_BYTES: u64 = 320 * 1024 * 1024;
41const MIN_TRANSCODE_DIMENSION: u32 = 512;
42const JPEG_QUALITIES: [u8; 5] = [85, 75, 65, 55, 45];
43
44/// Safe metadata returned alongside native image content.
45#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
46pub struct ReadImageMetadata {
47    pub source_bytes: usize,
48    pub delivered_bytes: usize,
49    pub base64_bytes: usize,
50    pub source_width: u32,
51    pub source_height: u32,
52    pub delivered_width: u32,
53    pub delivered_height: u32,
54    pub source_mime_type: String,
55    pub delivered_mime_type: String,
56    pub transcoded: bool,
57    pub deduplicated: bool,
58    /// A short content fingerprint for correlation, never file contents.
59    pub content_fingerprint: String,
60}
61
62/// Detailed server-facing outcome. Direct library callers retain the historical
63/// (`mime_type`, `base64_data`) wrapper below.
64#[derive(Clone, Debug, Eq, PartialEq)]
65pub enum ReadImageDelivery {
66    Image { mime_type: String, base64_data: String, metadata: ReadImageMetadata },
67    AlreadyDelivered { metadata: ReadImageMetadata },
68}
69
70struct ImageSource {
71    bytes: Vec<u8>,
72    path: PathBuf,
73    format: ImageFormat,
74    mime_type: String,
75    width: u32,
76    height: u32,
77    fingerprint: String,
78    short_fingerprint: String,
79}
80
81fn image_error(path: &Path, message: impl Into<String>) -> WinxError {
82    WinxError::FileAccessError { path: path.to_path_buf(), message: message.into() }
83}
84
85fn supported_mime(format: ImageFormat) -> Option<&'static str> {
86    match format {
87        ImageFormat::Jpeg => Some("image/jpeg"),
88        ImageFormat::Png => Some("image/png"),
89        ImageFormat::Gif => Some("image/gif"),
90        ImageFormat::WebP => Some("image/webp"),
91        _ => None,
92    }
93}
94
95fn decode_limits() -> Limits {
96    let mut limits = Limits::default();
97    limits.max_image_width = Some(MAX_SOURCE_DIMENSION);
98    limits.max_image_height = Some(MAX_SOURCE_DIMENSION);
99    limits.max_alloc = Some(MAX_DECODE_ALLOC_BYTES);
100    limits
101}
102
103fn fingerprint(bytes: &[u8]) -> (String, String) {
104    use std::fmt::Write as _;
105
106    let digest = Sha256::digest(bytes);
107    let mut full = String::with_capacity(digest.len() * 2);
108    for byte in digest {
109        let _ = write!(full, "{byte:02x}");
110    }
111    let short = full[..16].to_string();
112    (full, short)
113}
114
115fn read_image_source_impl(
116    file_path: &str,
117    cwd: &Path,
118    workspace_root: &Path,
119    extra_roots: Option<&[PathBuf]>,
120) -> Result<ImageSource> {
121    debug!(file_path, "reading image source");
122    let expanded = expand_user(file_path);
123    let requested = if Path::new(&expanded).is_absolute() {
124        PathBuf::from(&expanded)
125    } else {
126        cwd.join(&expanded)
127    };
128    let validated = match extra_roots {
129        Some(extra_roots) => validate_path_with_roots(&requested, workspace_root, extra_roots),
130        None => validate_path_in_workspace(&requested, workspace_root),
131    };
132    let path = validated.map_err(|error| WinxError::PathSecurityError {
133        path: requested,
134        message: error.to_string(),
135    })?;
136
137    if !path.is_file() {
138        return Err(image_error(&path, "file does not exist or is not a regular file"));
139    }
140    let size = std::fs::metadata(&path).map_err(|error| {
141        image_error(&path, format!("could not inspect image metadata: {error}"))
142    })?;
143    if size.len() > MAX_SOURCE_BYTES {
144        return Err(WinxError::FileTooLarge { path, size: size.len(), max_size: MAX_SOURCE_BYTES });
145    }
146
147    let bytes = std::fs::read(&path)
148        .map_err(|error| image_error(&path, format!("could not read image: {error}")))?;
149    let actual_size = u64::try_from(bytes.len()).unwrap_or(u64::MAX);
150    if actual_size > MAX_SOURCE_BYTES {
151        return Err(WinxError::FileTooLarge {
152            path,
153            size: actual_size,
154            max_size: MAX_SOURCE_BYTES,
155        });
156    }
157
158    let format = image::guess_format(&bytes)
159        .map_err(|error| image_error(&path, format!("unsupported or malformed image: {error}")))?;
160    let mime_type = supported_mime(format).ok_or_else(|| {
161        image_error(
162            &path,
163            format!("image format {format:?} is not supported; use JPEG, PNG, GIF, or WebP"),
164        )
165    })?;
166    let mut reader = ImageReader::with_format(Cursor::new(bytes.as_slice()), format);
167    reader.limits(decode_limits());
168    let (width, height) = reader
169        .into_dimensions()
170        .map_err(|error| image_error(&path, format!("could not read image dimensions: {error}")))?;
171    let pixels = u64::from(width).saturating_mul(u64::from(height));
172    if pixels > MAX_SOURCE_PIXELS {
173        return Err(image_error(
174            &path,
175            format!(
176                "decoded image would contain {pixels} pixels; the safety limit is \
177                 {MAX_SOURCE_PIXELS} pixels"
178            ),
179        ));
180    }
181
182    let (fingerprint, short_fingerprint) = fingerprint(&bytes);
183    Ok(ImageSource {
184        bytes,
185        path,
186        format,
187        mime_type: mime_type.to_string(),
188        width,
189        height,
190        fingerprint,
191        short_fingerprint,
192    })
193}
194
195fn decode_source(source: &ImageSource) -> Result<DynamicImage> {
196    let mut reader = ImageReader::with_format(Cursor::new(source.bytes.as_slice()), source.format);
197    reader.limits(decode_limits());
198    reader
199        .decode()
200        .map_err(|error| image_error(&source.path, format!("could not decode image: {error}")))
201}
202
203fn flatten_onto_white(image: &DynamicImage) -> image::RgbImage {
204    if !image.color().has_alpha() {
205        return image.to_rgb8();
206    }
207
208    let rgba = image.to_rgba8();
209    let capacity = (rgba.width() as usize).saturating_mul(rgba.height() as usize).saturating_mul(3);
210    let mut rgb = Vec::with_capacity(capacity);
211    for pixel in rgba.pixels() {
212        let alpha = u16::from(pixel[3]);
213        let inverse = 255_u16.saturating_sub(alpha);
214        for channel in &pixel.0[..3] {
215            let blended = (u16::from(*channel) * alpha + 255 * inverse + 127) / 255;
216            rgb.push(u8::try_from(blended).unwrap_or(255));
217        }
218    }
219    image::RgbImage::from_raw(rgba.width(), rgba.height(), rgb)
220        .unwrap_or_else(|| image::RgbImage::new(rgba.width(), rgba.height()))
221}
222
223fn encode_jpeg(image: &image::RgbImage, quality: u8, path: &Path) -> Result<Vec<u8>> {
224    let mut encoded = Vec::new();
225    JpegEncoder::new_with_quality(&mut encoded, quality)
226        .encode(image.as_raw(), image.width(), image.height(), ExtendedColorType::Rgb8)
227        .map_err(|error| image_error(path, format!("could not encode bounded JPEG: {error}")))?;
228    Ok(encoded)
229}
230
231fn metadata_for_duplicate(source: &ImageSource) -> ReadImageMetadata {
232    ReadImageMetadata {
233        source_bytes: source.bytes.len(),
234        delivered_bytes: 0,
235        base64_bytes: 0,
236        source_width: source.width,
237        source_height: source.height,
238        delivered_width: 0,
239        delivered_height: 0,
240        source_mime_type: source.mime_type.clone(),
241        delivered_mime_type: source.mime_type.clone(),
242        transcoded: false,
243        deduplicated: true,
244        content_fingerprint: source.short_fingerprint.clone(),
245    }
246}
247
248fn encode_delivery(source: ImageSource) -> Result<(String, String, ReadImageMetadata)> {
249    let needs_transcode = source.bytes.len() > MAX_DELIVERED_BYTES
250        || source.width > MAX_DELIVERED_DIMENSION
251        || source.height > MAX_DELIVERED_DIMENSION;
252
253    let (mime_type, delivered, delivered_width, delivered_height, transcoded) = if needs_transcode {
254        let decoded = decode_source(&source)?;
255        let mut edge = MAX_DELIVERED_DIMENSION;
256        let bounded = loop {
257            let candidate = if decoded.width() > edge || decoded.height() > edge {
258                decoded.resize(edge, edge, FilterType::Triangle)
259            } else {
260                decoded.clone()
261            };
262            let rgb = flatten_onto_white(&candidate);
263            let mut accepted = None;
264            for quality in JPEG_QUALITIES {
265                let encoded = encode_jpeg(&rgb, quality, &source.path)?;
266                if encoded.len() <= MAX_DELIVERED_BYTES {
267                    accepted = Some((encoded, rgb.width(), rgb.height()));
268                    break;
269                }
270            }
271            if let Some(accepted) = accepted {
272                break accepted;
273            }
274            if edge <= MIN_TRANSCODE_DIMENSION {
275                return Err(image_error(
276                    &source.path,
277                    format!(
278                        "could not fit image within the {MAX_DELIVERED_BYTES}-byte delivery budget"
279                    ),
280                ));
281            }
282            edge = (edge * 3 / 4).max(MIN_TRANSCODE_DIMENSION);
283        };
284        ("image/jpeg".to_string(), bounded.0, bounded.1, bounded.2, true)
285    } else {
286        // `guess_format` and `into_dimensions` validate only the container
287        // header. Decode once before passing original bytes through so a
288        // truncated/corrupt payload is not mislabeled as a usable MCP image.
289        drop(decode_source(&source)?);
290        (source.mime_type.clone(), source.bytes.clone(), source.width, source.height, false)
291    };
292
293    let base64_data = general_purpose::STANDARD.encode(&delivered);
294    let metadata = ReadImageMetadata {
295        source_bytes: source.bytes.len(),
296        delivered_bytes: delivered.len(),
297        base64_bytes: base64_data.len(),
298        source_width: source.width,
299        source_height: source.height,
300        delivered_width,
301        delivered_height,
302        source_mime_type: source.mime_type,
303        delivered_mime_type: mime_type.clone(),
304        transcoded,
305        deduplicated: false,
306        content_fingerprint: source.short_fingerprint,
307    };
308    Ok((mime_type, base64_data, metadata))
309}
310
311#[cfg(test)]
312fn read_image_from_path_with_roots(
313    file_path: &str,
314    cwd: &Path,
315    workspace_root: &Path,
316    extra_roots: &[PathBuf],
317) -> Result<(String, String, ReadImageMetadata)> {
318    let source = read_image_source_impl(file_path, cwd, workspace_root, Some(extra_roots))?;
319    encode_delivery(source)
320}
321
322/// Server-facing image delivery with content-based, per-session deduplication.
323#[instrument(level = "info", skip(bash_state_arc, read_image))]
324pub async fn handle_tool_call_detailed(
325    bash_state_arc: &Arc<Mutex<Option<BashState>>>,
326    read_image: ReadImage,
327) -> Result<ReadImageDelivery> {
328    info!(file_path = %read_image.file_path, force = read_image.force, "ReadImage tool called");
329    let (cwd, workspace_root) = {
330        let guard = bash_state_arc.lock().await;
331        let Some(state) = guard.as_ref() else {
332            error!("BashState not initialized");
333            return Err(WinxError::BashStateNotInitialized);
334        };
335        (state.cwd.clone(), state.workspace_root.clone())
336    };
337
338    let file_path = read_image.file_path;
339    let source = tokio::task::spawn_blocking(move || {
340        read_image_source_impl(&file_path, &cwd, &workspace_root, None)
341    })
342    .await
343    .map_err(|error| {
344        WinxError::CommandExecutionError(format!("ReadImage source worker failed: {error}"))
345    })??;
346
347    if !read_image.force {
348        let already_delivered = {
349            let mut guard = bash_state_arc.lock().await;
350            let state = guard.as_mut().ok_or(WinxError::BashStateNotInitialized)?;
351            state.image_was_delivered(&source.fingerprint)
352        };
353        if already_delivered {
354            return Ok(ReadImageDelivery::AlreadyDelivered {
355                metadata: metadata_for_duplicate(&source),
356            });
357        }
358    }
359
360    let delivered_fingerprint = source.fingerprint.clone();
361    let (mime_type, base64_data, metadata) =
362        tokio::task::spawn_blocking(move || encode_delivery(source)).await.map_err(|error| {
363            WinxError::CommandExecutionError(format!("ReadImage encode worker failed: {error}"))
364        })??;
365    {
366        let mut guard = bash_state_arc.lock().await;
367        let state = guard.as_mut().ok_or(WinxError::BashStateNotInitialized)?;
368        state.record_image_delivery(delivered_fingerprint);
369    }
370    Ok(ReadImageDelivery::Image { mime_type, base64_data, metadata })
371}
372
373/// Backwards-compatible text-free library API. Direct callers explicitly get
374/// an image every time; MCP callers use `handle_tool_call_detailed` so repeat
375/// deliveries can be compacted safely.
376pub async fn handle_tool_call(
377    bash_state_arc: &Arc<Mutex<Option<BashState>>>,
378    mut read_image: ReadImage,
379) -> Result<(String, String)> {
380    read_image.force = true;
381    match handle_tool_call_detailed(bash_state_arc, read_image).await? {
382        ReadImageDelivery::Image { mime_type, base64_data, .. } => Ok((mime_type, base64_data)),
383        ReadImageDelivery::AlreadyDelivered { .. } => Err(WinxError::CommandExecutionError(
384            "forced ReadImage delivery was unexpectedly deduplicated".to_string(),
385        )),
386    }
387}
388
389#[cfg(test)]
390mod tests {
391    #![allow(clippy::unwrap_used)]
392
393    use image::{Rgba, RgbaImage};
394    use tempfile::TempDir;
395
396    use super::*;
397
398    fn write_png(path: &Path, width: u32, height: u32) {
399        let image = RgbaImage::from_pixel(width, height, Rgba([210, 40, 80, 255]));
400        DynamicImage::ImageRgba8(image).save_with_format(path, ImageFormat::Png).unwrap();
401    }
402
403    #[test]
404    fn reads_valid_image_inside_workspace() {
405        let workspace = TempDir::new().unwrap();
406        let image = workspace.path().join("shot.png");
407        write_png(&image, 2, 2);
408        let (mime, base64, metadata) = read_image_from_path_with_roots(
409            image.to_str().unwrap(),
410            workspace.path(),
411            workspace.path(),
412            &[],
413        )
414        .unwrap();
415        assert_eq!(mime, "image/png");
416        assert!(!base64.is_empty(), "expected non-empty base64");
417        assert!(!metadata.transcoded);
418    }
419
420    #[test]
421    fn rejects_image_outside_workspace() {
422        let workspace = TempDir::new().unwrap();
423        let outside = TempDir::new().unwrap();
424        let secret = outside.path().join("secret.png");
425        write_png(&secret, 2, 2);
426        let error = read_image_from_path_with_roots(
427            secret.to_str().unwrap(),
428            workspace.path(),
429            workspace.path(),
430            &[],
431        );
432        assert!(matches!(error, Err(WinxError::PathSecurityError { .. })));
433    }
434
435    #[test]
436    fn large_dimensions_are_transcoded_into_the_delivery_budget() {
437        let workspace = TempDir::new().unwrap();
438        let image = workspace.path().join("wide.png");
439        write_png(&image, MAX_DELIVERED_DIMENSION + 100, 32);
440        let (mime, base64, metadata) = read_image_from_path_with_roots(
441            image.to_str().unwrap(),
442            workspace.path(),
443            workspace.path(),
444            &[],
445        )
446        .unwrap();
447        assert_eq!(mime, "image/jpeg");
448        assert!(metadata.transcoded);
449        assert!(metadata.delivered_bytes <= MAX_DELIVERED_BYTES);
450        assert!(metadata.delivered_width <= MAX_DELIVERED_DIMENSION);
451        assert_eq!(metadata.base64_bytes, base64.len());
452    }
453
454    #[test]
455    fn rejects_non_image_payload_instead_of_claiming_it_is_jpeg() {
456        let workspace = TempDir::new().unwrap();
457        let text = workspace.path().join("payload.txt");
458        std::fs::write(&text, "not an image").unwrap();
459        let error = read_image_from_path_with_roots(
460            text.to_str().unwrap(),
461            workspace.path(),
462            workspace.path(),
463            &[],
464        )
465        .unwrap_err();
466        assert!(error.to_string().contains("unsupported or malformed image"), "{error}");
467    }
468
469    #[test]
470    fn rejects_truncated_image_after_a_valid_header() {
471        let workspace = TempDir::new().unwrap();
472        let image = workspace.path().join("truncated.png");
473        write_png(&image, 32, 32);
474        let mut bytes = std::fs::read(&image).unwrap();
475        bytes.truncate(bytes.len() / 2);
476        std::fs::write(&image, bytes).unwrap();
477
478        let error = read_image_from_path_with_roots(
479            image.to_str().unwrap(),
480            workspace.path(),
481            workspace.path(),
482            &[],
483        )
484        .unwrap_err();
485        assert!(error.to_string().contains("could not decode image"), "{error}");
486    }
487
488    #[tokio::test]
489    async fn repeated_delivery_is_compacted_unless_forced() {
490        let workspace = TempDir::new().unwrap();
491        let image = workspace.path().join("repeat.png");
492        write_png(&image, 2, 2);
493        let mut state = BashState::new();
494        state.cwd = workspace.path().canonicalize().map(crate::utils::path::simplified).unwrap();
495        state.workspace_root = state.cwd.clone();
496        state.current_thread_id = "image-cache".to_string();
497        state.initialized = true;
498        let state = Arc::new(Mutex::new(Some(state)));
499        let request = ReadImage {
500            file_path: image.to_string_lossy().into_owned(),
501            thread_id: "image-cache".to_string(),
502            force: false,
503        };
504
505        assert!(matches!(
506            handle_tool_call_detailed(&state, request.clone()).await.unwrap(),
507            ReadImageDelivery::Image { .. }
508        ));
509        assert!(matches!(
510            handle_tool_call_detailed(&state, request.clone()).await.unwrap(),
511            ReadImageDelivery::AlreadyDelivered { .. }
512        ));
513        assert!(matches!(
514            handle_tool_call_detailed(&state, ReadImage { force: true, ..request }).await.unwrap(),
515            ReadImageDelivery::Image { .. }
516        ));
517    }
518}
winx_code_agent/tools/read_image.rs

winx_code_agent/tools/
read_image.rs
