llm_browser_testkit/
vision.rs

1//! Vision support — screenshot capture, downscaling, and JPEG encoding for
2//! LLM visual assertions.
3//!
4//! When an `assert` step sets `screenshot = true`, the runner captures the
5//! current viewport as PNG, downscales it so its longest edge is at most
6//! [`default_max_dimension`] (the page's `[config]
7//! screenshot_max_dimension` wins when set), and encodes it as a JPEG data
8//! URL. The image is sent to the vision endpoint as an OpenAI-compatible
9//! `image_url` content part alongside the text prompt.
10//!
11//! Downscaling happens in Rust (no page JS, no fragile `canvas` evaluation):
12//! the PNG bytes are decoded with the `image` crate, resized with Lanczos
13//! filtering, and re-encoded as quality-85 JPEG.
14
15use base64::engine::general_purpose::STANDARD;
16use base64::Engine as _;
17use headless_chrome::Tab;
18
19/// Default longest edge (px) of screenshots sent to vision endpoints.
20pub const DEFAULT_MAX_DIMENSION: u32 = 1400;
21/// JPEG quality used for the encoded screenshot.
22const JPEG_QUALITY: u8 = 85;
23
24/// Captures the current viewport and returns a JPEG data URL suitable for
25/// the OpenAI-compatible `image_url` content part.
26///
27/// The image is downscaled so its longest edge is at most
28/// `max_dimension` (no upscaling; `0` disables resizing).
29///
30/// # Errors
31///
32/// Returns a description when the CDP screenshot capture, PNG decode,
33/// resize, or JPEG encode fails.
34pub fn capture_screenshot_data_url(tab: &Tab, max_dimension: u32) -> Result<String, String> {
35    let png = tab
36        .capture_screenshot(
37            headless_chrome::protocol::cdp::Page::CaptureScreenshotFormatOption::Png,
38            None,
39            None,
40            true,
41        )
42        .map_err(|e| format!("screenshot capture failed: {e}"))?;
43
44    let img =
45        image::load_from_memory(&png).map_err(|e| format!("screenshot decode failed: {e}"))?;
46
47    let (width, height) = (img.width(), img.height());
48    let longest = width.max(height);
49    let resized = if max_dimension > 0 && longest > max_dimension {
50        let scale = f64::from(max_dimension) / f64::from(longest);
51        #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
52        let new_width = (f64::from(width) * scale).round().max(1.0) as u32;
53        #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
54        let new_height = (f64::from(height) * scale).round().max(1.0) as u32;
55        img.resize(new_width, new_height, image::imageops::FilterType::Lanczos3)
56    } else {
57        img
58    };
59
60    let mut jpeg = Vec::new();
61    {
62        let mut encoder =
63            image::codecs::jpeg::JpegEncoder::new_with_quality(&mut jpeg, JPEG_QUALITY);
64        encoder
65            .encode_image(&resized)
66            .map_err(|e| format!("screenshot encode failed: {e}"))?;
67    }
68
69    Ok(format!("data:image/jpeg;base64,{}", STANDARD.encode(jpeg)))
70}
71
72#[cfg(test)]
73mod tests {
74    use super::DEFAULT_MAX_DIMENSION;
75
76    #[test]
77    fn default_max_dimension_sane() {
78        // 1280x720 viewports stay untouched; 1920-wide screens are scaled
79        // down to keep vision tokens reasonable.
80        const {
81            assert!(DEFAULT_MAX_DIMENSION >= 1280 && DEFAULT_MAX_DIMENSION <= 1600);
82        }
83    }
84}
llm_browser_testkit/vision.rs

llm_browser_testkit/
vision.rs
