| 1 | use super::*; |
| 2 | use codewhale_models::Role; |
| 3 | |
| 4 | /// A 1x1 PNG, as bytes rather than a fixture file so the encoding tests |
| 5 | /// have no filesystem dependency. |
| 6 | pub(crate) const PNG_1X1: &[u8] = &[ |
| 7 | 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00, 0x00, 0x0d, 0x49, 0x48, 0x44, 0x52, |
| 8 | 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x06, 0x00, 0x00, 0x00, 0x1f, 0x15, 0xc4, |
| 9 | 0x89, 0x00, 0x00, 0x00, 0x0a, 0x49, 0x44, 0x41, 0x54, 0x78, 0x9c, 0x63, 0x00, 0x01, 0x00, 0x00, |
| 10 | 0x05, 0x00, 0x01, 0x0d, 0x0a, 0x2d, 0xb4, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4e, 0x44, 0xae, |
| 11 | 0x42, 0x60, 0x82, |
| 12 | ]; |
| 13 | |
| 14 | #[test] |
| 15 | fn sniffs_every_accepted_format_from_magic_bytes() { |
| 16 | assert_eq!(sniff_media_type(PNG_1X1), Some("image/png")); |
| 17 | assert_eq!( |
| 18 | sniff_media_type(&[0xff, 0xd8, 0xff, 0xe0, 0x00]), |
| 19 | Some("image/jpeg") |
| 20 | ); |
| 21 | assert_eq!(sniff_media_type(b"GIF89a....."), Some("image/gif")); |
| 22 | assert_eq!(sniff_media_type(b"GIF87a....."), Some("image/gif")); |
| 23 | assert_eq!( |
| 24 | sniff_media_type(b"RIFF\x00\x00\x00\x00WEBPVP8 "), |
| 25 | Some("image/webp") |
| 26 | ); |
| 27 | } |
| 28 | |
| 29 | /// Noise defeats compression, so the shrink ladder has to re-encode for real. |
| 30 | fn noise_payload(width: u32, height: u32) -> Vec<u8> { |
| 31 | use image::ImageEncoder as _; |
| 32 | let mut pixels = image::RgbImage::new(width, height); |
| 33 | for (x, y, pixel) in pixels.enumerate_pixels_mut() { |
| 34 | *pixel = image::Rgb([ |
| 35 | (x.wrapping_mul(31) ^ y.wrapping_mul(17)) as u8, |
| 36 | (x.wrapping_mul(7) ^ y.wrapping_mul(29)) as u8, |
| 37 | (x.wrapping_add(y).wrapping_mul(13)) as u8, |
| 38 | ]); |
| 39 | } |
| 40 | let mut bytes = Vec::new(); |
| 41 | image::codecs::png::PngEncoder::new_with_quality( |
| 42 | &mut bytes, |
| 43 | image::codecs::png::CompressionType::Fast, |
| 44 | image::codecs::png::FilterType::NoFilter, |
| 45 | ) |
| 46 | .write_image( |
| 47 | pixels.as_raw(), |
| 48 | width, |
| 49 | height, |
| 50 | image::ExtendedColorType::Rgb8, |
| 51 | ) |
| 52 | .expect("encode fixture png"); |
| 53 | bytes |
| 54 | } |
| 55 | |
| 56 | fn image_blocks_fixture() -> Vec<codewhale_models::Message> { |
| 57 | let payload = STANDARD.encode(noise_payload(320, 320)); |
| 58 | vec![ |
| 59 | codewhale_models::Message { |
| 60 | role: Role::User, |
| 61 | content: vec![ |
| 62 | ContentBlock::Text { |
| 63 | text: "look at this".to_string(), |
| 64 | cache_control: None, |
| 65 | }, |
| 66 | ContentBlock::ImageUrl { |
| 67 | image_url: ImageUrlContent { |
| 68 | url: format!("data:image/png;base64,{payload}"), |
| 69 | }, |
| 70 | }, |
| 71 | ], |
| 72 | }, |
| 73 | codewhale_models::Message { |
| 74 | role: Role::User, |
| 75 | content: vec![ContentBlock::ToolResult { |
| 76 | execution_id: None, |
| 77 | tool_use_id: "call_1".to_string(), |
| 78 | content: "Read image file [image/png]".to_string(), |
| 79 | is_error: None, |
| 80 | content_blocks: Some(vec![serde_json::json!({ |
| 81 | "type": "image", |
| 82 | "mime_type": "image/png", |
| 83 | "data": payload, |
| 84 | })]), |
| 85 | }], |
| 86 | }, |
| 87 | ] |
| 88 | } |
| 89 | |
| 90 | #[test] |
| 91 | fn compaction_shrink_rewrites_both_image_carriers_under_budget() { |
| 92 | let budget = 64 * 1024; |
| 93 | let mut messages = image_blocks_fixture(); |
| 94 | let outcome = shrink_images_for_request_with_budget(&mut messages, budget); |
| 95 | assert_eq!(outcome.images, 2, "both carriers are rewritten"); |
| 96 | assert!(outcome.bytes_after < outcome.bytes_before); |
| 97 | |
| 98 | let ContentBlock::ImageUrl { image_url } = &messages[0].content[1] else { |
| 99 | panic!("image block survives the shrink"); |
| 100 | }; |
| 101 | let (mime, payload) = parse_data_url(&image_url.url).expect("data url"); |
| 102 | assert!(matches!(mime, "image/png" | "image/jpeg"), "{mime}"); |
| 103 | let bytes = STANDARD.decode(payload).expect("base64"); |
| 104 | assert!(bytes.len() <= budget, "{} <= {budget}", bytes.len()); |
| 105 | assert_eq!(sniff_media_type(&bytes), Some(mime)); |
| 106 | |
| 107 | let ContentBlock::ToolResult { content_blocks, .. } = &messages[1].content[0] else { |
| 108 | panic!("tool result survives the shrink"); |
| 109 | }; |
| 110 | let block = &content_blocks.as_ref().expect("blocks")[0]; |
| 111 | assert_eq!(block["type"], "image"); |
| 112 | let nested = STANDARD |
| 113 | .decode(block["data"].as_str().expect("data")) |
| 114 | .expect("nested base64"); |
| 115 | assert!(nested.len() <= budget); |
| 116 | assert_eq!( |
| 117 | sniff_media_type(&nested), |
| 118 | block["mime_type"].as_str(), |
| 119 | "the mime must match the rewritten bytes" |
| 120 | ); |
| 121 | } |
| 122 | |
| 123 | #[test] |
| 124 | fn compaction_shrink_leaves_an_under_budget_history_untouched() { |
| 125 | let mut messages = image_blocks_fixture(); |
| 126 | let before = messages.clone(); |
| 127 | let outcome = shrink_images_for_request_with_budget(&mut messages, 32 * 1024 * 1024); |
| 128 | assert_eq!(outcome.images, 0, "a fitting request is never rewritten"); |
| 129 | assert_eq!( |
| 130 | outcome.images_seen, 2, |
| 131 | "the request still carried images, and the ladder must be able to tell" |
| 132 | ); |
| 133 | assert_eq!(messages, before); |
| 134 | } |
| 135 | |
| 136 | #[test] |
| 137 | fn compaction_placeholder_keeps_the_image_visible_as_text() { |
| 138 | let mut messages = image_blocks_fixture(); |
| 139 | let replaced = replace_images_with_placeholders( |
| 140 | &mut messages, |
| 141 | "the summary request exceeded the provider's request-body limit (HTTP 413)", |
| 142 | ); |
| 143 | assert_eq!(replaced, 2); |
| 144 | |
| 145 | let ContentBlock::Text { text, .. } = &messages[0].content[1] else { |
| 146 | panic!("image_url becomes a note"); |
| 147 | }; |
| 148 | assert!(text.contains("omitted from this summary pass"), "{text}"); |
| 149 | assert!(text.contains("do not describe"), "{text}"); |
| 150 | |
| 151 | let ContentBlock::ToolResult { |
| 152 | content, |
| 153 | content_blocks, |
| 154 | .. |
| 155 | } = &messages[1].content[0] |
| 156 | else { |
| 157 | panic!("tool result survives"); |
| 158 | }; |
| 159 | assert!(content_blocks.is_none(), "no base64 may remain"); |
| 160 | assert!(content.contains("1 image(s)"), "{content}"); |
| 161 | assert!(content.contains("Read image file [image/png]"), "{content}"); |
| 162 | } |
| 163 | |
| 164 | #[test] |
| 165 | fn sniffing_ignores_the_extension_and_believes_the_bytes() { |
| 166 | // A JPEG named .png must be declared image/jpeg, or the provider |
| 167 | // rejects the media-type mismatch. |
| 168 | let jpeg = [0xff, 0xd8, 0xff, 0xe0, 0x11, 0x22]; |
| 169 | let attached = encode_image_bytes(&jpeg, "screenshot.png").expect("attach"); |
| 170 | assert_eq!(attached.media_type, "image/jpeg"); |
| 171 | assert!(attached.data_url.starts_with("data:image/jpeg;base64,")); |
| 172 | } |
| 173 | |
| 174 | #[test] |
| 175 | fn rich_tool_images_reject_corruption_mime_spoofing_and_decode_bombs() { |
| 176 | use crate::tools::spec::{RichToolResult, ToolResult}; |
| 177 | use codewhale_tools::ToolResultContentBlock; |
| 178 | let mut oversized_header = PNG_1X1.to_vec(); |
| 179 | oversized_header[16..20].copy_from_slice(&(MAX_IMAGE_DIMENSION + 1).to_be_bytes()); |
| 180 | for (mime, bytes) in [ |
| 181 | ("image/png", &PNG_1X1[..8]), |
| 182 | ("image/jpeg", PNG_1X1), |
| 183 | ("image/png", oversized_header.as_slice()), |
| 184 | ] { |
| 185 | assert!(prepare_tool_image_bytes(bytes, mime).block.is_none()); |
| 186 | let rich = bound_rich_tool_result(RichToolResult::with_content_blocks( |
| 187 | ToolResult::success("capture receipt"), |
| 188 | vec![ToolResultContentBlock::Image { |
| 189 | mime_type: mime.into(), |
| 190 | data: STANDARD.encode(bytes), |
| 191 | }], |
| 192 | )); |
| 193 | assert!(rich.result.success); |
| 194 | assert!(rich.content_blocks.is_empty()); |
| 195 | assert!(rich.result.content.starts_with("capture receipt")); |
| 196 | assert!( |
| 197 | rich.result |
| 198 | .content |
| 199 | .contains("1 tool-result image block(s) omitted") |
| 200 | ); |
| 201 | let stored = vec![ |
| 202 | serde_json::json!({"type":"image","mime_type":mime,"data":STANDARD.encode(bytes)}), |
| 203 | ]; |
| 204 | let (image, omitted) = provider_tool_result_image_refs(Some(&stored)); |
| 205 | assert!( |
| 206 | image.is_none(), |
| 207 | "restored history must not bypass validation" |
| 208 | ); |
| 209 | assert_eq!(omitted, 1); |
| 210 | } |
| 211 | } |
| 212 | |
| 213 | #[test] |
| 214 | fn lowercase_read_image_preparation_is_typed_and_bounded() { |
| 215 | let prepared = prepare_tool_image_bytes(PNG_1X1, "image/png"); |
| 216 | assert_eq!(prepared.note, "Read image file [image/png]"); |
| 217 | let codewhale_tools::ToolResultContentBlock::Image { mime_type, data } = |
| 218 | prepared.block.expect("typed image"); |
| 219 | assert_eq!(mime_type, "image/png"); |
| 220 | assert_eq!(STANDARD.decode(data).expect("base64"), PNG_1X1); |
| 221 | |
| 222 | let omitted = prepare_tool_image_bytes(b"BMnot-a-safe-bitmap", "image/bmp"); |
| 223 | assert!(omitted.block.is_none()); |
| 224 | assert!(omitted.note.contains("Image omitted"), "{}", omitted.note); |
| 225 | } |
| 226 | |
| 227 | #[test] |
| 228 | fn blind_route_removes_nested_tool_result_image() { |
| 229 | let mut messages = vec![codewhale_models::Message { |
| 230 | role: Role::User, |
| 231 | content: vec![ContentBlock::ToolResult { |
| 232 | execution_id: None, |
| 233 | tool_use_id: "call-image".to_string(), |
| 234 | content: "Read image file [image/png]".to_string(), |
| 235 | is_error: None, |
| 236 | content_blocks: Some(vec![serde_json::json!({ |
| 237 | "type": "image", |
| 238 | "mime_type": "image/png", |
| 239 | "data": "QUJD", |
| 240 | })]), |
| 241 | }], |
| 242 | }]; |
| 243 | |
| 244 | assert_eq!( |
| 245 | strip_images_when_unsupported(&mut messages, SupportState::Unsupported, "text-only",), |
| 246 | 1 |
| 247 | ); |
| 248 | let ContentBlock::ToolResult { |
| 249 | content, |
| 250 | content_blocks, |
| 251 | .. |
| 252 | } = &messages[0].content[0] |
| 253 | else { |
| 254 | panic!("tool result") |
| 255 | }; |
| 256 | assert!(content_blocks.is_none()); |
| 257 | assert!(content.contains("text-only"), "{content}"); |
| 258 | } |
| 259 | |
| 260 | #[test] |
| 261 | fn copy_projection_never_contains_inline_image_bytes() { |
| 262 | const SENTINEL: &str = "U0VOU0lUSVZFX0JBU0U2NA=="; |
| 263 | let projected = safe_tool_result_content_blocks(Some(&[serde_json::json!({ |
| 264 | "type": "image", |
| 265 | "mime_type": "image/png", |
| 266 | "data": SENTINEL, |
| 267 | })])) |
| 268 | .expect("projection"); |
| 269 | let encoded = serde_json::to_string(&projected).expect("json"); |
| 270 | assert!(!encoded.contains(SENTINEL), "{encoded}"); |
| 271 | assert!( |
| 272 | encoded.contains("inline_or_local_image_payload"), |
| 273 | "{encoded}" |
| 274 | ); |
| 275 | } |
| 276 | |
| 277 | #[test] |
| 278 | fn riff_that_is_not_webp_is_not_an_image() { |
| 279 | // A WAV file is also RIFF. Matching on "RIFF" alone would attach audio. |
| 280 | assert_eq!(sniff_media_type(b"RIFF\x00\x00\x00\x00WAVEfmt "), None); |
| 281 | } |
| 282 | |
| 283 | #[test] |
| 284 | fn encodes_a_png_to_a_data_url_that_round_trips() { |
| 285 | let attached = encode_image_bytes(PNG_1X1, "shot.png").expect("attach"); |
| 286 | assert_eq!(attached.media_type, "image/png"); |
| 287 | assert_eq!(attached.source_bytes, PNG_1X1.len()); |
| 288 | |
| 289 | let (media_type, payload) = parse_data_url(&attached.data_url).expect("parse"); |
| 290 | assert_eq!(media_type, "image/png"); |
| 291 | assert_eq!(STANDARD.decode(payload).expect("decode"), PNG_1X1); |
| 292 | } |
| 293 | |
| 294 | #[test] |
| 295 | fn rejects_a_file_over_the_size_limit() { |
| 296 | let oversized = vec![0u8; MAX_IMAGE_BYTES + 1]; |
| 297 | let error = encode_image_bytes(&oversized, "huge.png").expect_err("must reject"); |
| 298 | assert!( |
| 299 | matches!(error, ImageAttachError::TooLarge { .. }), |
| 300 | "got {error:?}" |
| 301 | ); |
| 302 | let rendered = error.to_string(); |
| 303 | assert!(rendered.contains("5.0 MB"), "{rendered}"); |
| 304 | assert!(rendered.contains("huge.png"), "{rendered}"); |
| 305 | } |
| 306 | |
| 307 | #[test] |
| 308 | fn accepts_a_file_exactly_at_the_size_limit() { |
| 309 | // The boundary is inclusive; an off-by-one here would reject images |
| 310 | // the providers accept. |
| 311 | let mut at_limit = PNG_1X1.to_vec(); |
| 312 | at_limit.resize(MAX_IMAGE_BYTES, 0); |
| 313 | assert!(encode_image_bytes(&at_limit, "edge.png").is_ok()); |
| 314 | } |
| 315 | |
| 316 | #[test] |
| 317 | fn rejects_a_real_image_in_an_unsupported_format_by_name() { |
| 318 | for (bytes, name) in [ |
| 319 | (b"BM\x00\x00\x00\x00".as_slice(), "BMP"), |
| 320 | (b"II\x2a\x00extra".as_slice(), "TIFF"), |
| 321 | (b"MM\x00\x2aextra".as_slice(), "TIFF"), |
| 322 | (b"<svg xmlns=".as_slice(), "SVG"), |
| 323 | (b"%PDF-1.7".as_slice(), "PDF"), |
| 324 | ] { |
| 325 | let error = encode_image_bytes(bytes, "f").expect_err("must reject"); |
| 326 | match error { |
| 327 | ImageAttachError::UnsupportedFormat { detected, .. } => { |
| 328 | assert_eq!(detected, name); |
| 329 | } |
| 330 | other => panic!("expected UnsupportedFormat for {name}, got {other:?}"), |
| 331 | } |
| 332 | } |
| 333 | } |
| 334 | |
| 335 | #[test] |
| 336 | fn rejects_a_file_that_is_not_an_image_at_all() { |
| 337 | let error = encode_image_bytes(b"#!/bin/sh\necho hi\n", "script.png").expect_err("must reject"); |
| 338 | assert!( |
| 339 | matches!(error, ImageAttachError::NotAnImage { .. }), |
| 340 | "got {error:?}" |
| 341 | ); |
| 342 | } |
| 343 | |
| 344 | #[test] |
| 345 | fn rejects_an_empty_file() { |
| 346 | let error = encode_image_bytes(b"", "empty.png").expect_err("must reject"); |
| 347 | assert!( |
| 348 | matches!(error, ImageAttachError::Empty { .. }), |
| 349 | "got {error:?}" |
| 350 | ); |
| 351 | } |
| 352 | |
| 353 | #[test] |
| 354 | fn parses_and_rejects_data_urls() { |
| 355 | assert_eq!( |
| 356 | parse_data_url("data:image/png;base64,QUJD"), |
| 357 | Some(("image/png", "QUJD")) |
| 358 | ); |
| 359 | // Not base64-tagged: Anthropic has no shape for a raw data URL. |
| 360 | assert_eq!(parse_data_url("data:image/png,QUJD"), None); |
| 361 | // Remote URLs are a different source type, not a malformed data URL. |
| 362 | assert_eq!(parse_data_url("https://example.com/a.png"), None); |
| 363 | // Degenerate forms must not produce an empty base64 payload that the |
| 364 | // provider would reject with an opaque error. |
| 365 | assert_eq!(parse_data_url("data:;base64,QUJD"), None); |
| 366 | assert_eq!(parse_data_url("data:image/png;base64,"), None); |
| 367 | assert_eq!(parse_data_url("data:image/png;base64"), None); |
| 368 | } |
| 369 | |
| 370 | #[test] |
| 371 | fn classifies_remote_urls() { |
| 372 | assert!(is_remote_image_url("https://example.com/a.png")); |
| 373 | assert!(is_remote_image_url("http://example.com/a.png")); |
| 374 | assert!(!is_remote_image_url("data:image/png;base64,QUJD")); |
| 375 | assert!(!is_remote_image_url("file:///tmp/a.png")); |
| 376 | } |
| 377 | |
| 378 | fn message_with_image(url: &str) -> codewhale_models::Message { |
| 379 | codewhale_models::Message { |
| 380 | role: Role::User, |
| 381 | content: vec![ |
| 382 | ContentBlock::ImageUrl { |
| 383 | image_url: ImageUrlContent { |
| 384 | url: url.to_string(), |
| 385 | }, |
| 386 | }, |
| 387 | ContentBlock::Text { |
| 388 | text: "what is this?".to_string(), |
| 389 | cache_control: None, |
| 390 | }, |
| 391 | ], |
| 392 | } |
| 393 | } |
| 394 | |
| 395 | fn write_png(dir: &std::path::Path, name: &str) -> std::path::PathBuf { |
| 396 | let path = dir.join(name); |
| 397 | std::fs::write(&path, PNG_1X1).expect("write fixture"); |
| 398 | path |
| 399 | } |
| 400 | |
| 401 | #[test] |
| 402 | fn expands_a_placeholder_into_an_image_block() { |
| 403 | let dir = tempfile::tempdir().expect("tempdir"); |
| 404 | let path = write_png(dir.path(), "shot.png"); |
| 405 | let text = format!("look at this\n[Attached image: {}]", path.display()); |
| 406 | |
| 407 | let expanded = expand_attachment_blocks(&text); |
| 408 | |
| 409 | assert!(expanded.notices.is_empty(), "{expanded:?}"); |
| 410 | // Bracketed: open tag naming the path, the image, close tag. |
| 411 | assert_eq!(expanded.blocks.len(), 3, "{expanded:?}"); |
| 412 | match &expanded.blocks[0] { |
| 413 | ContentBlock::Text { text, .. } => { |
| 414 | assert!(text.starts_with("<image path=\""), "{text}"); |
| 415 | assert!(text.contains("shot.png"), "{text}"); |
| 416 | } |
| 417 | other => panic!("expected an opening tag, got {other:?}"), |
| 418 | } |
| 419 | match &expanded.blocks[1] { |
| 420 | ContentBlock::ImageUrl { image_url } => { |
| 421 | assert!(image_url.url.starts_with("data:image/png;base64,")); |
| 422 | } |
| 423 | other => panic!("expected an image block, got {other:?}"), |
| 424 | } |
| 425 | assert_eq!( |
| 426 | expanded.blocks[2], |
| 427 | ContentBlock::Text { |
| 428 | text: "</image>".to_string(), |
| 429 | cache_control: None |
| 430 | } |
| 431 | ); |
| 432 | } |
| 433 | |
| 434 | #[test] |
| 435 | fn expands_multiple_placeholders_in_order() { |
| 436 | let dir = tempfile::tempdir().expect("tempdir"); |
| 437 | let first = write_png(dir.path(), "one.png"); |
| 438 | let second = write_png(dir.path(), "two.png"); |
| 439 | std::fs::write(&second, [0xff, 0xd8, 0xff, 0xe0, 0x01]).expect("write jpeg"); |
| 440 | let text = format!( |
| 441 | "[Attached image: {}]\nand\n[Attached image: {}]", |
| 442 | first.display(), |
| 443 | second.display() |
| 444 | ); |
| 445 | |
| 446 | let expanded = expand_attachment_blocks(&text); |
| 447 | |
| 448 | let media: Vec<_> = expanded |
| 449 | .blocks |
| 450 | .iter() |
| 451 | .filter_map(|block| match block { |
| 452 | ContentBlock::ImageUrl { image_url } => Some( |
| 453 | parse_data_url(&image_url.url) |
| 454 | .expect("data url") |
| 455 | .0 |
| 456 | .to_string(), |
| 457 | ), |
| 458 | _ => None, |
| 459 | }) |
| 460 | .collect(); |
| 461 | assert_eq!(media, vec!["image/png", "image/jpeg"]); |
| 462 | |
| 463 | // Each image carries its own path tag, so the model can tell two |
| 464 | // screenshots in one turn apart. |
| 465 | let tags: Vec<_> = expanded |
| 466 | .blocks |
| 467 | .iter() |
| 468 | .filter_map(|block| match block { |
| 469 | ContentBlock::Text { text, .. } if text.starts_with("<image path=") => { |
| 470 | Some(text.clone()) |
| 471 | } |
| 472 | _ => None, |
| 473 | }) |
| 474 | .collect(); |
| 475 | assert_eq!(tags.len(), 2, "{tags:?}"); |
| 476 | assert!(tags[0].contains("one.png"), "{tags:?}"); |
| 477 | assert!(tags[1].contains("two.png"), "{tags:?}"); |
| 478 | } |
| 479 | |
| 480 | #[test] |
| 481 | fn ingest_does_not_consult_model_capability() { |
| 482 | // Capability is a route property and is re-decided per request. If |
| 483 | // ingest started gating on it, attaching under a text-only model would |
| 484 | // destroy the image for the rest of the session. |
| 485 | let dir = tempfile::tempdir().expect("tempdir"); |
| 486 | let path = write_png(dir.path(), "shot.png"); |
| 487 | let text = format!("[Attached image: {}]", path.display()); |
| 488 | |
| 489 | let expanded = expand_attachment_blocks(&text); |
| 490 | |
| 491 | assert_eq!(expanded.blocks.len(), 3); |
| 492 | assert!(expanded.notices.is_empty()); |
| 493 | } |
| 494 | |
| 495 | #[test] |
| 496 | fn a_blind_route_gets_text_in_place_of_every_image() { |
| 497 | let mut messages = vec![ |
| 498 | message_with_image("data:image/png;base64,QUJD"), |
| 499 | codewhale_models::Message { |
| 500 | role: Role::Assistant, |
| 501 | content: vec![ContentBlock::Text { |
| 502 | text: "sure".to_string(), |
| 503 | cache_control: None, |
| 504 | }], |
| 505 | }, |
| 506 | ]; |
| 507 | |
| 508 | let stripped = |
| 509 | strip_images_when_unsupported(&mut messages, SupportState::Unsupported, "deepseek-chat"); |
| 510 | |
| 511 | assert_eq!(stripped, 1); |
| 512 | assert!( |
| 513 | !messages[0] |
| 514 | .content |
| 515 | .iter() |
| 516 | .any(|block| matches!(block, ContentBlock::ImageUrl { .. })), |
| 517 | "no image may survive to a route that cannot read one" |
| 518 | ); |
| 519 | match &messages[0].content[0] { |
| 520 | ContentBlock::Text { text, .. } => { |
| 521 | assert!(text.contains("deepseek-chat"), "{text}"); |
| 522 | assert!(text.contains("/model"), "{text}"); |
| 523 | assert!(text.contains("omitted"), "{text}"); |
| 524 | } |
| 525 | other => panic!("expected replacement text, got {other:?}"), |
| 526 | } |
| 527 | } |
| 528 | |
| 529 | #[test] |
| 530 | fn a_supported_or_unknown_route_keeps_its_images() { |
| 531 | // Unknown is the common case: models.dev has no modality data for most |
| 532 | // routes. Stripping there would make the feature dead on arrival for |
| 533 | // self-hosted and custom providers. |
| 534 | for vision in [SupportState::Supported, SupportState::Unknown] { |
| 535 | let mut messages = vec![message_with_image("data:image/png;base64,QUJD")]; |
| 536 | |
| 537 | let stripped = strip_images_when_unsupported(&mut messages, vision, "some-model"); |
| 538 | |
| 539 | assert_eq!(stripped, 0, "{vision:?} must not strip"); |
| 540 | assert!( |
| 541 | messages[0] |
| 542 | .content |
| 543 | .iter() |
| 544 | .any(|block| matches!(block, ContentBlock::ImageUrl { .. })), |
| 545 | "{vision:?} must keep the image" |
| 546 | ); |
| 547 | } |
| 548 | } |
| 549 | |
| 550 | #[test] |
| 551 | fn offline_seed_route_for_a_vision_model_keeps_the_image() { |
| 552 | // #6396: the bundled seed lists Claude as text-only. An offline cold |
| 553 | // start resolves the route from that seed alone; the image must still go. |
| 554 | let route = codewhale_config::route::RouteResolver::new() |
| 555 | .resolve(&codewhale_config::route::RouteRequest { |
| 556 | explicit_provider: Some(codewhale_config::ProviderKind::Anthropic), |
| 557 | model_selector: Some(codewhale_config::route::LogicalModelRef::from( |
| 558 | "claude-opus-5", |
| 559 | )), |
| 560 | saved_provider_model: None, |
| 561 | base_url_override: None, |
| 562 | limit_overrides: Vec::new(), |
| 563 | }) |
| 564 | .expect("bundled Anthropic route resolves offline"); |
| 565 | let mut messages = vec![message_with_image("data:image/png;base64,QUJD")]; |
| 566 | |
| 567 | let stripped = strip_images_when_unsupported( |
| 568 | &mut messages, |
| 569 | route.capabilities().image_input, |
| 570 | "claude-opus-5", |
| 571 | ); |
| 572 | |
| 573 | assert_eq!(stripped, 0); |
| 574 | } |
| 575 | |
| 576 | #[test] |
| 577 | fn stripping_replaces_every_image_across_every_message() { |
| 578 | let mut messages = vec![ |
| 579 | message_with_image("data:image/png;base64,AAAA"), |
| 580 | message_with_image("data:image/jpeg;base64,BBBB"), |
| 581 | ]; |
| 582 | |
| 583 | let stripped = strip_images_when_unsupported(&mut messages, SupportState::Unsupported, "blind"); |
| 584 | |
| 585 | assert_eq!( |
| 586 | stripped, 2, |
| 587 | "a per-message early return would miss the second" |
| 588 | ); |
| 589 | } |
| 590 | |
| 591 | #[test] |
| 592 | fn a_missing_file_becomes_a_notice_not_a_dropped_turn() { |
| 593 | let text = "[Attached image: /nonexistent/definitely-not-here.png]"; |
| 594 | |
| 595 | let expanded = expand_attachment_blocks(text); |
| 596 | |
| 597 | assert!(expanded.blocks.is_empty()); |
| 598 | assert_eq!(expanded.notices.len(), 1); |
| 599 | assert!( |
| 600 | expanded.notices[0].contains("definitely-not-here.png"), |
| 601 | "{:?}", |
| 602 | expanded.notices |
| 603 | ); |
| 604 | } |
| 605 | |
| 606 | #[test] |
| 607 | fn one_bad_attachment_does_not_suppress_a_good_one() { |
| 608 | let dir = tempfile::tempdir().expect("tempdir"); |
| 609 | let good = write_png(dir.path(), "good.png"); |
| 610 | let text = format!( |
| 611 | "[Attached image: /nope/missing.png]\n[Attached image: {}]", |
| 612 | good.display() |
| 613 | ); |
| 614 | |
| 615 | let expanded = expand_attachment_blocks(&text); |
| 616 | |
| 617 | assert_eq!(expanded.blocks.len(), 3); |
| 618 | assert_eq!(expanded.notices.len(), 1); |
| 619 | } |
| 620 | |
| 621 | #[test] |
| 622 | fn video_attachments_are_left_as_text() { |
| 623 | let text = "[Attached video: /tmp/clip.mp4]"; |
| 624 | |
| 625 | let expanded = expand_attachment_blocks(text); |
| 626 | |
| 627 | assert!(expanded.blocks.is_empty(), "{expanded:?}"); |
| 628 | assert!(expanded.notices.is_empty(), "{expanded:?}"); |
| 629 | } |
| 630 | |
| 631 | #[test] |
| 632 | fn text_with_no_attachments_produces_nothing() { |
| 633 | let expanded = expand_attachment_blocks("just a normal question"); |
| 634 | assert!(expanded.blocks.is_empty()); |
| 635 | assert!(expanded.notices.is_empty()); |
| 636 | } |
| 637 | |
| 638 | #[test] |
| 639 | fn notice_block_names_the_failure_and_forbids_guessing() { |
| 640 | assert_eq!(notice_block(&[]), None); |
| 641 | let block = |
| 642 | notice_block(&["Cannot attach a.png: the file is empty".to_string()]).expect("block"); |
| 643 | match block { |
| 644 | ContentBlock::Text { text, .. } => { |
| 645 | assert!(text.contains("<attachment_notice>"), "{text}"); |
| 646 | assert!(text.contains("a.png"), "{text}"); |
| 647 | assert!(text.contains("Do not describe"), "{text}"); |
| 648 | } |
| 649 | other => panic!("expected text, got {other:?}"), |
| 650 | } |
| 651 | } |
| 652 | |
| 653 | #[test] |
| 654 | fn attach_from_path_reports_an_unreadable_file() { |
| 655 | let error = attach_image_from_path(Path::new("/nonexistent/x.png")).expect_err("must fail"); |
| 656 | assert!( |
| 657 | matches!(error, ImageAttachError::Unreadable { .. }), |
| 658 | "got {error:?}" |
| 659 | ); |
| 660 | } |
| 661 | |
| 662 | pub(crate) fn runtime_image_fixture(color: u8) -> codewhale_protocol::runtime::RuntimeImageInput { |
| 663 | let image = image::DynamicImage::ImageRgba8(image::RgbaImage::from_pixel( |
| 664 | 2, |
| 665 | 2, |
| 666 | image::Rgba([color, 31, 99, 255]), |
| 667 | )); |
| 668 | let mut bytes = std::io::Cursor::new(Vec::new()); |
| 669 | image.write_to(&mut bytes, image::ImageFormat::Png).unwrap(); |
| 670 | codewhale_protocol::runtime::RuntimeImageInput { |
| 671 | mime: "image/png".into(), |
| 672 | data_base64: STANDARD.encode(bytes.into_inner()), |
| 673 | } |
| 674 | } |
| 675 | |
| 676 | #[test] |
| 677 | fn runtime_image_validation_preserves_exact_bytes_and_rejects_corruption() { |
| 678 | let input = runtime_image_fixture(7); |
| 679 | let blocks = prepare_runtime_images(std::slice::from_ref(&input)).unwrap(); |
| 680 | assert_eq!( |
| 681 | runtime_images_from_blocks(&blocks).unwrap().as_slice(), |
| 682 | std::slice::from_ref(&input) |
| 683 | ); |
| 684 | for bad in [ |
| 685 | codewhale_protocol::runtime::RuntimeImageInput { |
| 686 | mime: "image/jpeg".into(), |
| 687 | ..input.clone() |
| 688 | }, |
| 689 | codewhale_protocol::runtime::RuntimeImageInput { |
| 690 | data_base64: "not base64".into(), |
| 691 | ..input.clone() |
| 692 | }, |
| 693 | codewhale_protocol::runtime::RuntimeImageInput { |
| 694 | data_base64: String::new(), |
| 695 | ..input.clone() |
| 696 | }, |
| 697 | // Existing signature sniffing alone accepted this truncated PNG. |
| 698 | codewhale_protocol::runtime::RuntimeImageInput { |
| 699 | data_base64: STANDARD.encode(b"\x89PNG\r\n\x1a\n"), |
| 700 | ..input.clone() |
| 701 | }, |
| 702 | ] { |
| 703 | assert!(prepare_runtime_images(&[bad]).is_err()); |
| 704 | } |
| 705 | } |
| 706 | |
| 707 | #[test] |
| 708 | fn runtime_image_validation_bounds_count_encoded_size_and_decode_dimensions() { |
| 709 | let input = runtime_image_fixture(7); |
| 710 | assert!(prepare_runtime_images(&vec![input.clone(); 11]).is_err()); |
| 711 | assert!( |
| 712 | prepare_runtime_images(&[codewhale_protocol::runtime::RuntimeImageInput { |
| 713 | data_base64: "A".repeat(MAX_IMAGE_BYTES.div_ceil(3) * 4 + 1), |
| 714 | ..input |
| 715 | }]) |
| 716 | .is_err() |
| 717 | ); |
| 718 | let wide = image::DynamicImage::ImageRgba8(image::RgbaImage::new(MAX_IMAGE_DIMENSION + 1, 1)); |
| 719 | let mut bytes = std::io::Cursor::new(Vec::new()); |
| 720 | wide.write_to(&mut bytes, image::ImageFormat::Png).unwrap(); |
| 721 | let input = codewhale_protocol::runtime::RuntimeImageInput { |
| 722 | mime: "image/png".into(), |
| 723 | data_base64: STANDARD.encode(bytes.into_inner()), |
| 724 | }; |
| 725 | assert!(prepare_runtime_images(&[input]).is_err()); |
| 726 | } |
| 727 | |
| 728 | pub(crate) fn runtime_image_fixture_bytes( |
| 729 | size: usize, |
| 730 | ) -> codewhale_protocol::runtime::RuntimeImageInput { |
| 731 | let mut image = runtime_image_fixture(42); |
| 732 | let mut bytes = STANDARD.decode(&image.data_base64).unwrap(); |
| 733 | bytes.resize(size, 0); |
| 734 | image.data_base64 = STANDARD.encode(bytes); |
| 735 | image |
| 736 | } |
| 737 | |
| 738 | #[test] |
| 739 | fn runtime_image_network_four_mib_and_historical_five_mib_bounds_are_distinct() { |
| 740 | let at_network_limit = runtime_image_fixture_bytes(MAX_RUNTIME_IMAGE_BYTES); |
| 741 | assert!(prepare_runtime_images(&[at_network_limit]).is_ok()); |
| 742 | let historical = runtime_image_fixture_bytes(MAX_RUNTIME_IMAGE_BYTES + 1); |
| 743 | assert!(prepare_runtime_images(std::slice::from_ref(&historical)).is_err()); |
| 744 | let stored = prepare_stored_images(std::slice::from_ref(&historical)).unwrap(); |
| 745 | assert_eq!(runtime_images_from_blocks(&stored).unwrap(), [historical]); |
| 746 | assert!(prepare_stored_images(&[runtime_image_fixture_bytes(MAX_IMAGE_BYTES)]).is_ok()); |
| 747 | assert!(prepare_stored_images(&[runtime_image_fixture_bytes(MAX_IMAGE_BYTES + 1)]).is_err()); |
| 748 | let three_mib = runtime_image_fixture_bytes(3 * 1024 * 1024); |
| 749 | assert!(prepare_runtime_images(&[three_mib.clone(), three_mib.clone()]).is_err()); |
| 750 | assert!(prepare_stored_images(&[three_mib.clone(), three_mib]).is_ok()); |
| 751 | assert!(prepare_stored_images(&vec![runtime_image_fixture(1); 11]).is_ok()); |
| 752 | } |
| 753 | |
| 754 | fn noise_png(width: u32, height: u32) -> Vec<u8> { |
| 755 | // A cheap xorshift so the PNG does not compress: a stand-in for a busy |
| 756 | // Retina screenshot that exceeds the 5 MiB inline limit as PNG. |
| 757 | let mut state = 0x2545_f491_u32; |
| 758 | let buffer = image::RgbImage::from_fn(width, height, |_, _| { |
| 759 | state ^= state << 13; |
| 760 | state ^= state >> 17; |
| 761 | state ^= state << 5; |
| 762 | // Low-amplitude gray noise: too busy for PNG, like screenshot |
| 763 | // texture, yet flat enough to stay on the PNG rungs once fitted. |
| 764 | let level = 96 + (state & 15) as u8; |
| 765 | image::Rgb([level, level, level]) |
| 766 | }); |
| 767 | let mut bytes = Vec::new(); |
| 768 | image::DynamicImage::ImageRgb8(buffer) |
| 769 | .write_to( |
| 770 | &mut std::io::Cursor::new(&mut bytes), |
| 771 | image::ImageFormat::Png, |
| 772 | ) |
| 773 | .expect("encode noise png"); |
| 774 | bytes |
| 775 | } |
| 776 | |
| 777 | fn attached_dimensions(attached: &AttachedImage) -> (u32, u32) { |
| 778 | let (_, payload) = parse_data_url(&attached.data_url).expect("data url"); |
| 779 | let bytes = STANDARD.decode(payload).expect("base64"); |
| 780 | image::load_from_memory(&bytes) |
| 781 | .expect("decodable attachment") |
| 782 | .to_rgb8() |
| 783 | .dimensions() |
| 784 | } |
| 785 | |
| 786 | #[test] |
| 787 | fn oversized_screenshot_is_downscaled_at_attach_time() { |
| 788 | let png = noise_png(2880, 1800); |
| 789 | assert!( |
| 790 | png.len() > MAX_IMAGE_BYTES, |
| 791 | "fixture must exceed the inline limit" |
| 792 | ); |
| 793 | let dir = tempfile::tempdir().expect("tempdir"); |
| 794 | let path = dir.path().join("retina.png"); |
| 795 | std::fs::write(&path, &png).expect("write fixture"); |
| 796 | |
| 797 | let attached = attach_image_from_path(&path).expect("large screenshot attaches"); |
| 798 | assert!(attached.source_bytes <= MAX_IMAGE_BYTES); |
| 799 | let (width, height) = attached_dimensions(&attached); |
| 800 | assert_eq!(width, ATTACH_MAX_EDGE_PX); |
| 801 | assert!( |
| 802 | (1279..=1281).contains(&height), |
| 803 | "aspect ratio is kept: {height}" |
| 804 | ); |
| 805 | } |
| 806 | |
| 807 | #[test] |
| 808 | fn small_file_with_a_long_edge_is_fitted_to_the_attach_edge() { |
| 809 | let wide = image::RgbImage::from_pixel(4000, 200, image::Rgb([30, 30, 30])); |
| 810 | let mut png = Vec::new(); |
| 811 | image::DynamicImage::ImageRgb8(wide) |
| 812 | .write_to(&mut std::io::Cursor::new(&mut png), image::ImageFormat::Png) |
| 813 | .expect("encode"); |
| 814 | assert!(png.len() < MAX_IMAGE_BYTES); |
| 815 | let dir = tempfile::tempdir().expect("tempdir"); |
| 816 | let path = dir.path().join("wide.png"); |
| 817 | std::fs::write(&path, &png).expect("write fixture"); |
| 818 | |
| 819 | let attached = attach_image_from_path(&path).expect("attach"); |
| 820 | assert_eq!(attached.media_type, "image/png", "flat content stays PNG"); |
| 821 | let (width, height) = attached_dimensions(&attached); |
| 822 | assert_eq!(width, ATTACH_MAX_EDGE_PX); |
| 823 | assert!( |
| 824 | (102..=103).contains(&height), |
| 825 | "aspect ratio is kept: {height}" |
| 826 | ); |
| 827 | } |
| 828 | |
| 829 | #[test] |
| 830 | fn small_image_attaches_byte_for_byte() { |
| 831 | let dir = tempfile::tempdir().expect("tempdir"); |
| 832 | let path = dir.path().join("dot.png"); |
| 833 | std::fs::write(&path, PNG_1X1).expect("write fixture"); |
| 834 | let attached = attach_image_from_path(&path).expect("attach"); |
| 835 | let (_, payload) = parse_data_url(&attached.data_url).expect("data url"); |
| 836 | assert_eq!(STANDARD.decode(payload).expect("base64"), PNG_1X1); |
| 837 | } |
| 838 | |
| 839 | #[test] |
| 840 | fn only_images_since_the_latest_prompt_count_as_this_turns() { |
| 841 | let image = || ContentBlock::ImageUrl { |
| 842 | image_url: ImageUrlContent { |
| 843 | url: "data:image/png;base64,AAAA".to_string(), |
| 844 | }, |
| 845 | }; |
| 846 | let text = |text: &str| ContentBlock::Text { |
| 847 | text: text.to_string(), |
| 848 | cache_control: None, |
| 849 | }; |
| 850 | let old_turn = codewhale_models::Message { |
| 851 | role: Role::User, |
| 852 | content: vec![text("earlier screenshot"), image()], |
| 853 | }; |
| 854 | let reply = codewhale_models::Message { |
| 855 | role: Role::Assistant, |
| 856 | content: vec![text("seen")], |
| 857 | }; |
| 858 | let prompt = |content| codewhale_models::Message { |
| 859 | role: Role::User, |
| 860 | content, |
| 861 | }; |
| 862 | let tool_image = codewhale_models::Message { |
| 863 | role: Role::User, |
| 864 | content: vec![ContentBlock::ToolResult { |
| 865 | execution_id: None, |
| 866 | tool_use_id: "call".to_string(), |
| 867 | content: "Read image".to_string(), |
| 868 | is_error: None, |
| 869 | content_blocks: Some(vec![serde_json::json!({"type": "image"})]), |
| 870 | }], |
| 871 | }; |
| 872 | |
| 873 | let history_only = vec![old_turn.clone(), reply.clone(), prompt(vec![text("go on")])]; |
| 874 | assert_eq!(images_since_last_user_prompt(&history_only), 0); |
| 875 | |
| 876 | let fresh = vec![ |
| 877 | old_turn, |
| 878 | reply, |
| 879 | prompt(vec![text("look"), image()]), |
| 880 | tool_image, |
| 881 | ]; |
| 882 | assert_eq!(images_since_last_user_prompt(&fresh), 2); |
| 883 | } |
| 884 |