Skip to main content

agentmux_srv/backend/attachments/
process.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Sniff, decode and derive the two files an attachment needs besides its
5//! original: a small thumbnail for the composer and transcript, and the
6//! "send-copy" the agent actually receives.
7//! SPEC_AGENT_PANE_IMAGE_ATTACHMENTS_2026_09_26.md §6.4.
8//!
9//! Everything here is synchronous and CPU-bound; callers run it on a
10//! blocking thread.
11
12use std::borrow::Cow;
13use std::io::BufWriter;
14use std::path::Path;
15
16use image::codecs::jpeg::JpegEncoder;
17use image::codecs::png::PngEncoder;
18use image::imageops::FilterType;
19use image::{DynamicImage, ImageDecoder, ImageEncoder, ImageFormat, ImageReader, Limits};
20
21/// Largest side accepted before decoding (decompression-bomb guard).
22pub const MAX_SIDE: u32 = 16_384;
23/// Largest pixel count accepted before decoding.
24pub const MAX_PIXELS: u64 = 200_000_000;
25/// Long edge of the thumbnail.
26pub const THUMB_EDGE: u32 = 256;
27/// Upper bound on the send-copy's size. Bedrock and Vertex cap images at
28/// 5 MB; the Anthropic API itself at 10 MB base64.
29pub const SEND_MAX_BYTES: u64 = 5 * 1024 * 1024;
30
31/// What the first bytes of a file say it is. Extensions lie (a HEIC saved as
32/// `.jpg` is common), so every decision below goes through this.
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum Sniffed {
35    Image(ImageFormat),
36    Heic,
37    Avif,
38    Svg,
39    NotImage,
40}
41
42pub fn sniff(head: &[u8]) -> Sniffed {
43    if head.starts_with(&[0x89, b'P', b'N', b'G']) {
44        return Sniffed::Image(ImageFormat::Png);
45    }
46    if head.starts_with(&[0xFF, 0xD8, 0xFF]) {
47        return Sniffed::Image(ImageFormat::Jpeg);
48    }
49    if head.starts_with(b"GIF8") {
50        return Sniffed::Image(ImageFormat::Gif);
51    }
52    if head.len() >= 12 && &head[0..4] == b"RIFF" && &head[8..12] == b"WEBP" {
53        return Sniffed::Image(ImageFormat::WebP);
54    }
55    if head.starts_with(b"BM") && head.len() >= 14 {
56        return Sniffed::Image(ImageFormat::Bmp);
57    }
58    if head.starts_with(b"II*\0") || head.starts_with(b"MM\0*") {
59        return Sniffed::Image(ImageFormat::Tiff);
60    }
61    if head.len() >= 12 && &head[4..8] == b"ftyp" {
62        let brand = &head[8..12];
63        if matches!(
64            brand,
65            b"heic" | b"heix" | b"hevc" | b"heim" | b"heis" | b"hevm" | b"mif1" | b"msf1"
66        ) {
67            return Sniffed::Heic;
68        }
69        if matches!(brand, b"avif" | b"avis") {
70            return Sniffed::Avif;
71        }
72    }
73    let text = String::from_utf8_lossy(&head[..head.len().min(1024)]);
74    let trimmed = text.trim_start_matches('\u{feff}').trim_start();
75    if trimmed.starts_with("<svg") || (trimmed.starts_with("<?xml") && trimmed.contains("<svg")) {
76        return Sniffed::Svg;
77    }
78    Sniffed::NotImage
79}
80
81pub fn mime_of(format: ImageFormat) -> &'static str {
82    match format {
83        ImageFormat::Png => "image/png",
84        ImageFormat::Jpeg => "image/jpeg",
85        ImageFormat::Gif => "image/gif",
86        ImageFormat::WebP => "image/webp",
87        ImageFormat::Bmp => "image/bmp",
88        ImageFormat::Tiff => "image/tiff",
89        _ => "application/octet-stream",
90    }
91}
92
93pub fn ext_of(format: ImageFormat) -> &'static str {
94    match format {
95        ImageFormat::Png => "png",
96        ImageFormat::Jpeg => "jpg",
97        ImageFormat::Gif => "gif",
98        ImageFormat::WebP => "webp",
99        ImageFormat::Bmp => "bmp",
100        ImageFormat::Tiff => "tif",
101        _ => "bin",
102    }
103}
104
105/// Why an image couldn't be processed. `code` values are the
106/// `AttachmentFailedEvent::code` vocabulary.
107#[derive(Debug, Clone, PartialEq, Eq)]
108pub struct ProcessError {
109    pub code: &'static str,
110    pub message: String,
111}
112
113impl ProcessError {
114    fn new(code: &'static str, message: impl Into<String>) -> Self {
115        Self {
116            code,
117            message: message.into(),
118        }
119    }
120}
121
122impl std::fmt::Display for ProcessError {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        f.write_str(&self.message)
125    }
126}
127
128/// One derived file, encoded in memory.
129#[derive(Debug, Clone)]
130pub struct Encoded {
131    pub bytes: Vec<u8>,
132    pub format: ImageFormat,
133    pub width: u32,
134    pub height: u32,
135}
136
137/// Everything the pipeline derives from one original.
138#[derive(Debug, Clone)]
139pub struct Derived {
140    pub width: u32,
141    pub height: u32,
142    pub first_frame_only: bool,
143    pub thumb: Encoded,
144    pub send: Encoded,
145}
146
147/// Width and height from the file header, without decoding pixels. Fails
148/// with `too_large_dimensions` past the bomb guard.
149pub fn header_dimensions(path: &Path) -> Result<(u32, u32), ProcessError> {
150    let (w, h) = ImageReader::open(path)
151        .map_err(|e| ProcessError::new("io", e.to_string()))?
152        .with_guessed_format()
153        .map_err(|e| ProcessError::new("io", e.to_string()))?
154        .into_dimensions()
155        .map_err(|e| ProcessError::new("decode", format!("Couldn't read the image: {e}")))?;
156    check_dimensions(w, h)?;
157    Ok((w, h))
158}
159
160/// Bytes of memory a decode of a `w`×`h` image is budgeted at: the RGBA8
161/// frame plus one working copy (a resize or colour conversion).
162pub fn decode_cost_bytes(w: u32, h: u32) -> u64 {
163    (w as u64) * (h as u64) * 8
164}
165
166/// Decode `path` (already sniffed as `format`) and produce the thumbnail and
167/// the send-copy. `send_max_edge` is the send-copy's long edge.
168pub fn derive(
169    path: &Path,
170    format: ImageFormat,
171    send_max_edge: u32,
172) -> Result<Derived, ProcessError> {
173    // Header-only size check before any pixel is decoded.
174    header_dimensions(path)?;
175
176    let mut reader = ImageReader::open(path).map_err(|e| ProcessError::new("io", e.to_string()))?;
177    reader.set_format(format);
178    let mut limits = Limits::default();
179    limits.max_image_width = Some(MAX_SIDE);
180    limits.max_image_height = Some(MAX_SIDE);
181    // 200 MP of RGBA16 is 1.6 GB; RGBA8 is 800 MB. Allow the 8-bit case.
182    limits.max_alloc = Some(MAX_PIXELS * 4 + 64 * 1024 * 1024);
183    reader.limits(limits);
184    let mut decoder = reader
185        .into_decoder()
186        .map_err(|e| ProcessError::new("decode", format!("Couldn't read the image: {e}")))?;
187    let orientation = decoder.orientation().ok();
188    let mut img = DynamicImage::from_decoder(decoder)
189        .map_err(|e| ProcessError::new("decode", format!("Couldn't decode the image: {e}")))?;
190    if let Some(o) = orientation {
191        img.apply_orientation(o);
192    }
193
194    let first_frame_only = format == ImageFormat::Gif && gif_is_animated(path);
195    let alpha = has_real_alpha(&img);
196    let lossless_source = matches!(
197        format,
198        ImageFormat::Png | ImageFormat::Bmp | ImageFormat::Gif | ImageFormat::Tiff
199    );
200
201    let thumb = encode_thumbnail(&img, alpha)?;
202    let send = encode_send_copy(&img, alpha, lossless_source, send_max_edge.max(64))?;
203
204    Ok(Derived {
205        width: img.width(),
206        height: img.height(),
207        first_frame_only,
208        thumb,
209        send,
210    })
211}
212
213pub fn check_dimensions(w: u32, h: u32) -> Result<(), ProcessError> {
214    if w == 0 || h == 0 {
215        return Err(ProcessError::new("decode", "The image has no pixels."));
216    }
217    if w > MAX_SIDE || h > MAX_SIDE || (w as u64) * (h as u64) > MAX_PIXELS {
218        return Err(ProcessError::new(
219            "too_large_dimensions",
220            format!(
221                "The image is {w}×{h}. The largest supported is {MAX_SIDE} px on a side and {} megapixels.",
222                MAX_PIXELS / 1_000_000
223            ),
224        ));
225    }
226    Ok(())
227}
228
229/// True when some pixel is not fully opaque. Many screenshots carry an alpha
230/// channel that is 255 everywhere; those can still become JPEG.
231fn has_real_alpha(img: &DynamicImage) -> bool {
232    if !img.color().has_alpha() {
233        return false;
234    }
235    // Sample a bounded thumbnail rather than every pixel of a large image.
236    let probe = if img.width().max(img.height()) > 512 {
237        img.thumbnail(512, 512)
238    } else {
239        img.clone()
240    };
241    probe.to_rgba8().pixels().any(|p| p.0[3] < 255)
242}
243
244fn gif_is_animated(path: &Path) -> bool {
245    use image::AnimationDecoder;
246    let Ok(f) = std::fs::File::open(path) else {
247        return false;
248    };
249    let Ok(dec) = image::codecs::gif::GifDecoder::new(std::io::BufReader::new(f)) else {
250        return false;
251    };
252    dec.into_frames().take(2).count() > 1
253}
254
255fn fit(w: u32, h: u32, edge: u32) -> (u32, u32) {
256    let long = w.max(h);
257    if long <= edge {
258        return (w, h);
259    }
260    let scale = edge as f64 / long as f64;
261    (
262        ((w as f64 * scale).round() as u32).max(1),
263        ((h as f64 * scale).round() as u32).max(1),
264    )
265}
266
267fn encode_thumbnail(img: &DynamicImage, alpha: bool) -> Result<Encoded, ProcessError> {
268    let (tw, th) = fit(img.width(), img.height(), THUMB_EDGE);
269    let small: Cow<'_, DynamicImage> = if (tw, th) == (img.width(), img.height()) {
270        Cow::Borrowed(img)
271    } else {
272        Cow::Owned(img.thumbnail(tw, th))
273    };
274    if alpha {
275        encode_png(&small)
276    } else {
277        encode_jpeg(&small, 80)
278    }
279}
280
281fn encode_send_copy(
282    img: &DynamicImage,
283    alpha: bool,
284    lossless_source: bool,
285    max_edge: u32,
286) -> Result<Encoded, ProcessError> {
287    let mut edge = max_edge;
288    // Shrink by 20% per attempt until the file fits. Four attempts take a
289    // 2000 px image down to ~820 px, which fits any realistic content.
290    for _ in 0..5 {
291        let (sw, sh) = fit(img.width(), img.height(), edge);
292        // Borrow at full size: cloning a large decoded frame would double
293        // the memory this job was budgeted for.
294        let scaled: Cow<'_, DynamicImage> = if (sw, sh) == (img.width(), img.height()) {
295            Cow::Borrowed(img)
296        } else {
297            Cow::Owned(img.resize_exact(sw, sh, FilterType::CatmullRom))
298        };
299        let out = if alpha {
300            encode_png(&scaled)?
301        } else if lossless_source {
302            // Screenshots and UI captures stay sharp as PNG; fall back to
303            // JPEG only when PNG is too big (photos saved as PNG).
304            let png = encode_png(&scaled)?;
305            if png.bytes.len() as u64 <= SEND_MAX_BYTES {
306                png
307            } else {
308                encode_jpeg(&scaled, 85)?
309            }
310        } else {
311            encode_jpeg(&scaled, 85)?
312        };
313        if out.bytes.len() as u64 <= SEND_MAX_BYTES {
314            return Ok(out);
315        }
316        edge = ((sw.max(sh) as f64) * 0.8) as u32;
317    }
318    Err(ProcessError::new(
319        "decode",
320        "Couldn't make a small enough copy of this image to send.",
321    ))
322}
323
324fn encode_png(img: &DynamicImage) -> Result<Encoded, ProcessError> {
325    // 16-bit and float images become 8-bit: no model uses the extra depth.
326    let rgba;
327    let rgb;
328    let (buf, color): (&[u8], image::ExtendedColorType) = if img.color().has_alpha() {
329        rgba = img.to_rgba8();
330        (rgba.as_raw(), image::ExtendedColorType::Rgba8)
331    } else {
332        rgb = img.to_rgb8();
333        (rgb.as_raw(), image::ExtendedColorType::Rgb8)
334    };
335    let mut out = Vec::new();
336    PngEncoder::new(BufWriter::new(&mut out))
337        .write_image(buf, img.width(), img.height(), color)
338        .map_err(|e| ProcessError::new("decode", format!("Couldn't encode PNG: {e}")))?;
339    Ok(Encoded {
340        bytes: out,
341        format: ImageFormat::Png,
342        width: img.width(),
343        height: img.height(),
344    })
345}
346
347fn encode_jpeg(img: &DynamicImage, quality: u8) -> Result<Encoded, ProcessError> {
348    let rgb = img.to_rgb8();
349    let mut out = Vec::new();
350    JpegEncoder::new_with_quality(&mut out, quality)
351        .write_image(
352            rgb.as_raw(),
353            img.width(),
354            img.height(),
355            image::ExtendedColorType::Rgb8,
356        )
357        .map_err(|e| ProcessError::new("decode", format!("Couldn't encode JPEG: {e}")))?;
358    Ok(Encoded {
359        bytes: out,
360        format: ImageFormat::Jpeg,
361        width: img.width(),
362        height: img.height(),
363    })
364}
365
366#[cfg(test)]
367mod tests {
368    use super::*;
369    use image::{Rgb, RgbImage, Rgba, RgbaImage};
370
371    fn write(dir: &Path, name: &str, img: &DynamicImage, fmt: ImageFormat) -> std::path::PathBuf {
372        let p = dir.join(name);
373        img.save_with_format(&p, fmt).unwrap();
374        p
375    }
376
377    #[test]
378    fn sniff_recognises_formats_by_content_not_name() {
379        assert_eq!(
380            sniff(&[0x89, b'P', b'N', b'G', 0, 0]),
381            Sniffed::Image(ImageFormat::Png)
382        );
383        assert_eq!(
384            sniff(&[0xFF, 0xD8, 0xFF, 0xE0]),
385            Sniffed::Image(ImageFormat::Jpeg)
386        );
387        assert_eq!(sniff(b"GIF89a......"), Sniffed::Image(ImageFormat::Gif));
388        assert_eq!(
389            sniff(b"RIFF\0\0\0\0WEBPVP8 "),
390            Sniffed::Image(ImageFormat::WebP)
391        );
392        assert_eq!(sniff(b"\0\0\0\x18ftypheic\0\0\0\0"), Sniffed::Heic);
393        assert_eq!(sniff(b"\0\0\0\x18ftypmif1\0\0\0\0"), Sniffed::Heic);
394        assert_eq!(sniff(b"\0\0\0\x1cftypavif\0\0\0\0"), Sniffed::Avif);
395        assert_eq!(sniff(b"\xef\xbb\xbf  <svg xmlns=..."), Sniffed::Svg);
396        assert_eq!(sniff(b"<?xml version=\"1.0\"?><svg>"), Sniffed::Svg);
397        assert_eq!(sniff(b"fn main() {}"), Sniffed::NotImage);
398        assert_eq!(sniff(b""), Sniffed::NotImage);
399    }
400
401    #[test]
402    fn dimension_guard_rejects_bombs_before_decode() {
403        assert!(check_dimensions(4000, 3000).is_ok());
404        assert_eq!(
405            check_dimensions(20_000, 10).unwrap_err().code,
406            "too_large_dimensions"
407        );
408        assert_eq!(
409            check_dimensions(16_000, 16_000).unwrap_err().code,
410            "too_large_dimensions"
411        );
412        assert_eq!(check_dimensions(0, 10).unwrap_err().code, "decode");
413    }
414
415    #[test]
416    fn large_photo_is_downscaled_to_the_send_edge_as_jpeg() {
417        let dir = tempfile::tempdir().unwrap();
418        let img = DynamicImage::ImageRgb8(RgbImage::from_fn(3000, 1500, |x, y| {
419            Rgb([(x % 256) as u8, (y % 256) as u8, ((x + y) % 256) as u8])
420        }));
421        let p = write(dir.path(), "photo.jpg", &img, ImageFormat::Jpeg);
422        let d = derive(&p, ImageFormat::Jpeg, 2000).unwrap();
423        assert_eq!((d.width, d.height), (3000, 1500));
424        assert_eq!(d.send.format, ImageFormat::Jpeg);
425        assert_eq!((d.send.width, d.send.height), (2000, 1000));
426        assert!(d.send.bytes.len() as u64 <= SEND_MAX_BYTES);
427        assert_eq!(d.thumb.width.max(d.thumb.height), THUMB_EDGE);
428    }
429
430    #[test]
431    fn small_screenshot_stays_png_at_full_size() {
432        let dir = tempfile::tempdir().unwrap();
433        let img =
434            DynamicImage::ImageRgba8(RgbaImage::from_pixel(800, 600, Rgba([30, 30, 30, 255])));
435        let p = write(dir.path(), "shot.png", &img, ImageFormat::Png);
436        let d = derive(&p, ImageFormat::Png, 2000).unwrap();
437        assert_eq!(d.send.format, ImageFormat::Png);
438        assert_eq!((d.send.width, d.send.height), (800, 600));
439        // Fully opaque alpha doesn't force a PNG thumbnail.
440        assert_eq!(d.thumb.format, ImageFormat::Jpeg);
441    }
442
443    #[test]
444    fn transparency_is_kept_as_png() {
445        let dir = tempfile::tempdir().unwrap();
446        let img = DynamicImage::ImageRgba8(RgbaImage::from_fn(300, 300, |x, _| {
447            Rgba([255, 0, 0, if x < 150 { 0 } else { 255 }])
448        }));
449        let p = write(dir.path(), "logo.png", &img, ImageFormat::Png);
450        let d = derive(&p, ImageFormat::Png, 2000).unwrap();
451        assert_eq!(d.send.format, ImageFormat::Png);
452        assert_eq!(d.thumb.format, ImageFormat::Png);
453    }
454
455    #[test]
456    fn send_copy_carries_no_exif() {
457        let dir = tempfile::tempdir().unwrap();
458        let img = DynamicImage::ImageRgb8(RgbImage::from_pixel(64, 64, Rgb([1, 2, 3])));
459        let p = write(dir.path(), "a.jpg", &img, ImageFormat::Jpeg);
460        let d = derive(&p, ImageFormat::Jpeg, 2000).unwrap();
461        // An APP1 "Exif" segment would appear right after SOI.
462        assert!(!d.send.bytes.windows(4).any(|w| w == b"Exif"));
463    }
464
465    #[test]
466    fn truncated_file_fails_cleanly() {
467        let dir = tempfile::tempdir().unwrap();
468        let img = DynamicImage::ImageRgb8(RgbImage::from_pixel(200, 200, Rgb([9, 9, 9])));
469        let p = write(dir.path(), "t.png", &img, ImageFormat::Png);
470        let bytes = std::fs::read(&p).unwrap();
471        std::fs::write(&p, &bytes[..bytes.len() / 3]).unwrap();
472        let e = derive(&p, ImageFormat::Png, 2000).unwrap_err();
473        assert_eq!(e.code, "decode");
474    }
475
476    #[test]
477    fn fit_keeps_aspect_and_never_upscales() {
478        assert_eq!(fit(4000, 2000, 2000), (2000, 1000));
479        assert_eq!(fit(100, 50, 2000), (100, 50));
480        assert_eq!(fit(1, 5000, 256), (1, 256));
481    }
482}