| 1 | //! Turning a local image file into a wire-ready image content block. |
| 2 | //! |
| 3 | //! CodeWhale's message model has been multimodal for a long time — |
| 4 | //! [`ContentBlock::ImageUrl`] round-trips |
| 5 | //! through session persistence, compaction and all three wire builders. What it |
| 6 | //! never had was a *faucet*: nothing outside `#[cfg(test)]` ever constructed |
| 7 | //! one, so a user who attached a screenshot got the literal text |
| 8 | //! `[Attached image: /path/to/shot.png]` and a model that correctly concluded it |
| 9 | //! could not see the picture. |
| 10 | //! |
| 11 | //! This module is that faucet. It reads a file, proves it is an image the |
| 12 | //! providers actually accept, holds it to a size budget, and emits a |
| 13 | //! `data:` URL. |
| 14 | //! |
| 15 | //! # Two stages, because the two failure kinds differ |
| 16 | //! |
| 17 | //! [`expand_attachment_blocks`] runs when the message is built and decides the |
| 18 | //! *permanent* questions: does this file exist, is it really an image, is it |
| 19 | //! small enough. Those answers cannot change, so they are baked into history. |
| 20 | //! |
| 21 | //! [`strip_images_when_unsupported`] runs per outbound request and decides the |
| 22 | //! one *contingent* question: can the model this request is going to actually |
| 23 | //! see images. Routes change mid-session, so answering that at build time |
| 24 | //! would mean attaching a screenshot under a text-only model and losing it |
| 25 | //! permanently, even after switching to a vision model. History keeps the |
| 26 | //! image; each request is normalized against its own route. |
| 27 | //! |
| 28 | //! # Provider neutrality |
| 29 | //! |
| 30 | //! There is exactly one internal representation — a `data:<media-type>;base64,…` |
| 31 | //! URL on `ContentBlock::ImageUrl` — and each wire builder projects it: |
| 32 | //! |
| 33 | //! | wire format | shape | |
| 34 | //! |---|---| |
| 35 | //! | Chat Completions | `{"type":"image_url","image_url":{"url":"data:…"}}` | |
| 36 | //! | Responses | `{"type":"input_image","image_url":"data:…"}` | |
| 37 | //! | Anthropic Messages | `{"type":"image","source":{"type":"base64","media_type":…,"data":…}}` | |
| 38 | //! |
| 39 | //! The Anthropic split lives in [`parse_data_url`], which |
| 40 | //! `client::anthropic` calls. Anthropic is the reason the accepted-format list |
| 41 | //! below is not simply "whatever the `image` crate can decode": it accepts only |
| 42 | //! PNG/JPEG/GIF/WebP, and so, therefore, do we. Refusing a BMP here with a |
| 43 | //! readable message beats letting one through to a provider-side 400. |
| 44 | |
| 45 | use std::path::Path; |
| 46 | |
| 47 | use base64::{Engine as _, engine::general_purpose::STANDARD}; |
| 48 | |
| 49 | use crate::model_profile::SupportState; |
| 50 | use crate::models::{ContentBlock, ImageUrlContent}; |
| 51 | |
| 52 | /// Largest source image accepted, in bytes, before base64 expansion. |
| 53 | /// |
| 54 | /// Base64 inflates by 4/3, so this admits roughly 6.7 MB of request body per |
| 55 | /// image. The number is Anthropic's documented per-image ceiling; keeping the |
| 56 | /// tightest provider limit as the shared limit is what makes a "CodeWhale |
| 57 | /// accepted it" verdict portable across routes. |
| 58 | pub const MAX_IMAGE_BYTES: usize = 5 * 1024 * 1024; |
| 59 | |
| 60 | /// Why a file could not be attached as an image. |
| 61 | /// |
| 62 | /// Every variant renders to a sentence naming the file and the reason. These |
| 63 | /// strings reach both the user (as a command error) and the model (as an |
| 64 | /// in-band notice), so they say what to do next rather than only what failed. |
| 65 | #[derive(Debug, Clone, PartialEq, Eq)] |
| 66 | pub enum ImageAttachError { |
| 67 | /// The file could not be read at all. |
| 68 | Unreadable { path: String, reason: String }, |
| 69 | /// The file is zero bytes. |
| 70 | Empty { path: String }, |
| 71 | /// Over [`MAX_IMAGE_BYTES`]. |
| 72 | TooLarge { path: String, bytes: usize }, |
| 73 | /// Magic bytes identify a format no provider in the set accepts. |
| 74 | UnsupportedFormat { path: String, detected: String }, |
| 75 | /// Magic bytes match nothing we recognize as an image. |
| 76 | NotAnImage { path: String }, |
| 77 | } |
| 78 | |
| 79 | impl std::fmt::Display for ImageAttachError { |
| 80 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 81 | match self { |
| 82 | Self::Unreadable { path, reason } => { |
| 83 | write!(f, "Cannot attach {path}: {reason}") |
| 84 | } |
| 85 | Self::Empty { path } => { |
| 86 | write!(f, "Cannot attach {path}: the file is empty") |
| 87 | } |
| 88 | Self::TooLarge { path, bytes } => write!( |
| 89 | f, |
| 90 | "Cannot attach {path}: {} exceeds the {} per-image limit. \ |
| 91 | Downscale or crop it first.", |
| 92 | human_bytes(*bytes), |
| 93 | human_bytes(MAX_IMAGE_BYTES), |
| 94 | ), |
| 95 | Self::UnsupportedFormat { path, detected } => write!( |
| 96 | f, |
| 97 | "Cannot attach {path}: {detected} is not accepted by vision \ |
| 98 | models. Convert it to PNG, JPEG, GIF or WebP.", |
| 99 | ), |
| 100 | Self::NotAnImage { path } => write!( |
| 101 | f, |
| 102 | "Cannot attach {path}: the file is not a PNG, JPEG, GIF or \ |
| 103 | WebP image (its contents do not match any of those formats).", |
| 104 | ), |
| 105 | } |
| 106 | } |
| 107 | } |
| 108 | |
| 109 | impl std::error::Error for ImageAttachError {} |
| 110 | |
| 111 | fn human_bytes(bytes: usize) -> String { |
| 112 | if bytes >= 1024 * 1024 { |
| 113 | format!("{:.1} MB", bytes as f64 / (1024.0 * 1024.0)) |
| 114 | } else if bytes >= 1024 { |
| 115 | format!("{:.1} KB", bytes as f64 / 1024.0) |
| 116 | } else { |
| 117 | format!("{bytes} bytes") |
| 118 | } |
| 119 | } |
| 120 | |
| 121 | /// An image that is ready to go on the wire. |
| 122 | #[derive(Debug, Clone, PartialEq, Eq)] |
| 123 | pub struct AttachedImage { |
| 124 | /// e.g. `"image/png"`. |
| 125 | pub media_type: &'static str, |
| 126 | /// `data:<media_type>;base64,<payload>`. |
| 127 | pub data_url: String, |
| 128 | /// Size of the source file, before base64. |
| 129 | pub source_bytes: usize, |
| 130 | } |
| 131 | |
| 132 | impl AttachedImage { |
| 133 | /// The content block this image becomes in a message. |
| 134 | #[must_use] |
| 135 | pub fn content_block(&self) -> ContentBlock { |
| 136 | ContentBlock::ImageUrl { |
| 137 | image_url: ImageUrlContent { |
| 138 | url: self.data_url.clone(), |
| 139 | }, |
| 140 | } |
| 141 | } |
| 142 | } |
| 143 | |
| 144 | /// Identify an image format from its leading bytes. |
| 145 | /// |
| 146 | /// Extension sniffing is not enough here: the extension is attacker- and |
| 147 | /// typo-controlled, while what the provider validates is the payload. A |
| 148 | /// `.png` holding JPEG bytes must be declared `image/jpeg` or the request is |
| 149 | /// rejected with a media-type mismatch that reads like a CodeWhale bug. |
| 150 | /// |
| 151 | /// Returns `None` for anything that is not one of the four accepted formats; |
| 152 | /// [`detect_rejected_format`] names the near-misses so the error can be |
| 153 | /// specific. |
| 154 | #[must_use] |
| 155 | pub fn sniff_media_type(bytes: &[u8]) -> Option<&'static str> { |
| 156 | if bytes.starts_with(b"\x89PNG\r\n\x1a\n") { |
| 157 | return Some("image/png"); |
| 158 | } |
| 159 | if bytes.starts_with(b"\xff\xd8\xff") { |
| 160 | return Some("image/jpeg"); |
| 161 | } |
| 162 | if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") { |
| 163 | return Some("image/gif"); |
| 164 | } |
| 165 | if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" { |
| 166 | return Some("image/webp"); |
| 167 | } |
| 168 | None |
| 169 | } |
| 170 | |
| 171 | /// Name a format we can recognize but deliberately refuse. |
| 172 | /// |
| 173 | /// These are real images, so `NotAnImage` would be a lie and would send the |
| 174 | /// user looking for a corrupt file. Naming the format points at the fix |
| 175 | /// (convert it) instead. |
| 176 | #[must_use] |
| 177 | pub fn detect_rejected_format(bytes: &[u8]) -> Option<&'static str> { |
| 178 | if bytes.starts_with(b"BM") { |
| 179 | return Some("BMP"); |
| 180 | } |
| 181 | if bytes.starts_with(b"II\x2a\x00") || bytes.starts_with(b"MM\x00\x2a") { |
| 182 | return Some("TIFF"); |
| 183 | } |
| 184 | if bytes.len() >= 12 && bytes.starts_with(b"\0\0\0") && &bytes[4..8] == b"ftyp" { |
| 185 | return Some("HEIC/AVIF"); |
| 186 | } |
| 187 | if bytes.starts_with(b"<svg") || bytes.starts_with(b"<?xml") { |
| 188 | return Some("SVG"); |
| 189 | } |
| 190 | if bytes.starts_with(b"%PDF") { |
| 191 | return Some("PDF"); |
| 192 | } |
| 193 | None |
| 194 | } |
| 195 | |
| 196 | /// Validate and encode raw bytes that were read from `path`. |
| 197 | /// |
| 198 | /// Split from [`attach_image_from_path`] so the whole policy — order of |
| 199 | /// checks, limits, format verdicts — is testable without touching a |
| 200 | /// filesystem. |
| 201 | pub fn encode_image_bytes(bytes: &[u8], path: &str) -> Result<AttachedImage, ImageAttachError> { |
| 202 | if bytes.is_empty() { |
| 203 | return Err(ImageAttachError::Empty { |
| 204 | path: path.to_string(), |
| 205 | }); |
| 206 | } |
| 207 | // Size is checked before format so a huge file is rejected on the cheap |
| 208 | // fact rather than after we have decided what it is. |
| 209 | if bytes.len() > MAX_IMAGE_BYTES { |
| 210 | return Err(ImageAttachError::TooLarge { |
| 211 | path: path.to_string(), |
| 212 | bytes: bytes.len(), |
| 213 | }); |
| 214 | } |
| 215 | let Some(media_type) = sniff_media_type(bytes) else { |
| 216 | return Err(match detect_rejected_format(bytes) { |
| 217 | Some(detected) => ImageAttachError::UnsupportedFormat { |
| 218 | path: path.to_string(), |
| 219 | detected: detected.to_string(), |
| 220 | }, |
| 221 | None => ImageAttachError::NotAnImage { |
| 222 | path: path.to_string(), |
| 223 | }, |
| 224 | }); |
| 225 | }; |
| 226 | let payload = STANDARD.encode(bytes); |
| 227 | Ok(AttachedImage { |
| 228 | media_type, |
| 229 | data_url: format!("data:{media_type};base64,{payload}"), |
| 230 | source_bytes: bytes.len(), |
| 231 | }) |
| 232 | } |
| 233 | |
| 234 | /// Read, validate and encode an image file. |
| 235 | pub fn attach_image_from_path(path: &Path) -> Result<AttachedImage, ImageAttachError> { |
| 236 | let display = path.display().to_string(); |
| 237 | // Check the size from metadata first so a multi-gigabyte file is refused |
| 238 | // without being read into memory. |
| 239 | if let Ok(meta) = std::fs::metadata(path) { |
| 240 | let len = meta.len(); |
| 241 | if len > MAX_IMAGE_BYTES as u64 { |
| 242 | return Err(ImageAttachError::TooLarge { |
| 243 | path: display, |
| 244 | bytes: usize::try_from(len).unwrap_or(usize::MAX), |
| 245 | }); |
| 246 | } |
| 247 | } |
| 248 | let bytes = std::fs::read(path).map_err(|error| ImageAttachError::Unreadable { |
| 249 | path: display.clone(), |
| 250 | reason: error.to_string(), |
| 251 | })?; |
| 252 | encode_image_bytes(&bytes, &display) |
| 253 | } |
| 254 | |
| 255 | /// Split a `data:<media-type>;base64,<payload>` URL. |
| 256 | /// |
| 257 | /// Anthropic's Messages API models an image as a tagged `source` rather than a |
| 258 | /// URL, so the native route has to take the data URL back apart. Returns |
| 259 | /// `None` for `http(s)` URLs and for anything malformed, which the caller |
| 260 | /// renders as a remote source or a visible degradation respectively. |
| 261 | #[must_use] |
| 262 | pub fn parse_data_url(url: &str) -> Option<(&str, &str)> { |
| 263 | let rest = url.strip_prefix("data:")?; |
| 264 | let (header, payload) = rest.split_once(',')?; |
| 265 | let media_type = header.strip_suffix(";base64")?; |
| 266 | if media_type.is_empty() || payload.is_empty() { |
| 267 | return None; |
| 268 | } |
| 269 | Some((media_type, payload)) |
| 270 | } |
| 271 | |
| 272 | /// Whether a URL is one a provider can fetch for itself. |
| 273 | #[must_use] |
| 274 | pub fn is_remote_image_url(url: &str) -> bool { |
| 275 | url.starts_with("https://") || url.starts_with("http://") |
| 276 | } |
| 277 | |
| 278 | /// The outcome of expanding a user turn's attachment placeholders. |
| 279 | #[derive(Debug, Clone, Default, PartialEq)] |
| 280 | pub struct ExpandedAttachments { |
| 281 | /// Image blocks to append to the user message, in placeholder order. |
| 282 | pub blocks: Vec<ContentBlock>, |
| 283 | /// One line per attachment that could not be sent. These are shown to the |
| 284 | /// user and also handed to the model, because a model that is not told an |
| 285 | /// image was dropped will confidently discuss it from the filename. |
| 286 | pub notices: Vec<String>, |
| 287 | } |
| 288 | |
| 289 | /// Build the image blocks for a user turn from its `[Attached image: …]` lines. |
| 290 | /// |
| 291 | /// This is the ingest half of the composer's placeholder design: the buffer |
| 292 | /// holds a path-bearing text line (which survives editing, history and session |
| 293 | /// reload for free), and the bytes are read here, once, as the message is |
| 294 | /// built. |
| 295 | /// |
| 296 | /// Only *permanent* failures are decided here — a file that is missing, |
| 297 | /// oversized, or not an image will still be all of those things next turn, so |
| 298 | /// baking the verdict into history costs nothing. Whether the *model* can see |
| 299 | /// images is deliberately not decided here; that is contingent on the active |
| 300 | /// route and is re-decided per request by |
| 301 | /// [`strip_images_when_unsupported`]. |
| 302 | /// |
| 303 | /// Each image is bracketed by text tags naming its path. Without them a turn |
| 304 | /// carrying three screenshots gives the model three anonymous images in a row |
| 305 | /// and no way to say which is which. |
| 306 | /// |
| 307 | /// Never returns an error: a turn with one bad attachment should still be |
| 308 | /// sent, with the failure stated in-band rather than swallowed. |
| 309 | #[must_use] |
| 310 | pub fn expand_attachment_blocks(text: &str) -> ExpandedAttachments { |
| 311 | let references = crate::tui::file_mention::media_attachment_references(text); |
| 312 | let mut out = ExpandedAttachments::default(); |
| 313 | for reference in references { |
| 314 | if reference.kind != "image" { |
| 315 | // Video and any future kind: left as the text reference it |
| 316 | // already was. Silently ignoring it here is not a drop — the |
| 317 | // path is still in the prompt, exactly as before this module |
| 318 | // existed. |
| 319 | continue; |
| 320 | } |
| 321 | match attach_image_from_path(Path::new(&reference.path)) { |
| 322 | Ok(image) => { |
| 323 | out.blocks |
| 324 | .push(tag_block(&format!("<image path=\"{}\">", reference.path))); |
| 325 | out.blocks.push(image.content_block()); |
| 326 | out.blocks.push(tag_block("</image>")); |
| 327 | } |
| 328 | Err(error) => out.notices.push(error.to_string()), |
| 329 | } |
| 330 | } |
| 331 | out |
| 332 | } |
| 333 | |
| 334 | fn tag_block(text: &str) -> ContentBlock { |
| 335 | ContentBlock::Text { |
| 336 | text: text.to_string(), |
| 337 | cache_control: None, |
| 338 | } |
| 339 | } |
| 340 | |
| 341 | /// Replace every image in a request with text when the route cannot see them. |
| 342 | /// |
| 343 | /// Capability is a property of the *route*, not of the attachment, and the |
| 344 | /// route changes freely mid-session. Deciding this when the message is built |
| 345 | /// would burn the answer into history: attach a screenshot while a text-only |
| 346 | /// model is selected, switch to a vision model, and the image would be gone |
| 347 | /// for good. So history always keeps the real image and each outbound request |
| 348 | /// is normalized against the model it is actually going to. |
| 349 | /// |
| 350 | /// Only a known `Unsupported` strips. Most routes report `Unknown` because |
| 351 | /// models.dev has no modality data for them, and treating unknown as "no" |
| 352 | /// would make the feature dead on arrival for exactly the self-hosted and |
| 353 | /// custom routes that most need it — so `Unknown` sends the image and lets the |
| 354 | /// provider be the authority. |
| 355 | /// |
| 356 | /// The image is replaced in place rather than removed, so the model is told |
| 357 | /// why it is looking at a gap instead of being left to invent one. |
| 358 | pub fn strip_images_when_unsupported( |
| 359 | messages: &mut [crate::models::Message], |
| 360 | vision: SupportState, |
| 361 | model: &str, |
| 362 | ) -> usize { |
| 363 | if vision != SupportState::Unsupported { |
| 364 | return 0; |
| 365 | } |
| 366 | let mut stripped = 0; |
| 367 | for message in messages.iter_mut() { |
| 368 | for block in &mut message.content { |
| 369 | if matches!(block, ContentBlock::ImageUrl { .. }) { |
| 370 | *block = ContentBlock::Text { |
| 371 | text: format!( |
| 372 | "[image content omitted: the active model ({model}) does \ |
| 373 | not accept image input. Switch to a vision-capable \ |
| 374 | model with /model to see it.]" |
| 375 | ), |
| 376 | cache_control: None, |
| 377 | }; |
| 378 | stripped += 1; |
| 379 | } |
| 380 | } |
| 381 | } |
| 382 | stripped |
| 383 | } |
| 384 | |
| 385 | /// Render dropped-attachment notices as a block the model will read. |
| 386 | /// |
| 387 | /// Wrapped in a tag rather than appended as bare prose so the model can tell |
| 388 | /// the difference between the user saying something and the harness reporting |
| 389 | /// on itself. |
| 390 | #[must_use] |
| 391 | pub fn notice_block(notices: &[String]) -> Option<ContentBlock> { |
| 392 | if notices.is_empty() { |
| 393 | return None; |
| 394 | } |
| 395 | let body = notices.join("\n"); |
| 396 | Some(ContentBlock::Text { |
| 397 | text: format!( |
| 398 | "<attachment_notice>\n{body}\nDo not describe these images from \ |
| 399 | memory or from their filenames; ask the user to re-share them.\n\ |
| 400 | </attachment_notice>" |
| 401 | ), |
| 402 | cache_control: None, |
| 403 | }) |
| 404 | } |
| 405 | |
| 406 | #[cfg(test)] |
| 407 | mod tests { |
| 408 | use super::*; |
| 409 | |
| 410 | /// A 1x1 PNG, as bytes rather than a fixture file so the encoding tests |
| 411 | /// have no filesystem dependency. |
| 412 | const PNG_1X1: &[u8] = &[ |
| 413 | 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00, 0x00, 0x0d, 0x49, 0x48, 0x44, |
| 414 | 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x06, 0x00, 0x00, 0x00, 0x1f, |
| 415 | 0x15, 0xc4, 0x89, 0x00, 0x00, 0x00, 0x0a, 0x49, 0x44, 0x41, 0x54, 0x78, 0x9c, 0x63, 0x00, |
| 416 | 0x01, 0x00, 0x00, 0x05, 0x00, 0x01, 0x0d, 0x0a, 0x2d, 0xb4, 0x00, 0x00, 0x00, 0x00, 0x49, |
| 417 | 0x45, 0x4e, 0x44, 0xae, 0x42, 0x60, 0x82, |
| 418 | ]; |
| 419 | |
| 420 | #[test] |
| 421 | fn sniffs_every_accepted_format_from_magic_bytes() { |
| 422 | assert_eq!(sniff_media_type(PNG_1X1), Some("image/png")); |
| 423 | assert_eq!( |
| 424 | sniff_media_type(&[0xff, 0xd8, 0xff, 0xe0, 0x00]), |
| 425 | Some("image/jpeg") |
| 426 | ); |
| 427 | assert_eq!(sniff_media_type(b"GIF89a....."), Some("image/gif")); |
| 428 | assert_eq!(sniff_media_type(b"GIF87a....."), Some("image/gif")); |
| 429 | assert_eq!( |
| 430 | sniff_media_type(b"RIFF\x00\x00\x00\x00WEBPVP8 "), |
| 431 | Some("image/webp") |
| 432 | ); |
| 433 | } |
| 434 | |
| 435 | #[test] |
| 436 | fn sniffing_ignores_the_extension_and_believes_the_bytes() { |
| 437 | // A JPEG named .png must be declared image/jpeg, or the provider |
| 438 | // rejects the media-type mismatch. |
| 439 | let jpeg = [0xff, 0xd8, 0xff, 0xe0, 0x11, 0x22]; |
| 440 | let attached = encode_image_bytes(&jpeg, "screenshot.png").expect("attach"); |
| 441 | assert_eq!(attached.media_type, "image/jpeg"); |
| 442 | assert!(attached.data_url.starts_with("data:image/jpeg;base64,")); |
| 443 | } |
| 444 | |
| 445 | #[test] |
| 446 | fn riff_that_is_not_webp_is_not_an_image() { |
| 447 | // A WAV file is also RIFF. Matching on "RIFF" alone would attach audio. |
| 448 | assert_eq!(sniff_media_type(b"RIFF\x00\x00\x00\x00WAVEfmt "), None); |
| 449 | } |
| 450 | |
| 451 | #[test] |
| 452 | fn encodes_a_png_to_a_data_url_that_round_trips() { |
| 453 | let attached = encode_image_bytes(PNG_1X1, "shot.png").expect("attach"); |
| 454 | assert_eq!(attached.media_type, "image/png"); |
| 455 | assert_eq!(attached.source_bytes, PNG_1X1.len()); |
| 456 | |
| 457 | let (media_type, payload) = parse_data_url(&attached.data_url).expect("parse"); |
| 458 | assert_eq!(media_type, "image/png"); |
| 459 | assert_eq!(STANDARD.decode(payload).expect("decode"), PNG_1X1); |
| 460 | } |
| 461 | |
| 462 | #[test] |
| 463 | fn rejects_a_file_over_the_size_limit() { |
| 464 | let oversized = vec![0u8; MAX_IMAGE_BYTES + 1]; |
| 465 | let error = encode_image_bytes(&oversized, "huge.png").expect_err("must reject"); |
| 466 | assert!( |
| 467 | matches!(error, ImageAttachError::TooLarge { .. }), |
| 468 | "got {error:?}" |
| 469 | ); |
| 470 | let rendered = error.to_string(); |
| 471 | assert!(rendered.contains("5.0 MB"), "{rendered}"); |
| 472 | assert!(rendered.contains("huge.png"), "{rendered}"); |
| 473 | } |
| 474 | |
| 475 | #[test] |
| 476 | fn accepts_a_file_exactly_at_the_size_limit() { |
| 477 | // The boundary is inclusive; an off-by-one here would reject images |
| 478 | // the providers accept. |
| 479 | let mut at_limit = PNG_1X1.to_vec(); |
| 480 | at_limit.resize(MAX_IMAGE_BYTES, 0); |
| 481 | assert!(encode_image_bytes(&at_limit, "edge.png").is_ok()); |
| 482 | } |
| 483 | |
| 484 | #[test] |
| 485 | fn rejects_a_real_image_in_an_unsupported_format_by_name() { |
| 486 | for (bytes, name) in [ |
| 487 | (b"BM\x00\x00\x00\x00".as_slice(), "BMP"), |
| 488 | (b"II\x2a\x00extra".as_slice(), "TIFF"), |
| 489 | (b"MM\x00\x2aextra".as_slice(), "TIFF"), |
| 490 | (b"<svg xmlns=".as_slice(), "SVG"), |
| 491 | (b"%PDF-1.7".as_slice(), "PDF"), |
| 492 | ] { |
| 493 | let error = encode_image_bytes(bytes, "f").expect_err("must reject"); |
| 494 | match error { |
| 495 | ImageAttachError::UnsupportedFormat { detected, .. } => { |
| 496 | assert_eq!(detected, name); |
| 497 | } |
| 498 | other => panic!("expected UnsupportedFormat for {name}, got {other:?}"), |
| 499 | } |
| 500 | } |
| 501 | } |
| 502 | |
| 503 | #[test] |
| 504 | fn rejects_a_file_that_is_not_an_image_at_all() { |
| 505 | let error = |
| 506 | encode_image_bytes(b"#!/bin/sh\necho hi\n", "script.png").expect_err("must reject"); |
| 507 | assert!( |
| 508 | matches!(error, ImageAttachError::NotAnImage { .. }), |
| 509 | "got {error:?}" |
| 510 | ); |
| 511 | } |
| 512 | |
| 513 | #[test] |
| 514 | fn rejects_an_empty_file() { |
| 515 | let error = encode_image_bytes(b"", "empty.png").expect_err("must reject"); |
| 516 | assert!( |
| 517 | matches!(error, ImageAttachError::Empty { .. }), |
| 518 | "got {error:?}" |
| 519 | ); |
| 520 | } |
| 521 | |
| 522 | #[test] |
| 523 | fn parses_and_rejects_data_urls() { |
| 524 | assert_eq!( |
| 525 | parse_data_url("data:image/png;base64,QUJD"), |
| 526 | Some(("image/png", "QUJD")) |
| 527 | ); |
| 528 | // Not base64-tagged: Anthropic has no shape for a raw data URL. |
| 529 | assert_eq!(parse_data_url("data:image/png,QUJD"), None); |
| 530 | // Remote URLs are a different source type, not a malformed data URL. |
| 531 | assert_eq!(parse_data_url("https://example.com/a.png"), None); |
| 532 | // Degenerate forms must not produce an empty base64 payload that the |
| 533 | // provider would reject with an opaque error. |
| 534 | assert_eq!(parse_data_url("data:;base64,QUJD"), None); |
| 535 | assert_eq!(parse_data_url("data:image/png;base64,"), None); |
| 536 | assert_eq!(parse_data_url("data:image/png;base64"), None); |
| 537 | } |
| 538 | |
| 539 | #[test] |
| 540 | fn classifies_remote_urls() { |
| 541 | assert!(is_remote_image_url("https://example.com/a.png")); |
| 542 | assert!(is_remote_image_url("http://example.com/a.png")); |
| 543 | assert!(!is_remote_image_url("data:image/png;base64,QUJD")); |
| 544 | assert!(!is_remote_image_url("file:///tmp/a.png")); |
| 545 | } |
| 546 | |
| 547 | fn message_with_image(url: &str) -> crate::models::Message { |
| 548 | crate::models::Message { |
| 549 | role: "user".to_string(), |
| 550 | content: vec![ |
| 551 | ContentBlock::ImageUrl { |
| 552 | image_url: ImageUrlContent { |
| 553 | url: url.to_string(), |
| 554 | }, |
| 555 | }, |
| 556 | ContentBlock::Text { |
| 557 | text: "what is this?".to_string(), |
| 558 | cache_control: None, |
| 559 | }, |
| 560 | ], |
| 561 | } |
| 562 | } |
| 563 | |
| 564 | fn write_png(dir: &std::path::Path, name: &str) -> std::path::PathBuf { |
| 565 | let path = dir.join(name); |
| 566 | std::fs::write(&path, PNG_1X1).expect("write fixture"); |
| 567 | path |
| 568 | } |
| 569 | |
| 570 | #[test] |
| 571 | fn expands_a_placeholder_into_an_image_block() { |
| 572 | let dir = tempfile::tempdir().expect("tempdir"); |
| 573 | let path = write_png(dir.path(), "shot.png"); |
| 574 | let text = format!("look at this\n[Attached image: {}]", path.display()); |
| 575 | |
| 576 | let expanded = expand_attachment_blocks(&text); |
| 577 | |
| 578 | assert!(expanded.notices.is_empty(), "{expanded:?}"); |
| 579 | // Bracketed: open tag naming the path, the image, close tag. |
| 580 | assert_eq!(expanded.blocks.len(), 3, "{expanded:?}"); |
| 581 | match &expanded.blocks[0] { |
| 582 | ContentBlock::Text { text, .. } => { |
| 583 | assert!(text.starts_with("<image path=\""), "{text}"); |
| 584 | assert!(text.contains("shot.png"), "{text}"); |
| 585 | } |
| 586 | other => panic!("expected an opening tag, got {other:?}"), |
| 587 | } |
| 588 | match &expanded.blocks[1] { |
| 589 | ContentBlock::ImageUrl { image_url } => { |
| 590 | assert!(image_url.url.starts_with("data:image/png;base64,")); |
| 591 | } |
| 592 | other => panic!("expected an image block, got {other:?}"), |
| 593 | } |
| 594 | assert_eq!( |
| 595 | expanded.blocks[2], |
| 596 | ContentBlock::Text { |
| 597 | text: "</image>".to_string(), |
| 598 | cache_control: None |
| 599 | } |
| 600 | ); |
| 601 | } |
| 602 | |
| 603 | #[test] |
| 604 | fn expands_multiple_placeholders_in_order() { |
| 605 | let dir = tempfile::tempdir().expect("tempdir"); |
| 606 | let first = write_png(dir.path(), "one.png"); |
| 607 | let second = write_png(dir.path(), "two.png"); |
| 608 | std::fs::write(&second, [0xff, 0xd8, 0xff, 0xe0, 0x01]).expect("write jpeg"); |
| 609 | let text = format!( |
| 610 | "[Attached image: {}]\nand\n[Attached image: {}]", |
| 611 | first.display(), |
| 612 | second.display() |
| 613 | ); |
| 614 | |
| 615 | let expanded = expand_attachment_blocks(&text); |
| 616 | |
| 617 | let media: Vec<_> = expanded |
| 618 | .blocks |
| 619 | .iter() |
| 620 | .filter_map(|block| match block { |
| 621 | ContentBlock::ImageUrl { image_url } => Some( |
| 622 | parse_data_url(&image_url.url) |
| 623 | .expect("data url") |
| 624 | .0 |
| 625 | .to_string(), |
| 626 | ), |
| 627 | _ => None, |
| 628 | }) |
| 629 | .collect(); |
| 630 | assert_eq!(media, vec!["image/png", "image/jpeg"]); |
| 631 | |
| 632 | // Each image carries its own path tag, so the model can tell two |
| 633 | // screenshots in one turn apart. |
| 634 | let tags: Vec<_> = expanded |
| 635 | .blocks |
| 636 | .iter() |
| 637 | .filter_map(|block| match block { |
| 638 | ContentBlock::Text { text, .. } if text.starts_with("<image path=") => { |
| 639 | Some(text.clone()) |
| 640 | } |
| 641 | _ => None, |
| 642 | }) |
| 643 | .collect(); |
| 644 | assert_eq!(tags.len(), 2, "{tags:?}"); |
| 645 | assert!(tags[0].contains("one.png"), "{tags:?}"); |
| 646 | assert!(tags[1].contains("two.png"), "{tags:?}"); |
| 647 | } |
| 648 | |
| 649 | #[test] |
| 650 | fn ingest_does_not_consult_model_capability() { |
| 651 | // Capability is a route property and is re-decided per request. If |
| 652 | // ingest started gating on it, attaching under a text-only model would |
| 653 | // destroy the image for the rest of the session. |
| 654 | let dir = tempfile::tempdir().expect("tempdir"); |
| 655 | let path = write_png(dir.path(), "shot.png"); |
| 656 | let text = format!("[Attached image: {}]", path.display()); |
| 657 | |
| 658 | let expanded = expand_attachment_blocks(&text); |
| 659 | |
| 660 | assert_eq!(expanded.blocks.len(), 3); |
| 661 | assert!(expanded.notices.is_empty()); |
| 662 | } |
| 663 | |
| 664 | #[test] |
| 665 | fn a_blind_route_gets_text_in_place_of_every_image() { |
| 666 | let mut messages = vec![ |
| 667 | message_with_image("data:image/png;base64,QUJD"), |
| 668 | crate::models::Message { |
| 669 | role: "assistant".to_string(), |
| 670 | content: vec![ContentBlock::Text { |
| 671 | text: "sure".to_string(), |
| 672 | cache_control: None, |
| 673 | }], |
| 674 | }, |
| 675 | ]; |
| 676 | |
| 677 | let stripped = strip_images_when_unsupported( |
| 678 | &mut messages, |
| 679 | SupportState::Unsupported, |
| 680 | "deepseek-chat", |
| 681 | ); |
| 682 | |
| 683 | assert_eq!(stripped, 1); |
| 684 | assert!( |
| 685 | !messages[0] |
| 686 | .content |
| 687 | .iter() |
| 688 | .any(|block| matches!(block, ContentBlock::ImageUrl { .. })), |
| 689 | "no image may survive to a route that cannot read one" |
| 690 | ); |
| 691 | match &messages[0].content[0] { |
| 692 | ContentBlock::Text { text, .. } => { |
| 693 | assert!(text.contains("deepseek-chat"), "{text}"); |
| 694 | assert!(text.contains("/model"), "{text}"); |
| 695 | assert!(text.contains("omitted"), "{text}"); |
| 696 | } |
| 697 | other => panic!("expected replacement text, got {other:?}"), |
| 698 | } |
| 699 | } |
| 700 | |
| 701 | #[test] |
| 702 | fn a_supported_or_unknown_route_keeps_its_images() { |
| 703 | // Unknown is the common case: models.dev has no modality data for most |
| 704 | // routes. Stripping there would make the feature dead on arrival for |
| 705 | // self-hosted and custom providers. |
| 706 | for vision in [SupportState::Supported, SupportState::Unknown] { |
| 707 | let mut messages = vec![message_with_image("data:image/png;base64,QUJD")]; |
| 708 | |
| 709 | let stripped = strip_images_when_unsupported(&mut messages, vision, "some-model"); |
| 710 | |
| 711 | assert_eq!(stripped, 0, "{vision:?} must not strip"); |
| 712 | assert!( |
| 713 | messages[0] |
| 714 | .content |
| 715 | .iter() |
| 716 | .any(|block| matches!(block, ContentBlock::ImageUrl { .. })), |
| 717 | "{vision:?} must keep the image" |
| 718 | ); |
| 719 | } |
| 720 | } |
| 721 | |
| 722 | #[test] |
| 723 | fn stripping_replaces_every_image_across_every_message() { |
| 724 | let mut messages = vec![ |
| 725 | message_with_image("data:image/png;base64,AAAA"), |
| 726 | message_with_image("data:image/jpeg;base64,BBBB"), |
| 727 | ]; |
| 728 | |
| 729 | let stripped = |
| 730 | strip_images_when_unsupported(&mut messages, SupportState::Unsupported, "blind"); |
| 731 | |
| 732 | assert_eq!( |
| 733 | stripped, 2, |
| 734 | "a per-message early return would miss the second" |
| 735 | ); |
| 736 | } |
| 737 | |
| 738 | #[test] |
| 739 | fn a_missing_file_becomes_a_notice_not_a_dropped_turn() { |
| 740 | let text = "[Attached image: /nonexistent/definitely-not-here.png]"; |
| 741 | |
| 742 | let expanded = expand_attachment_blocks(text); |
| 743 | |
| 744 | assert!(expanded.blocks.is_empty()); |
| 745 | assert_eq!(expanded.notices.len(), 1); |
| 746 | assert!( |
| 747 | expanded.notices[0].contains("definitely-not-here.png"), |
| 748 | "{:?}", |
| 749 | expanded.notices |
| 750 | ); |
| 751 | } |
| 752 | |
| 753 | #[test] |
| 754 | fn one_bad_attachment_does_not_suppress_a_good_one() { |
| 755 | let dir = tempfile::tempdir().expect("tempdir"); |
| 756 | let good = write_png(dir.path(), "good.png"); |
| 757 | let text = format!( |
| 758 | "[Attached image: /nope/missing.png]\n[Attached image: {}]", |
| 759 | good.display() |
| 760 | ); |
| 761 | |
| 762 | let expanded = expand_attachment_blocks(&text); |
| 763 | |
| 764 | assert_eq!(expanded.blocks.len(), 3); |
| 765 | assert_eq!(expanded.notices.len(), 1); |
| 766 | } |
| 767 | |
| 768 | #[test] |
| 769 | fn video_attachments_are_left_as_text() { |
| 770 | let text = "[Attached video: /tmp/clip.mp4]"; |
| 771 | |
| 772 | let expanded = expand_attachment_blocks(text); |
| 773 | |
| 774 | assert!(expanded.blocks.is_empty(), "{expanded:?}"); |
| 775 | assert!(expanded.notices.is_empty(), "{expanded:?}"); |
| 776 | } |
| 777 | |
| 778 | #[test] |
| 779 | fn text_with_no_attachments_produces_nothing() { |
| 780 | let expanded = expand_attachment_blocks("just a normal question"); |
| 781 | assert!(expanded.blocks.is_empty()); |
| 782 | assert!(expanded.notices.is_empty()); |
| 783 | } |
| 784 | |
| 785 | #[test] |
| 786 | fn notice_block_names_the_failure_and_forbids_guessing() { |
| 787 | assert_eq!(notice_block(&[]), None); |
| 788 | let block = |
| 789 | notice_block(&["Cannot attach a.png: the file is empty".to_string()]).expect("block"); |
| 790 | match block { |
| 791 | ContentBlock::Text { text, .. } => { |
| 792 | assert!(text.contains("<attachment_notice>"), "{text}"); |
| 793 | assert!(text.contains("a.png"), "{text}"); |
| 794 | assert!(text.contains("Do not describe"), "{text}"); |
| 795 | } |
| 796 | other => panic!("expected text, got {other:?}"), |
| 797 | } |
| 798 | } |
| 799 | |
| 800 | #[test] |
| 801 | fn attach_from_path_reports_an_unreadable_file() { |
| 802 | let error = attach_image_from_path(Path::new("/nonexistent/x.png")).expect_err("must fail"); |
| 803 | assert!( |
| 804 | matches!(error, ImageAttachError::Unreadable { .. }), |
| 805 | "got {error:?}" |
| 806 | ); |
| 807 | } |
| 808 | } |
| 809 |