diff --git a/.github/capture-benchmark.json b/.github/capture-benchmark.json new file mode 100644 index 0000000..29daf96 --- /dev/null +++ b/.github/capture-benchmark.json @@ -0,0 +1,51 @@ +{ + "run": "https://github.com/aktanazat/iris/actions/runs/35725521554", + "baseline_commit": "227d3109be29c87c2bed85d214c122795a640210", + "candidate_binary_sha256": "1e2b4839764588decaf7d809bdad2c9524e93a2b13991d603897944d763301bb", + "system": "macOS-26.6.2-arm64-arm-64bit-Mach-O", + "chrome": "152.0.7977.83", + "method": "One fresh MCP server per run. Measure from capture request to complete response. First capture is separate from the ten subsequent warm captures. All four commands exited 0.", + "capture": { + "fixture": "tests/fixtures/precise-capture.html", + "selector": ".capture-target", + "padding": 10, + "size": "320x240", + "scale": 1, + "format": "png", + "timeout_seconds": 30, + "decoded_width": 141, + "decoded_height": 81 + }, + "runs": [ + { + "label": "baseline-a", + "first_ms": 5693.806, + "warm_ms": [1369.315, 1080.249, 1059.016, 1252.061, 1137.677, 1054.079, 1261.106, 979.53, 948.367, 717.126] + }, + { + "label": "candidate-a", + "first_ms": 2488.757, + "warm_ms": [547.263, 606.339, 616.116, 618.652, 624.265, 572.997, 528.598, 460.88, 624.339, 689.279] + }, + { + "label": "candidate-b", + "first_ms": 2199.247, + "warm_ms": [653.233, 642.993, 579.069, 536.663, 546.382, 573.03, 635.321, 963.811, 707.962, 558.141] + }, + { + "label": "baseline-b", + "first_ms": 2097.308, + "warm_ms": [1159.602, 897.696, 857.371, 1086.449, 1021.004, 1128.011, 689.528, 1065.104, 920.49, 966.771] + } + ], + "summary": { + "warm_population_per_version": 20, + "baseline_warm_median_ms": 1056.5475, + "candidate_warm_median_ms": 611.2275, + "baseline_warm_p95_ms": 1261.106, + "candidate_warm_p95_ms": 707.962, + "p95_method": "nearest rank over the combined twenty warm samples per version", + "warm_median_reduction_percent": 42.14860193223685, + "limitation": "One fixture on one runner. First-capture times varied and do not establish a reliable startup improvement." + } +} diff --git a/Cargo.lock b/Cargo.lock index 22b606a..0c3f0b5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "aho-corasick" version = "1.1.5" @@ -138,6 +144,18 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bytemuck" +version = "1.25.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" + +[[package]] +name = "byteorder-lite" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f1fe948ff07f4bd06c30984e69f5b4899c516a3ef74f34df92a2df2ab535495" + [[package]] name = "bytes" version = "1.12.1" @@ -299,6 +317,15 @@ dependencies = [ "libc", ] +[[package]] +name = "crc32fast" +version = "1.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01a7799fd6b852db0e61728dde9a204c423b44d689dbd432522543614b490e78" +dependencies = [ + "cfg-if", +] + [[package]] name = "crypto-common" version = "0.1.7" @@ -404,12 +431,32 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "fdeflate" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e6853b52649d4ac5c0bd02320cddc5ba956bdb407c4b75a2c6b75bf51500f8c" +dependencies = [ + "simd-adler32", +] + [[package]] name = "find-msvc-tools" version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d45db016d36b838f563236e9193d0ee6ce38f3f68b6c94e914b4929c96bbb890" +[[package]] +name = "flate2" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" +dependencies = [ + "crc32fast", + "miniz_oxide 0.9.1", + "zlib-rs", +] + [[package]] name = "fnv" version = "1.0.7" @@ -786,6 +833,32 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "image" +version = "0.25.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85ab80394333c02fe689eaf900ab500fbd0c2213da414687ebf995a65d5a6104" +dependencies = [ + "bytemuck", + "byteorder-lite", + "image-webp", + "moxcms", + "num-traits", + "png", + "zune-core", + "zune-jpeg", +] + +[[package]] +name = "image-webp" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "525e9ff3e1a4be2fbea1fdf0e98686a6d98b4d8f937e1bf7402245af1909e8c3" +dependencies = [ + "byteorder-lite", + "quick-error", +] + [[package]] name = "indexmap" version = "2.14.0" @@ -813,6 +886,8 @@ dependencies = [ "chromiumoxide", "clap", "futures", + "image", + "num-traits", "rmcp", "serde", "serde_json", @@ -867,6 +942,26 @@ version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "miniz_oxide" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.2.2" @@ -878,6 +973,16 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "moxcms" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb85c154ba489f01b25c0d36ae69a87e4a1c73a72631fc6c0eb6dde34a73e44b" +dependencies = [ + "num-traits", + "pxfm", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -917,6 +1022,19 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" +[[package]] +name = "png" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60769b8b31b2a9f263dae2776c37b1b28ae246943cf719eb6946a1db05128a61" +dependencies = [ + "bitflags", + "crc32fast", + "fdeflate", + "flate2", + "miniz_oxide 0.8.9", +] + [[package]] name = "potential_utf" version = "0.1.6" @@ -944,6 +1062,18 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "pxfm" +version = "0.1.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d55d956fa96f5ec02be2e13af0e20391a5aa83d6a074e3ad368959d0fab299ea" + +[[package]] +name = "quick-error" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a993555f31e5a609f617c12db6250dedcac1b0a85076912c436e6fc9b2c8e6a3" + [[package]] name = "quote" version = "1.0.47" @@ -1241,6 +1371,12 @@ dependencies = [ "libc", ] +[[package]] +name = "simd-adler32" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" + [[package]] name = "slab" version = "0.4.12" @@ -1876,8 +2012,29 @@ dependencies = [ "syn 3.0.3", ] +[[package]] +name = "zlib-rs" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b268e58e7c693d7c271f93ffc4ba3b380412554231c85bf61ca7af91042a4112" + [[package]] name = "zmij" version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" + +[[package]] +name = "zune-core" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d56377fd46368984a170bc5aac5567e52ca5da874caa60bea39fcbca78fb658b" + +[[package]] +name = "zune-jpeg" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27bc9d5b815bc103f142aa054f561d9187d191692ec7c2d1e2b4737f8dbd7296" +dependencies = [ + "zune-core", +] diff --git a/Cargo.toml b/Cargo.toml index 93396eb..cd27cfc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,8 @@ base64 = "0.23" clap = { version = "4", features = ["derive", "env"] } chromiumoxide = { version = "0.8", default-features = false, features = ["tokio-runtime"] } futures = "0.3" +image = { version = "0.25", default-features = false, features = ["png", "jpeg", "webp"] } +num-traits = "0.2" rmcp = { version = "3.1.2", default-features = false, features = ["macros", "server", "transport-io"] } serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/README.md b/README.md index c1c46b9..879adf1 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ ChatGPT Image Aug 18, 2026, 12_50_00 PM -A camera for coding agents. One fast command or MCP tool call produces one trustworthy image. +A camera for coding agents. Capture pages, compare saved images, and inspect capture readiness without adding a model call. ``` iris example.com # 1440×900 @2x → example.com.png @@ -13,6 +13,9 @@ iris -o shots/ a.com b.com c.com # batch, captured concurrently cat urls.txt | iris - -o shots/ # batch from stdin (# comments ok) iris -o hero.jpg --wait-for 'h1' app.dev iris --selector 'h1' --json app.dev # machine-readable JSON Lines +iris --full --max-pixels 1500000 app.dev # lower density, unchanged page layout +iris --freeze-animations app.dev # finish entrances and pause loops +iris compare before.png after.png -o diff.png ``` `iris --full bridger.to` → @@ -33,7 +36,7 @@ cargo install iris-screenshot The crates.io package is named `iris-screenshot`; the installed command is `iris`. -The only runtime dependency is an installed Chrome-family browser (Chrome, Chromium, Edge, or Brave). +The only runtime dependency is an installed Chrome-family browser (Chrome, Chromium, Edge, Brave, or Helium). For any other Chromium build, set `CHROME=/path/to/browser` or pass `--chrome`. Building from source requires Rust 1.88 or newer. ## Give your coding agent eyes @@ -57,7 +60,13 @@ Or use the equivalent configuration in another MCP client: } ``` -The server exposes one tool, `capture`, which returns the image inline with structured metadata. It writes nothing by default; pass `output` when the agent also needs a file. +The server exposes three tools: + +- `capture` returns one image inline with structured metadata. It writes nothing unless `output` is supplied. +- `capture_batch` accepts a `captures` array of 1–16 capture requests. Up to four run concurrently. Results stay in input order; a failed item does not discard successful images. +- `compare` compares two saved PNG, JPEG, or WebP files. It returns changed-pixel counts and bounds, and can write a PNG difference image with `output`. + +All three run locally. They do not send images, page contents, or reports to an AI service. ```json { @@ -74,7 +83,7 @@ The server exposes one tool, `capture`, which returns the image inline with stru Bare localhost, `.localhost`, and loopback addresses use HTTP automatically. Other bare hosts use HTTPS. Run `iris mcp --help` to select a Chrome binary for the server. -MCP clients may need a fresh agent task before the newly registered `capture` tool appears. +MCP clients may need to restart the server and refresh tool discovery before new tools or parameters appear. ## Setup Prompt @@ -125,19 +134,19 @@ interaction scripting, or review tooling; use Iris only as the camera. ## What it does for you - Renders with your real installed Chrome, driven over the DevTools Protocol -- Waits for fonts, image loads, entrance animations, and — on `--full` — scroll-triggers lazy-loaded content before capturing +- Waits for fonts, image loads, and finite entrance animations under one shared three-second budget. Rendered-frame waits and full-page scrolling also count toward the capture timeout. - Captures the first matching element with `--selector`, automatically scrolling it into view and settling newly visible content before framing it - Serves the same capture engine to coding agents with `iris mcp`, returning pixels inline instead of making the agent locate a file -- Retina `@2x` output by default; full pages taller than Chrome's ~16k px render limit fall back to `@1x` automatically (the report tells you which you got) -- Image format from the `-o` extension or `--format`: `png` (default), `jpg`, `webp` (JPEG/WebP encode at quality 90) -- One browser process, concurrent tabs; a failed URL prints `✗` and never kills the batch (exit code 1 if anything failed) +- Retina `@2x` output by default; full pages taller than Chrome's ~16k px render limit fall back to `@1x` automatically (the report tells you which you got); `--max-pixels` lowers output density without resizing the page +- Image format from the `-o` extension or `--format`: `png` (default), `jpg`, `webp`; `--quality` controls JPEG/WebP encoding (0–100, default 90) +- One browser process, concurrent headless windows; a failed URL prints `✗` and never kills the batch (exit code 1 if anything failed) - Batch filenames derive from the URL (`example.com-pricing.png`); collisions get `-2`, `-3` suffixes - `--json` writes one JSON object per completed capture to stdout, in concurrent completion order; capture failures are JSON too and still produce exit code 1 Element capture is intentionally CSS-selector based: Iris captures the first match in document order. `--selector` conflicts with `--full`; `--padding` requires it. Cross-origin iframe contents and capturing every match are not supported. ```json -{"status":"ok","url":"https://example.com/","output":"/absolute/example.com.png","mode":"element","selector":"h1","padding":24,"css_width":180,"css_height":72,"scale":2.0,"format":"png","bytes":14231} +{"status":"ok","url":"https://example.com/","output":"/absolute/example.com.png","mode":"element","selector":"h1","padding":24,"css_width":180,"css_height":72,"scale":2.0,"format":"png","bytes":14231,"page":{"final_url":"https://example.com/","title":"Example","document_width":1440,"document_height":900,"pending_images":0,"failed_images":0,"fonts_pending":false,"running_animations":0},"timings":{"setup_ms":20,"navigation_ms":35,"ready_ms":100,"screenshot_ms":30,"total_ms":185},"warnings":[]} ``` ## Benchmarking @@ -166,6 +175,20 @@ Reference results (not a performance guarantee): All reference runs used the local `precise-capture.html` fixture, `.capture-target`, 24px padding, PNG, and scale 1. Record `iris --version`, the Chrome version, and the exact command with any published result. +A separate before/after comparison used a GitHub-hosted ARM64 runner with macOS 26.6.2 and Chrome 152.0.7977.83. The baseline was upstream `main` at `227d3109`; the candidate adds the capture controls and agent workflows below. + +| MCP capture time | Before | After | +| --- | ---: | ---: | +| Warm median, 20 captures per version | 1,057 ms | 611 ms | +| Warm p95, nearest rank | 1,261 ms | 708 ms | +| First capture, two fresh servers per version | 5,694 / 2,097 ms | 2,489 / 2,199 ms | + +The run order was baseline, candidate, candidate, baseline. Each server captured the same local fixture once, then ten more times: `.capture-target`, 10px padding, a 320×240 viewport, PNG, scale 1, and a 30-second timeout. Every image decoded to 141×81 pixels. + +The warm median fell by 42%. First-capture times varied and do not establish a reliable startup improvement. + +[Raw timing samples](.github/capture-benchmark.json) retain the run order and capture settings. The linked workflow contains the command receipts. + ## Flags ``` @@ -176,16 +199,37 @@ All reference runs used the local `precise-capture.html` fixture, `.capture-targ --padding nonnegative CSS-pixel padding around a selected element --dark emulate prefers-color-scheme: dark --format png | jpg | webp (a recognized --out extension wins) - --wait extra settle delay after smart waiting + --quality JPEG/WebP quality, 0–100 (default: 90) + --max-pixels cap output pixels by lowering density, not resizing the page + --freeze-animations finish finite animations and pause repeating ones + --wait extra delay before the final readiness check --wait-for wait until a selector exists before capturing --scale device scale factor (overrides the preset's) --jobs concurrent captures (default: min(4, URLs)) - --timeout per-page budget (default: 30) + --timeout capture budget (default: 30; MCP includes queue and browser startup) --chrome browser binary (auto-detected; also via $CHROME) --json emit one JSON object per completed capture ``` -`iris mcp [--chrome ]` serves the single-image `capture` tool over stdio. Batch capture, browser interaction scripting, diffs, and review workflows remain outside Iris. +Capture reports include the final URL, title, document dimensions, pending or failed images, font readiness, active animations, stage timings, and warnings. + +A completed image is not proof that every page asset loaded. An explicit `wait_for` or `wait_ms` can help with app-specific readiness. + +Timings cover the active capture. MCP queueing and cold browser startup are outside `timings.total_ms`, but inside `timeout_seconds`. + +Cancelling an MCP request stops queued work and in-progress captures. Comparisons check for cancellation between decoding steps and rows. A file write that has already started may finish. + +`freeze_animations` changes the state being photographed: finite animations finish and repeating animations pause. Use it for stable layout inspection, not for checking motion. Canvas, video, and JavaScript timers are not frozen. + +A pixel budget too small to retain at least one pixel on each axis returns an error. + +## Compare saved captures + +`iris compare before.png after.png --threshold 8 -o diff.png` prints a JSON report. A pixel changes when its largest channel difference exceeds the threshold. Transparent pixels are composited on white. + +Different dimensions are compared on the union canvas without rescaling. Pixels present in only one image count as changed. The optional PNG marks changed pixels red. + +Comparison errors exit nonzero. A nonzero change count is a result, not a command failure. ## License diff --git a/src/capture.rs b/src/capture.rs index c369353..0269ec8 100644 --- a/src/capture.rs +++ b/src/capture.rs @@ -1,22 +1,26 @@ use std::net::IpAddr; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::Duration; +use std::time::{Duration, Instant}; use anyhow::{Context, Result, anyhow, bail}; +use base64::Engine as _; +use base64::engine::general_purpose::STANDARD as BASE64; use chromiumoxide::browser::{Browser, BrowserConfig}; use chromiumoxide::cdp::browser_protocol::emulation::{ MediaFeature, SetDeviceMetricsOverrideParams, SetEmulatedMediaParams, SetUserAgentOverrideParams, }; use chromiumoxide::cdp::browser_protocol::page::{ - CaptureScreenshotFormat, Viewport as ScreenshotViewport, + CaptureScreenshotFormat, CaptureScreenshotParams, Viewport as ScreenshotViewport, }; +use chromiumoxide::cdp::browser_protocol::target::CreateTargetParams; use chromiumoxide::cdp::js_protocol::runtime::EvaluateParams; use chromiumoxide::error::CdpError; -use chromiumoxide::page::{Page, ScreenshotParams}; +use chromiumoxide::page::Page; use futures::StreamExt; -use serde::{Deserialize, Serialize}; +use num_traits::ToPrimitive; +use serde::{Deserialize, Serialize, de::DeserializeOwned}; use tokio::task::JoinHandle; use url::Url; @@ -147,17 +151,42 @@ pub struct Opts { pub wait_for: Option, pub timeout: Duration, pub format: Format, + pub quality: u8, + pub max_pixels: Option, + pub freeze_animations: bool, +} + +#[derive(Debug, Clone, Default, Serialize)] +pub struct Timings { + pub setup_ms: u64, + pub navigation_ms: u64, + pub ready_ms: u64, + pub screenshot_ms: u64, + pub total_ms: u64, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct PageInfo { + pub final_url: String, + pub title: String, + pub document_width: u32, + pub document_height: u32, + pub pending_images: u32, + pub failed_images: u32, + pub fonts_pending: bool, + pub running_animations: u32, } #[derive(Debug)] pub struct Shot { - /// Captured width and height in CSS pixels. + /// Capture rectangle in CSS pixels; scale records the actual image density. pub width: u32, pub height: u32, - /// Device scale factor actually used (full-page shots too tall for Chrome's - /// ~16k texture limit fall back to 1x). pub scale: f64, pub bytes: u64, + pub page: PageInfo, + pub timings: Timings, + pub warnings: Vec, } #[derive(Debug, Clone, Serialize)] @@ -176,6 +205,9 @@ pub struct SuccessReport { scale: f64, format: &'static str, bytes: u64, + page: PageInfo, + timings: Timings, + warnings: Vec, } #[derive(Debug, Clone, Serialize)] @@ -211,6 +243,9 @@ pub fn success_report( scale: shot.scale, format: format.ext(), bytes: shot.bytes, + page: shot.page.clone(), + timings: shot.timings.clone(), + warnings: shot.warnings.clone(), } } @@ -241,17 +276,24 @@ pub fn absolute_output(path: &Path) -> String { #[derive(Debug)] pub struct CapturedImage { pub shot: Shot, - pub data: Vec, + /// Keep Chrome's encoded bytes for MCP instead of decoding and encoding again. + pub data: String, } impl CapturedImage { + pub fn pixels(&self) -> Result> { + BASE64 + .decode(&self.data) + .context("Chrome returned invalid image data") + } + pub async fn write_to(&self, out: &Path) -> Result<()> { if let Some(parent) = out.parent().filter(|path| !path.as_os_str().is_empty()) { tokio::fs::create_dir_all(parent) .await .with_context(|| format!("failed to create {}", parent.display()))?; } - tokio::fs::write(out, &self.data) + tokio::fs::write(out, self.pixels()?) .await .with_context(|| format!("failed to write {}", out.display())) } @@ -285,6 +327,39 @@ pub struct Session { user_agent: Option, } +struct CaptureTab { + page: Page, + closing: bool, +} + +impl CaptureTab { + async fn close(mut self) { + let result = tokio::time::timeout(Duration::from_secs(2), self.page.clone().close()).await; + self.closing = matches!(result, Ok(Ok(_))); + } +} + +impl Drop for CaptureTab { + fn drop(&mut self) { + if !self.closing { + let page = self.page.clone(); + tokio::spawn(async move { + let _ = tokio::time::timeout(Duration::from_secs(2), page.close()).await; + }); + } + } +} + +impl Drop for Session { + fn drop(&mut self) { + self.handler.abort(); + } +} + +fn elapsed_ms(started: Instant) -> u64 { + u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX) +} + impl Session { pub async fn launch(chrome: Option, viewport: Viewport) -> Result { let profile_dir = ProfileDir::unique(); @@ -294,10 +369,12 @@ impl Session { if let Some(path) = chrome.or_else(find_chrome) { config = config.chrome_executable(path); } - let config = config.build().map_err(|e| anyhow!(e))?; + let config = config + .build() + .map_err(|e| anyhow!("{e}: set CHROME=/path/to/browser or pass --chrome"))?; let (browser, mut handler) = Browser::launch(config) .await - .context("failed to launch Chrome (install Google Chrome or pass --chrome)")?; + .context("failed to launch Chrome (install a Chrome-family browser, set CHROME, or pass --chrome)")?; // Newer Chrome versions emit CDP messages chromiumoxide can't deserialize // (Serde errors) — harmless, keep pumping. Transport errors mean the // connection is gone, so stop. @@ -323,14 +400,27 @@ impl Session { } pub async fn capture(&self, url: &str, opts: &Opts) -> Result { - let page = tokio::time::timeout(opts.timeout, self.browser.new_page("about:blank")) + let started = Instant::now(); + let deadline = tokio::time::Instant::now() + opts.timeout; + // Background tabs can stop producing the rendered frames these captures await. + let target = CreateTargetParams::builder() + .url("about:blank") + .new_window(true) + .build() + .map_err(|error| anyhow!(error))?; + let page = tokio::time::timeout_at(deadline, self.browser.new_page(target)) .await .map_err(|_| anyhow!("timed out opening a tab"))??; - let result = tokio::time::timeout(opts.timeout, self.pipeline(&page, url, opts)) - .await - .map_err(|_| anyhow!("timed out after {}s", opts.timeout.as_secs())) - .and_then(|r| r); - let _ = page.close().await; + let tab = CaptureTab { + page, + closing: false, + }; + let result = + tokio::time::timeout_at(deadline, self.pipeline(&tab.page, url, opts, started)) + .await + .map_err(|_| anyhow!("capture timed out after {}s", opts.timeout.as_secs())) + .and_then(|result| result); + tab.close().await; result } @@ -338,35 +428,31 @@ impl Session { !self.handler.is_finished() } - async fn pipeline(&self, page: &Page, url: &str, opts: &Opts) -> Result { - let started = std::time::Instant::now(); + async fn pipeline( + &self, + page: &Page, + url: &str, + opts: &Opts, + started: Instant, + ) -> Result { let v = opts.viewport; - page.execute( - SetDeviceMetricsOverrideParams::builder() - .width(v.width as i64) - .height(v.height as i64) - .device_scale_factor(v.scale) - .mobile(v.mobile) - .build() - .map_err(|e| anyhow!(e))?, - ) + page.execute(SetDeviceMetricsOverrideParams::new( + i64::from(v.width), + i64::from(v.height), + v.scale, + v.mobile, + )) .await?; - let ua = if v.mobile { - Some(IPHONE_UA.to_string()) + let user_agent = if v.mobile { + Some(IPHONE_UA) } else { - self.user_agent.clone() + self.user_agent.as_deref() }; - if let Some(ua) = ua { - page.execute( - SetUserAgentOverrideParams::builder() - .user_agent(ua) - .build() - .map_err(|e| anyhow!(e))?, - ) - .await?; + if let Some(user_agent) = user_agent { + page.execute(SetUserAgentOverrideParams::new(user_agent)) + .await?; } - if opts.dark { page.execute( SetEmulatedMediaParams::builder() @@ -378,169 +464,178 @@ impl Session { ) .await?; } - - page.goto(url).await?; - page.wait_for_navigation().await?; - - self.eval(page, SETTLE_JS.into()).await?; + let mut timings = Timings { + setup_ms: elapsed_ms(started), + ..Timings::default() + }; + let navigation = Instant::now(); + page.goto(url).await.context("navigation failed")?; + timings.navigation_ms = elapsed_ms(navigation); + + let ready = Instant::now(); + let remaining_ms = || { + u64::try_from( + opts.timeout + .saturating_sub(started.elapsed()) + .saturating_sub(Duration::from_millis(500)) + .as_millis(), + ) + .unwrap_or(u64::MAX) + }; if let Some(selector) = &opts.wait_for { - // Undercut the outer timeout so the descriptive selector error surfaces - // instead of a generic "timed out". - let budget = opts - .timeout - .saturating_sub(started.elapsed()) - .saturating_sub(Duration::from_millis(500)); - self.eval(page, wait_for_js(selector, budget.as_millis() as u64)) + self.eval::<()>(page, wait_for_js(selector, remaining_ms())?) .await?; } match &opts.mode { CaptureMode::Viewport => {} - CaptureMode::FullPage => { - self.eval(page, SCROLL_JS.into()).await?; - } + CaptureMode::FullPage => self.eval::<()>(page, SCROLL_JS.into()).await?, CaptureMode::Element { selector, .. } => { - let budget = opts - .timeout - .saturating_sub(started.elapsed()) - .saturating_sub(Duration::from_millis(500)); - self.eval( - page, - wait_and_scroll_js(selector, budget.as_millis() as u64), - ) - .await?; - // Scrolling can start image loads, IntersectionObservers, and - // entrance transitions that the initial settle could not see. - self.eval(page, SETTLE_JS.into()).await?; + self.eval::<()>(page, wait_and_scroll_js(selector, remaining_ms())?) + .await?; } } if opts.wait_ms > 0 { tokio::time::sleep(Duration::from_millis(opts.wait_ms)).await; - self.eval(page, SETTLE_JS.into()).await?; } - - let screenshot_params = || { - let mut params = ScreenshotParams::builder().format(opts.format.cdp()); - if opts.format != Format::Png { - params = params.quality(90); - } - params - }; - - let (width, height, scale, data) = match &opts.mode { - CaptureMode::Viewport => ( - v.width, - v.height, - v.scale, - page.screenshot(screenshot_params().build()).await?, + self.eval::<()>( + page, + format!( + "({SETTLE_JS})({}, {})", + matches!(opts.mode, CaptureMode::FullPage), + opts.freeze_animations ), - CaptureMode::FullPage => { - let doc_h = self - .eval_u32(page, DOC_HEIGHT_JS) - .await - .unwrap_or(v.height) - .max(v.height); - if doc_h as f64 * v.scale <= 16_000.0 { - // Retina full page: grow the viewport to the whole document so the - // scale factor still applies (CDP's captureBeyondViewport renders at 1x). - page.execute( - SetDeviceMetricsOverrideParams::builder() - .width(v.width as i64) - .height(doc_h as i64) - .device_scale_factor(v.scale) - .mobile(v.mobile) - .build() - .map_err(|e| anyhow!(e))?, - ) + ) + .await?; + let info: PageInfo = self.eval(page, PAGE_INFO_JS.into()).await?; + timings.ready_ms = elapsed_ms(ready); + + let clip = match &opts.mode { + CaptureMode::Viewport => { + let (x, y): (f64, f64) = self + .eval(page, "[window.scrollX, window.scrollY]".into()) .await?; - self.eval(page, SETTLE_JS.into()).await?; - ( - v.width, - doc_h, - v.scale, - page.screenshot(screenshot_params().build()).await?, - ) - } else { - ( - v.width, - doc_h, - 1.0, - page.screenshot(screenshot_params().full_page(true).build()) - .await?, - ) + ClipRect { + x, + y, + width: f64::from(v.width), + height: f64::from(v.height), } } + CaptureMode::FullPage => ClipRect { + x: 0.0, + y: 0.0, + width: f64::from(info.document_width.max(v.width)), + height: f64::from(info.document_height.max(v.height)), + }, CaptureMode::Element { selector, padding } => { - let value = self.eval(page, element_bounds_js(selector)).await?; - let bounds: ElementBounds = serde_json::from_value(value) - .context("failed to read selected element bounds")?; - let clip = round_clip(&bounds, *padding) - .with_context(|| format!("cannot capture selected element: {selector}"))?; - let cdp_clip = ScreenshotViewport::builder() - .x(clip.x) - .y(clip.y) - .width(clip.width) - .height(clip.height) - .scale(1.0) - .build() - .map_err(|e| anyhow!(e))?; - ( - clip.width as u32, - clip.height as u32, - v.scale, - page.screenshot( - screenshot_params() - .clip(cdp_clip) - .capture_beyond_viewport(true) - .build(), - ) - .await?, - ) + let bounds: ElementBounds = self.eval(page, element_bounds_js(selector)?).await?; + round_clip(&bounds, *padding) + .with_context(|| format!("cannot capture selected element: {selector}"))? } }; - - let bytes = data.len() as u64; + let mut scale = v.scale; + if let Some(max_pixels) = opts.max_pixels { + scale = scale.min( + (max_pixels + .to_f64() + .context("pixel budget is out of range")? + / (clip.width * clip.height)) + .sqrt(), + ); + scale = ((clip.width * scale).floor() / clip.width) + .min((clip.height * scale).floor() / clip.height); + if scale <= 0.0 { + bail!("pixel budget is too small for this capture's aspect ratio"); + } + } else if matches!(opts.mode, CaptureMode::FullPage) && clip.height * scale > 16_000.0 { + // Past Chrome's ~16k px render limit, full pages fall back to 1x. + scale = 1.0; + } + let capture_started = Instant::now(); + let mut params = CaptureScreenshotParams::builder() + .format(opts.format.cdp()) + .from_surface(true) + .capture_beyond_viewport(true) + .clip(ScreenshotViewport { + x: clip.x, + y: clip.y, + width: clip.width, + height: clip.height, + scale: scale / v.scale, + }); + if opts.format != Format::Png { + params = params.quality(i64::from(opts.quality)); + } + let data: String = page + .execute(params.build()) + .await + .context("screenshot failed")? + .result + .data + .into(); + timings.screenshot_ms = elapsed_ms(capture_started); + timings.total_ms = elapsed_ms(started); + let padding = data + .as_bytes() + .iter() + .rev() + .take_while(|&&byte| byte == b'=') + .count(); + let bytes = u64::try_from(data.len() / 4 * 3 - padding)?; + let mut warnings = Vec::new(); + if info.pending_images > 0 { + warnings.push(format!("{} images were still loading", info.pending_images)); + } + if info.failed_images > 0 { + warnings.push(format!("{} images failed to load", info.failed_images)); + } + if info.fonts_pending { + warnings.push("fonts were still loading".into()); + } + if info.running_animations > 0 { + warnings.push(format!( + "{} animations were still running", + info.running_animations + )); + } Ok(CapturedImage { shot: Shot { - width, - height, + width: clip + .width + .to_u32() + .context("capture width is out of range")?, + height: clip + .height + .to_u32() + .context("capture height is out of range")?, scale, bytes, + page: info, + timings, + warnings, }, data, }) } /// Run a JS expression (promises awaited); surface page-side exceptions as errors. - async fn eval(&self, page: &Page, js: String) -> Result { + async fn eval(&self, page: &Page, js: String) -> Result { let params = EvaluateParams::builder() .expression(js) .await_promise(true) .return_by_value(true) .build() - .map_err(|e| anyhow!(e))?; - let resp = page.execute(params).await?; - if let Some(details) = &resp.result.exception_details { - let msg = details + .map_err(|error| anyhow!(error))?; + let response = page.execute(params).await?.result; + if let Some(details) = response.exception_details { + let message = details .exception - .as_ref() - .and_then(|e| e.description.clone()) - .unwrap_or_else(|| details.text.clone()); - bail!("{}", msg.lines().next().unwrap_or("page script failed")); + .and_then(|exception| exception.description) + .unwrap_or(details.text); + bail!("{}", message.lines().next().unwrap_or("page script failed")); } - Ok(resp - .result - .result - .value - .clone() - .unwrap_or(serde_json::Value::Null)) - } - - async fn eval_u32(&self, page: &Page, js: &str) -> Result { - let value = self.eval(page, js.into()).await?; - value - .as_f64() - .map(|n| n as u32) - .ok_or_else(|| anyhow!("expected a number from page")) + serde_json::from_value(response.result.value.unwrap_or(serde_json::Value::Null)) + .context("unexpected response from capture script") } pub async fn close(mut self) { @@ -552,10 +647,53 @@ impl Session { let _ = self.browser.kill().await; } self.handler.abort(); + reap_profile_processes(self.profile_dir.path()).await; let _ = tokio::fs::remove_dir_all(self.profile_dir.path()).await; } } +/// Chromium helpers can outlive the browser: Helium leaves its network +/// service running, reparented to launchd. Stop every process still started +/// with this session's profile, politely first. +async fn reap_profile_processes(profile: &Path) { + let marker = format!("--user-data-dir={}", profile.display()); + for signal in ["-TERM", "-KILL"] { + let table = process_table(); + let pids = profile_pids(&table, &marker); + if pids.is_empty() { + return; + } + let _ = std::process::Command::new("kill") + .arg(signal) + .args(&pids) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status(); + let deadline = Instant::now() + Duration::from_secs(2); + while Instant::now() < deadline && !profile_pids(&process_table(), &marker).is_empty() { + tokio::time::sleep(Duration::from_millis(50)).await; + } + } +} + +fn process_table() -> String { + std::process::Command::new("ps") + .args(["axww", "-o", "pid=,command="]) + .output() + .map(|output| String::from_utf8_lossy(&output.stdout).into_owned()) + .unwrap_or_default() +} + +/// PIDs whose command line names exactly `marker` as an argument, so the +/// profile `...-1` never matches a sibling session's `...-10`. +fn profile_pids<'a>(table: &'a str, marker: &str) -> Vec<&'a str> { + table + .lines() + .filter(|line| line.split_whitespace().any(|arg| arg == marker)) + .filter_map(|line| line.split_whitespace().next()) + .collect() +} + pub fn parse_viewport(size: &str, scale: Option) -> Result { let (width, height, preset_scale, mobile) = match size { "desktop" => (1440, 900, 2.0, false), @@ -622,6 +760,7 @@ fn find_chrome() -> Option { "/Applications/Google Chrome Canary.app/Contents/MacOS/Google Chrome Canary", "/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge", "/Applications/Brave Browser.app/Contents/MacOS/Brave Browser", + "/Applications/Helium.app/Contents/MacOS/Helium", ]; #[cfg(windows)] const CANDIDATES: &[&str] = &[ @@ -638,62 +777,72 @@ fn find_chrome() -> Option { CANDIDATES.iter().map(PathBuf::from).find(|p| p.exists()) } -/// Fonts loaded, near-viewport images loaded (3s cap each), two frames painted, -/// then any running finite animations/transitions — entrance fade-ins — allowed -/// to finish (3s cap; infinite loops are skipped, they never settle). Off-screen -/// images are ignored: they don't appear in the capture, and lazy-loaded ones -/// would stall the wait forever. Full-page captures grow the viewport to the -/// whole document before the final settle, so everything counts as near there. -const SETTLE_JS: &str = r#"(async () => { - if (document.fonts) { try { await document.fonts.ready; } catch {} } - const near = (img) => { - const r = img.getBoundingClientRect(); - return r.top < innerHeight * 1.5 && r.bottom > -innerHeight * 0.5; - }; - await Promise.all(Array.from(document.images) - .filter(img => !img.complete && near(img)) - .map(img => new Promise(r => { - img.addEventListener('load', r, { once: true }); - img.addEventListener('error', r, { once: true }); - setTimeout(r, 3000); - }))); - await new Promise(r => requestAnimationFrame(() => requestAnimationFrame(r))); - const finite = document.getAnimations().filter(a => { +/// Wait once, after scrolling. Asset and finite-animation waits share the same cap. +const SETTLE_JS: &str = r#"async (fullPage, freeze) => { + const deadline = performance.now() + 3000; + const bounded = async (promise) => { + let timer; try { - return a.playState === 'running' && a.effect.getTiming().iterations !== Infinity; - } catch { return false; } - }); - await Promise.race([ - Promise.all(finite.map(a => a.finished.catch(() => {}))), - new Promise(r => setTimeout(r, 3000)), - ]); - await new Promise(r => requestAnimationFrame(() => requestAnimationFrame(r))); -})()"#; + await Promise.race([promise, new Promise(resolve => { + timer = setTimeout(resolve, Math.max(0, deadline - performance.now())); + })]); + } finally { clearTimeout(timer); } + }; + const frames = () => new Promise(resolve => requestAnimationFrame(() => requestAnimationFrame(resolve))); + const near = img => { + const rect = img.getBoundingClientRect(); + return fullPage || (rect.top < innerHeight * 1.5 && rect.bottom > -innerHeight * 0.5); + }; + const images = Array.from(document.images).filter(img => !img.complete && near(img)); + await bounded(Promise.all([ + document.fonts?.ready, + ...images.map(img => new Promise(resolve => { + img.addEventListener('load', resolve, { once: true }); + img.addEventListener('error', resolve, { once: true }); + })), + ])); + await frames(); + const animations = document.getAnimations().filter(animation => animation.playState === 'running'); + if (freeze) { + for (const animation of animations) { + try { + if (animation.effect?.getComputedTiming().endTime !== Infinity) animation.finish(); + else animation.pause(); + } catch { animation.pause(); } + } + } else { + const finite = animations.filter(animation => animation.effect?.getComputedTiming().endTime !== Infinity); + await bounded(Promise.all(finite.map(animation => animation.finished.catch(() => {})))); + } + await frames(); +}"#; -/// Step-scroll to the bottom so IntersectionObserver lazy-loading fires, then back to top. +/// Preserve the requested viewport while triggering intersection-based lazy loading. const SCROLL_JS: &str = r#"(async () => { - const height = () => Math.max( - document.body?.scrollHeight ?? 0, - document.documentElement.scrollHeight - ); - const step = Math.max(200, window.innerHeight); - for (let y = 0, guard = 0; y < height() && guard < 500; y += step, guard++) { - window.scrollTo(0, y); - await new Promise(r => setTimeout(r, 60)); + const height = () => Math.max(document.body?.scrollHeight ?? 0, document.documentElement.scrollHeight); + const step = Math.max(200, innerHeight * 0.9); + for (let y = 0, steps = 0; y < height() && steps < 500; y += step, steps++) { + scrollTo(0, y); + await new Promise(resolve => requestAnimationFrame(resolve)); } - window.scrollTo(0, 0); - await new Promise(r => setTimeout(r, 150)); - await new Promise(r => requestAnimationFrame(() => requestAnimationFrame(r))); + scrollTo(0, 0); + await new Promise(resolve => requestAnimationFrame(resolve)); })()"#; -const DOC_HEIGHT_JS: &str = r#"Math.max( - document.body?.scrollHeight ?? 0, - document.documentElement.scrollHeight -)"#; - -fn wait_for_js(selector: &str, budget_ms: u64) -> String { - let selector = serde_json::to_string(selector).expect("selector is serializable"); - format!( +const PAGE_INFO_JS: &str = r#"(() => ({ + final_url: location.href, + title: document.title, + document_width: Math.max(document.body?.scrollWidth ?? 0, document.documentElement.scrollWidth, innerWidth), + document_height: Math.max(document.body?.scrollHeight ?? 0, document.documentElement.scrollHeight, innerHeight), + pending_images: Array.from(document.images).filter(image => !image.complete).length, + failed_images: Array.from(document.images).filter(image => image.complete && image.currentSrc && !image.naturalWidth).length, + fonts_pending: document.fonts?.status === 'loading', + running_animations: document.getAnimations().filter(animation => animation.playState === 'running').length, +}))()"#; + +fn wait_for_js(selector: &str, budget_ms: u64) -> Result { + let selector = serde_json::to_string(selector)?; + Ok(format!( r#"(async () => {{ const deadline = Date.now() + {budget_ms}; const selector = {selector}; @@ -706,12 +855,12 @@ fn wait_for_js(selector: &str, budget_ms: u64) -> String { await new Promise(r => setTimeout(r, 100)); }} }})()"# - ) + )) } -fn wait_and_scroll_js(selector: &str, budget_ms: u64) -> String { - let selector = serde_json::to_string(selector).expect("selector is serializable"); - format!( +fn wait_and_scroll_js(selector: &str, budget_ms: u64) -> Result { + let selector = serde_json::to_string(selector)?; + Ok(format!( r#"(async () => {{ const deadline = Date.now() + {budget_ms}; const selector = {selector}; @@ -727,12 +876,12 @@ fn wait_and_scroll_js(selector: &str, budget_ms: u64) -> String { element.scrollIntoView({{ block: "center", inline: "center" }}); await new Promise(r => requestAnimationFrame(() => requestAnimationFrame(r))); }})()"# - ) + )) } -fn element_bounds_js(selector: &str) -> String { - let selector = serde_json::to_string(selector).expect("selector is serializable"); - format!( +fn element_bounds_js(selector: &str) -> Result { + let selector = serde_json::to_string(selector)?; + Ok(format!( r#"(() => {{ const selector = {selector}; let element; @@ -761,11 +910,11 @@ fn element_bounds_js(selector: &str) -> String { ) }}; }})()"# - ) + )) } fn round_clip(bounds: &ElementBounds, padding: u32) -> Result { - let padding = padding as f64; + let padding = f64::from(padding); let left = (bounds.x - padding).floor().clamp(0.0, bounds.doc_width); let top = (bounds.y - padding).floor().clamp(0.0, bounds.doc_height); let right = (bounds.x + bounds.width + padding) @@ -794,6 +943,20 @@ fn round_clip(bounds: &ElementBounds, padding: u32) -> Result { mod tests { use super::*; + #[test] + fn profile_pids_match_the_whole_profile_argument() { + let table = concat!( + " 101 /Apps/Helium --user-data-dir=/tmp/iris-chrome-7-1 --headless\n", + " 102 /Apps/Helium Helper --type=utility --user-data-dir=/tmp/iris-chrome-7-10\n", + " 103 /Apps/Helium Helper --type=renderer --user-data-dir=/tmp/iris-chrome-7-1\n", + ); + + assert_eq!( + profile_pids(table, "--user-data-dir=/tmp/iris-chrome-7-1"), + ["101", "103"] + ); + } + #[test] fn clip_rounds_outward_and_applies_padding() { let clip = round_clip( @@ -898,8 +1061,125 @@ mod tests { assert!(parse_viewport("desktop", Some(f64::NAN)).is_err()); } + #[tokio::test] + async fn full_page_capture_keeps_viewport_relative_heights() -> Result<()> { + let _browser_guard = crate::BROWSER_TEST_LOCK.lock().await; + let fixture = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/full-page-layout.html") + .canonicalize()?; + let url = Url::from_file_path(fixture).map_err(|_| anyhow!("invalid fixture URL"))?; + let viewport = Viewport { + width: 320, + height: 240, + scale: 2.0, + mobile: false, + }; + let session = Session::launch(None, viewport).await?; + let capture = session + .capture( + url.as_str(), + &page_opts(viewport, CaptureMode::FullPage, false), + ) + .await; + session.close().await; + let capture = capture?; + + // A 100vh header, 900px main, and 100px footer make a 1240px page. + assert_eq!( + (capture.shot.width, capture.shot.height, capture.shot.scale), + (320, 1240, 2.0) + ); + let image = image::load_from_memory(&capture.pixels()?)?.to_rgb8(); + assert_eq!(image.dimensions(), (640, 2480)); + // The header stays one viewport (240px, 480 device pixels) tall. + assert_eq!(image.get_pixel(320, 479).0, [22, 78, 154]); + assert_eq!(image.get_pixel(320, 480).0, [233, 138, 34]); + assert_eq!(image.get_pixel(320, 2479).0, [39, 133, 74]); + Ok(()) + } + + #[tokio::test] + async fn full_page_capture_past_the_render_limit_falls_back_to_1x() -> Result<()> { + let _browser_guard = crate::BROWSER_TEST_LOCK.lock().await; + let url = + "data:text/html,
"; + let viewport = Viewport { + width: 320, + height: 240, + scale: 2.0, + mobile: false, + }; + let session = Session::launch(None, viewport).await?; + let capture = session + .capture(url, &page_opts(viewport, CaptureMode::FullPage, false)) + .await; + session.close().await; + let capture = capture?; + + // 9000 CSS px at 2x would be 18,000 device px. + assert_eq!(capture.shot.scale, 1.0); + let image = image::load_from_memory(&capture.pixels()?)?.to_rgb8(); + assert_eq!(image.dimensions(), (320, 9000)); + Ok(()) + } + + #[tokio::test] + async fn full_page_pixel_budget_preserves_viewport_relative_layout() -> Result<()> { + let _browser_guard = crate::BROWSER_TEST_LOCK.lock().await; + let fixture = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/full-page-layout.html") + .canonicalize()?; + let url = Url::from_file_path(fixture).map_err(|_| anyhow!("invalid fixture URL"))?; + let viewport = Viewport { + width: 320, + height: 240, + scale: 2.0, + mobile: false, + }; + let session = Session::launch(None, viewport).await?; + let mut opts = page_opts(viewport, CaptureMode::FullPage, false); + opts.max_pixels = Some(100_000); + let capture = session.capture(url.as_str(), &opts).await; + session.close().await; + let image = image::load_from_memory(&capture?.pixels()?)?.to_rgb8(); + assert_eq!(image.dimensions(), (160, 620)); + assert_eq!(image.get_pixel(80, 119).0, [22, 78, 154]); + assert_eq!(image.get_pixel(80, 120).0, [233, 138, 34]); + assert_eq!(image.get_pixel(80, 619).0, [39, 133, 74]); + Ok(()) + } + + #[tokio::test] + async fn viewport_capture_preserves_fragment_scroll_position() -> Result<()> { + let _browser_guard = crate::BROWSER_TEST_LOCK.lock().await; + let fixture = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/full-page-layout.html") + .canonicalize()?; + let mut url = Url::from_file_path(fixture).map_err(|_| anyhow!("invalid fixture URL"))?; + url.set_fragment(Some("totals")); + let viewport = Viewport { + width: 320, + height: 240, + scale: 1.0, + mobile: false, + }; + let session = Session::launch(None, viewport).await?; + let capture = session + .capture( + url.as_str(), + &page_opts(viewport, CaptureMode::Viewport, false), + ) + .await; + session.close().await; + let image = image::load_from_memory(&capture?.pixels()?)?.to_rgb8(); + assert_eq!(image.get_pixel(160, 0).0, [233, 138, 34]); + assert_eq!(image.get_pixel(160, 239).0, [39, 133, 74]); + Ok(()) + } + #[tokio::test] async fn browser_element_capture_contract() -> Result<()> { + let _browser_guard = crate::BROWSER_TEST_LOCK.lock().await; let temp = std::env::temp_dir().join(format!("iris-browser-{}", std::process::id())); tokio::fs::create_dir_all(&temp).await?; let fixture = Path::new(env!("CARGO_MANIFEST_DIR")) @@ -983,7 +1263,7 @@ mod tests { (dark_viewport.shot.width, dark_viewport.shot.height), (320, 240) ); - assert_eq!(png_dimensions(&dark_viewport.data), (320, 240)); + assert_eq!(png_dimensions(&dark_viewport.pixels()?), (320, 240)); let full_page = session .capture( @@ -994,7 +1274,7 @@ mod tests { assert_eq!(full_page.shot.width, 320); assert!(full_page.shot.height > viewport.height); assert_eq!( - png_dimensions(&full_page.data), + png_dimensions(&full_page.pixels()?), (full_page.shot.width, full_page.shot.height) ); @@ -1011,7 +1291,7 @@ mod tests { ) .await?; assert_eq!((mobile.shot.width, mobile.shot.height), (390, 844)); - assert_eq!(png_dimensions(&mobile.data), (390, 844)); + assert_eq!(png_dimensions(&mobile.pixels()?), (390, 844)); let invalid = session .capture(url.as_str(), &element_opts(viewport, "[", 0, Format::Png)) @@ -1021,11 +1301,10 @@ mod tests { let mut missing_url = url.clone(); missing_url.set_query(Some("missing=1")); + let mut missing_opts = element_opts(viewport, ".capture-target", 0, Format::Png); + missing_opts.timeout = Duration::from_secs(3); let missing = session - .capture( - missing_url.as_str(), - &element_opts(viewport, ".capture-target", 0, Format::Png), - ) + .capture(missing_url.as_str(), &missing_opts) .await .unwrap_err(); assert!(format!("{missing:#}").contains("selector never appeared: .capture-target")); @@ -1054,8 +1333,11 @@ mod tests { dark: false, wait_ms: 0, wait_for: None, - timeout: Duration::from_secs(3), + timeout: Duration::from_secs(30), format, + quality: 90, + max_pixels: None, + freeze_animations: false, } } @@ -1066,13 +1348,19 @@ mod tests { dark, wait_ms: 0, wait_for: None, - timeout: Duration::from_secs(5), + timeout: Duration::from_secs(30), format: Format::Png, + quality: 90, + max_pixels: None, + freeze_animations: false, } } async fn capture_to(session: &Session, url: &str, path: &Path, opts: &Opts) -> Result { - let image = session.capture(url, opts).await?; + let image = session + .capture(url, opts) + .await + .with_context(|| format!("capture failed: {}", path.display()))?; image.write_to(path).await?; Ok(image.shot) } diff --git a/src/compare.rs b/src/compare.rs new file mode 100644 index 0000000..5f0319f --- /dev/null +++ b/src/compare.rs @@ -0,0 +1,487 @@ +//! Pixel comparison of two captures. +//! +//! Decoded pixels are composited on white and compared per channel. + +use std::fs::File; +use std::io::{BufWriter, Write}; +use std::path::Path; + +use anyhow::{Context, Result, bail, ensure}; +use image::codecs::png::{CompressionType, FilterType, PngEncoder}; +use image::{ExtendedColorType, ImageEncoder, ImageFormat, ImageReader, Limits, RgbaImage}; +use serde::Serialize; + +use crate::capture::absolute_output; + +/// Reject dimensions and decoded buffers too large to compare safely. +const MAX_DIMENSION: u32 = 32_768; +const MAX_DECODE_BYTES: u64 = 256 * 1024 * 1024; +const MAX_DIFF_BYTES: u64 = 512 * 1024 * 1024; + +/// Difference-image color for a changed pixel. +const CHANGED: [u8; 3] = [255, 0, 0]; +/// Difference-image color for union-canvas area neither capture covers. +const BLANK: [u8; 3] = [255, 255, 255]; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub struct Dimensions { + pub width: u32, + pub height: u32, +} + +/// Smallest rectangle containing every changed pixel. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub struct Bounds { + pub x: u32, + pub y: u32, + pub width: u32, + pub height: u32, +} + +#[derive(Debug, Clone, Serialize)] +pub struct Comparison { + pub before: Dimensions, + pub after: Dimensions, + pub same_dimensions: bool, + pub changed_pixels: u64, + /// Pixels in the rectangular union canvas. + pub total_pixels: u64, + pub changed_fraction: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub bounds: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub output: Option, +} + +/// Compares two captures pixel by pixel. +/// +/// A pixel changes when its largest composited channel gap is greater than +/// `threshold`. +/// +/// Different sizes use the union canvas without rescaling. A pixel covered by +/// only one capture changes. +/// +/// When `output` is present, writes a PNG with changes in red and unchanged +/// pixels faded from `after`. +pub fn compare( + before: &Path, + after: &Path, + threshold: u8, + output: Option<&Path>, + is_cancelled: impl Fn() -> bool, +) -> Result { + ensure!(!is_cancelled(), "comparison cancelled"); + let before_image = decode(before, "before")?; + ensure!(!is_cancelled(), "comparison cancelled"); + let after_image = decode(after, "after")?; + ensure!(!is_cancelled(), "comparison cancelled"); + + let (before_width, before_height) = before_image.dimensions(); + let (after_width, after_height) = after_image.dimensions(); + let before_width_usize = usize::try_from(before_width) + .context("before image width exceeds this platform's address space")?; + let before_height_usize = usize::try_from(before_height) + .context("before image height exceeds this platform's address space")?; + let after_width_usize = usize::try_from(after_width) + .context("after image width exceeds this platform's address space")?; + let after_height_usize = usize::try_from(after_height) + .context("after image height exceeds this platform's address space")?; + let union_width = before_width.max(after_width); + let union_height = before_height.max(after_height); + let union_width_usize = usize::try_from(union_width) + .context("difference image width exceeds this platform's address space")?; + let union_height_usize = usize::try_from(union_height) + .context("difference image height exceeds this platform's address space")?; + let total_pixels = u64::from(union_width) * u64::from(union_height); + + let mut diff = if output.is_some() { + let bytes = total_pixels * 3; + if bytes > MAX_DIFF_BYTES { + bail!( + "a {union_width}x{union_height} difference image needs {bytes} bytes of memory; \ + compare smaller captures or omit the difference output" + ); + } + let length = usize::try_from(bytes) + .context("difference image exceeds this platform's address space")?; + Some(vec![0; length]) + } else { + None + }; + + let before_raw = before_image.as_raw(); + let after_raw = after_image.as_raw(); + let diff_stride = union_width_usize * 3; + + let mut changed_pixels = 0_u64; + let mut change_box = ChangeBox::new(); + + for y in 0..union_height_usize { + ensure!(!is_cancelled(), "comparison cancelled"); + let before_row = row(before_raw, before_width_usize, before_height_usize, y); + let after_row = row(after_raw, after_width_usize, after_height_usize, y); + let overlap = match (before_row, after_row) { + (Some(_), Some(_)) => before_width_usize.min(after_width_usize), + _ => 0, + }; + let mut diff_row = diff + .as_mut() + .map(|buffer| &mut buffer[y * diff_stride..][..diff_stride]); + + let mut row_changed = 0_u64; + let mut first = 0; + let mut last = 0; + + if let (Some(before_row), Some(after_row)) = (before_row, after_row) { + let span = overlap * 4; + for (x, (before_pixel, after_pixel)) in before_row[..span] + .as_chunks::<4>() + .0 + .iter() + .zip(after_row[..span].as_chunks::<4>().0.iter()) + .enumerate() + { + let after_pixel = on_white(after_pixel); + let changed = distance(on_white(before_pixel), after_pixel) > threshold; + if changed { + if row_changed == 0 { + first = x; + } + last = x; + row_changed += 1; + } + if let Some(target) = diff_row.as_deref_mut() { + let color = if changed { CHANGED } else { ghost(after_pixel) }; + let at = x * 3; + target[at..at + 3].copy_from_slice(&color); + } + } + } + + let before_has_row = before_row.is_some(); + let after_has_row = after_row.is_some(); + let row_width = if diff_row.is_some() { + union_width_usize + } else { + before_row + .map_or(0, |row| row.len() / 4) + .max(after_row.map_or(0, |row| row.len() / 4)) + }; + for x in overlap..row_width { + let changed = (before_has_row && x < before_width_usize) + != (after_has_row && x < after_width_usize); + if changed { + if row_changed == 0 { + first = x; + } + last = x; + row_changed += 1; + } + if let Some(target) = diff_row.as_deref_mut() { + let at = x * 3; + target[at..at + 3].copy_from_slice(if changed { &CHANGED } else { &BLANK }); + } + } + + if row_changed > 0 { + changed_pixels += row_changed; + change_box.add_row(y, first, last); + } + } + + ensure!(!is_cancelled(), "comparison cancelled"); + let output = match (output, diff) { + (Some(path), Some(pixels)) => { + write_diff(path, union_width, union_height, &pixels)?; + Some(absolute_output(path)) + } + _ => None, + }; + let changed_fraction = if total_pixels == 0 { + 0.0 + } else { + let changed = + u32::try_from(changed_pixels).context("changed pixel count exceeds image limits")?; + let total = u32::try_from(total_pixels).context("pixel count exceeds image limits")?; + f64::from(changed) / f64::from(total) + }; + + Ok(Comparison { + before: Dimensions { + width: before_width, + height: before_height, + }, + after: Dimensions { + width: after_width, + height: after_height, + }, + same_dimensions: before_width == after_width && before_height == after_height, + changed_pixels, + total_pixels, + changed_fraction, + bounds: change_box.finish(), + output, + }) +} + +fn decode(path: &Path, role: &str) -> Result { + let mut reader = ImageReader::open(path) + .with_context(|| format!("failed to open {role} image {}", path.display()))? + .with_guessed_format() + .with_context(|| format!("failed to read {role} image {}", path.display()))?; + + // Detect the format from file content, not the extension. + match reader.format() { + Some(ImageFormat::Png | ImageFormat::Jpeg | ImageFormat::WebP) => {} + Some(other) => bail!( + "{role} image {} is {other:?}; iris compares PNG, JPEG and WebP", + path.display() + ), + None => bail!( + "{role} image {} is in no format iris knows; it compares PNG, JPEG and WebP", + path.display() + ), + } + + let mut limits = Limits::no_limits(); + limits.max_image_width = Some(MAX_DIMENSION); + limits.max_image_height = Some(MAX_DIMENSION); + limits.max_alloc = Some(MAX_DECODE_BYTES); + reader.limits(limits); + + let decoded = reader + .decode() + .with_context(|| format!("failed to decode {role} image {}", path.display()))?; + let rgba_bytes = u64::from(decoded.width()) * u64::from(decoded.height()) * 4; + if rgba_bytes > MAX_DECODE_BYTES { + bail!( + "{role} image needs {rgba_bytes} bytes after RGBA conversion; compare a smaller capture" + ); + } + Ok(decoded.into_rgba8()) +} + +/// One RGBA8 row, or `None` when `y` is outside the image. +#[inline] +fn row(raw: &[u8], width: usize, height: usize, y: usize) -> Option<&[u8]> { + if y >= height { + return None; + } + let stride = width * 4; + let start = y * stride; + Some(&raw[start..start + stride]) +} + +/// Composites an RGBA8 pixel on white. +#[inline] +fn on_white(pixel: &[u8]) -> [u8; 3] { + let alpha = u32::from(pixel[3]); + if alpha == 255 { + return [pixel[0], pixel[1], pixel[2]]; + } + let white = 255 * (255 - alpha); + [pixel[0], pixel[1], pixel[2]].map(|channel| { + // SAFETY: channel * alpha + white <= 255 * 255, so the rounded quotient fits in a u8. + u8::try_from((u32::from(channel) * alpha + white + 127) / 255) + .expect("alpha compositing must stay within a channel") + }) +} + +/// Largest per-channel gap between two composited pixels. +#[inline] +fn distance(before: [u8; 3], after: [u8; 3]) -> u8 { + before[0] + .abs_diff(after[0]) + .max(before[1].abs_diff(after[1])) + .max(before[2].abs_diff(after[2])) +} + +/// Fades an unchanged pixel in difference output. +#[inline] +fn ghost(pixel: [u8; 3]) -> [u8; 3] { + let luma = + (u32::from(pixel[0]) * 299 + u32::from(pixel[1]) * 587 + u32::from(pixel[2]) * 114) / 1000; + // SAFETY: weights sum to 1000, so luma and its faded value are within 0..=255. + let gray = u8::try_from(255 - (255 - luma) / 4).expect("luma must stay within a channel"); + [gray; 3] +} + +fn write_diff(path: &Path, width: u32, height: u32, pixels: &[u8]) -> Result<()> { + if let Some(parent) = path.parent().filter(|path| !path.as_os_str().is_empty()) { + std::fs::create_dir_all(parent) + .with_context(|| format!("failed to create {}", parent.display()))?; + } + + let file = File::create(path).with_context(|| format!("failed to write {}", path.display()))?; + let mut writer = BufWriter::new(file); + PngEncoder::new_with_quality(&mut writer, CompressionType::Fast, FilterType::Adaptive) + .write_image(pixels, width, height, ExtendedColorType::Rgb8) + .with_context(|| format!("failed to write {}", path.display()))?; + writer + .flush() + .with_context(|| format!("failed to write {}", path.display())) +} + +struct ChangeBox { + any: bool, + min_x: usize, + max_x: usize, + min_y: usize, + max_y: usize, +} + +impl ChangeBox { + const fn new() -> Self { + Self { + any: false, + min_x: usize::MAX, + max_x: 0, + min_y: 0, + max_y: 0, + } + } + + fn add_row(&mut self, y: usize, first: usize, last: usize) { + if !self.any { + self.any = true; + self.min_y = y; + } + self.max_y = y; + self.min_x = self.min_x.min(first); + self.max_x = self.max_x.max(last); + } + + fn finish(&self) -> Option { + if !self.any { + return None; + } + Some(Bounds { + x: u32::try_from(self.min_x).expect("change bounds are within image limits"), + y: u32::try_from(self.min_y).expect("change bounds are within image limits"), + width: u32::try_from(self.max_x - self.min_x + 1) + .expect("change bounds are within image limits"), + height: u32::try_from(self.max_y - self.min_y + 1) + .expect("change bounds are within image limits"), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::PathBuf; + + struct TempDir(PathBuf); + + impl TempDir { + fn new(label: &str) -> Self { + let path = + std::env::temp_dir().join(format!("iris-compare-{label}-{}", std::process::id())); + std::fs::create_dir_all(&path).expect("create temp dir"); + Self(path) + } + + fn join(&self, name: &str) -> PathBuf { + self.0.join(name) + } + } + + impl Drop for TempDir { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } + } + + /// Encodes a PNG fixture. + fn write_rgba(path: &Path, width: u32, height: u32, pixels: &[u8]) { + let mut writer = BufWriter::new(File::create(path).expect("create fixture")); + PngEncoder::new(&mut writer) + .write_image(pixels, width, height, ExtendedColorType::Rgba8) + .expect("encode fixture"); + writer.flush().expect("flush fixture"); + } + + #[test] + fn transparent_pixels_match_the_white_behind_them() { + let dir = TempDir::new("alpha"); + let before = dir.join("before.png"); + let after = dir.join("after.png"); + + // Transparent source colors do not affect composited output. + write_rgba(&before, 2, 1, &[0, 0, 0, 0, 17, 99, 200, 0]); + write_rgba(&after, 2, 1, &[255; 8]); + + let result = compare(&before, &after, 0, None, || false).expect("compare"); + + assert_eq!(result.changed_pixels, 0); + assert_eq!(result.bounds, None); + assert!(result.same_dimensions); + } + + #[test] + fn area_only_one_capture_covers_counts_as_changed() { + let dir = TempDir::new("union"); + let before = dir.join("before.png"); + let after = dir.join("after.png"); + + write_rgba(&before, 2, 2, &[255; 16]); + write_rgba(&after, 3, 2, &[255; 24]); + + let result = compare(&before, &after, 0, None, || false).expect("compare"); + + assert!(!result.same_dimensions); + // The extra column must change. + assert_eq!(result.total_pixels, 6); + assert_eq!(result.changed_pixels, 2); + assert_eq!( + result.bounds, + Some(Bounds { + x: 2, + y: 0, + width: 1, + height: 2 + }) + ); + } + + #[test] + fn threshold_forgives_its_own_distance_and_bounds_frame_the_rest() { + let dir = TempDir::new("threshold"); + let before = dir.join("before.png"); + let after = dir.join("after.png"); + + let white = [255u8; 36]; + let mut altered = white; + let at = |x: usize, y: usize| (y * 3 + x) * 4; + altered[at(1, 1)] = 245; + altered[at(2, 2)] = 0; + altered[at(2, 2) + 1] = 0; + + write_rgba(&before, 3, 3, &white); + write_rgba(&after, 3, 3, &altered); + + let forgiven = compare(&before, &after, 10, None, || false).expect("compare"); + assert_eq!(forgiven.changed_pixels, 1); + assert_eq!( + forgiven.bounds, + Some(Bounds { + x: 2, + y: 2, + width: 1, + height: 1 + }) + ); + + let counted = compare(&before, &after, 9, None, || false).expect("compare"); + assert_eq!(counted.changed_pixels, 2); + assert_eq!( + counted.bounds, + Some(Bounds { + x: 1, + y: 1, + width: 2, + height: 2 + }) + ); + } +} diff --git a/src/main.rs b/src/main.rs index 350ec49..665873f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,5 @@ mod capture; +mod compare; mod mcp; use std::collections::HashMap; @@ -11,8 +12,13 @@ use anyhow::{Context, Result, bail}; use capture::{CaptureMode, Format, Opts, Session, normalize_url, parse_viewport, success_report}; use clap::{Args, Parser, Subcommand}; use futures::StreamExt; +use num_traits::ToPrimitive; use url::Url; +// Keep short-deadline browser tests from competing for Chrome resources. +#[cfg(test)] +static BROWSER_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + /// Screenshot live websites. Minimal interface, powerful engine: /// smart waiting, lazy-load handling, retina output, concurrent capture. #[derive(Parser)] @@ -27,7 +33,7 @@ Examples: iris example.com 1440\u{d7}900 @2x \u{2192} example.com.png iris --full --dark tailwindcss.com full page, dark color scheme iris --selector '#hero' --padding 24 app.dev - iris mcp serve the capture tool over stdio + iris mcp serve capture and comparison tools over stdio cat urls.txt | iris - -o shots/ concurrent batch from stdin" )] struct Cli { @@ -42,6 +48,20 @@ struct Cli { enum Command { /// Serve Iris as a local Model Context Protocol camera Mcp(mcp::McpArgs), + /// Compare saved images and optionally write a visual difference map + Compare(CompareArgs), +} + +#[derive(Args)] +struct CompareArgs { + before: PathBuf, + after: PathBuf, + /// Ignore per-channel differences up to this value (0–255) + #[arg(long, default_value_t = 0)] + threshold: u8, + /// Save a PNG difference map + #[arg(short, long)] + output: Option, } #[derive(Args)] @@ -78,7 +98,19 @@ struct CaptureArgs { #[arg(long, value_parser = ["png", "jpg", "jpeg", "webp"], value_name = "FMT")] format: Option, - /// Extra settle delay in ms after smart waiting + /// JPEG/WebP quality (0–100) + #[arg(long, default_value_t = 90, value_parser = clap::value_parser!(u8).range(0..=100))] + quality: u8, + + /// Reduce image density to fit this pixel budget without changing page layout + #[arg(long, value_parser = clap::value_parser!(u64).range(1..))] + max_pixels: Option, + + /// Finish finite animations and pause repeating ones before capture + #[arg(long)] + freeze_animations: bool, + + /// Extra delay in ms before the final readiness check #[arg(long, default_value_t = 0, value_name = "MS")] wait: u64, @@ -110,10 +142,21 @@ struct CaptureArgs { #[tokio::main] async fn main() -> Result<()> { let cli = Cli::parse(); - if let Some(Command::Mcp(args)) = cli.command { - return mcp::run(args).await; + match cli.command { + Some(Command::Mcp(args)) => mcp::run(args).await, + Some(Command::Compare(args)) => { + let report = compare::compare( + &args.before, + &args.after, + args.threshold, + args.output.as_deref(), + || false, + )?; + println!("{}", serde_json::to_string(&report)?); + Ok(()) + } + None => run_capture(cli.capture).await, } - run_capture(cli.capture).await } async fn run_capture(cli: CaptureArgs) -> Result<()> { @@ -132,6 +175,9 @@ async fn run_capture(cli: CaptureArgs) -> Result<()> { wait_for: cli.wait_for, timeout: Duration::from_secs(cli.timeout.max(1)), format, + quality: cli.quality, + max_pixels: cli.max_pixels, + freeze_animations: cli.freeze_animations, }); let session = Arc::new(Session::launch(cli.chrome, viewport).await?); @@ -307,7 +353,10 @@ fn derived_name(url: &Url) -> String { fn human_size(bytes: u64) -> String { if bytes >= 1024 * 1024 { - format!("{:.1} MB", bytes as f64 / (1024.0 * 1024.0)) + format!( + "{:.1} MB", + bytes.to_f64().unwrap_or(f64::INFINITY) / (1024.0 * 1024.0) + ) } else { format!("{} KB", bytes.div_ceil(1024)) } @@ -318,7 +367,6 @@ mod tests { use clap::error::ErrorKind; use super::*; - use crate::capture::Shot; #[test] fn selector_conflicts_with_full_page() { @@ -342,32 +390,6 @@ mod tests { assert_eq!(negative.kind(), ErrorKind::ValueValidation); } - #[test] - fn capture_mode_is_constructed_once_from_validated_cli() { - let element = Cli::try_parse_from([ - "iris", - "example.com", - "--selector", - "main > h1", - "--padding", - "24", - ]) - .unwrap(); - assert_eq!( - capture_mode(&element.capture), - CaptureMode::Element { - selector: "main > h1".into(), - padding: 24, - } - ); - - let full = Cli::try_parse_from(["iris", "example.com", "--full"]).unwrap(); - assert_eq!(capture_mode(&full.capture), CaptureMode::FullPage); - - let viewport = Cli::try_parse_from(["iris", "example.com"]).unwrap(); - assert_eq!(capture_mode(&viewport.capture), CaptureMode::Viewport); - } - #[test] fn mcp_subcommand_does_not_require_a_url() { let cli = Cli::try_parse_from(["iris", "mcp"]).unwrap(); @@ -378,42 +400,4 @@ mod tests { .unwrap(); assert_eq!(conflict.kind(), ErrorKind::UnknownArgument); } - - #[test] - fn json_reports_have_stable_typed_fields() { - let mode = CaptureMode::Element { - selector: "h1".into(), - padding: 24, - }; - let shot = Shot { - width: 180, - height: 72, - scale: 2.0, - bytes: 14_231, - }; - let success = serde_json::to_string(&success_report( - "https://example.com/", - Some(std::path::Path::new("/tmp/example.png")), - &mode, - Format::Png, - &shot, - )) - .unwrap(); - assert_eq!( - success, - r#"{"status":"ok","url":"https://example.com/","output":"/tmp/example.png","mode":"element","selector":"h1","padding":24,"css_width":180,"css_height":72,"scale":2.0,"format":"png","bytes":14231}"# - ); - - let failure = serde_json::to_string(&capture::error_report( - "https://example.com/", - Some(std::path::Path::new("/tmp/example.png")), - &mode, - "selector never appeared: h1".into(), - )) - .unwrap(); - assert_eq!( - failure, - r#"{"status":"error","url":"https://example.com/","output":"/tmp/example.png","mode":"element","selector":"h1","padding":24,"error":"selector never appeared: h1"}"# - ); - } } diff --git a/src/mcp.rs b/src/mcp.rs index 9dd38a5..bacfafe 100644 --- a/src/mcp.rs +++ b/src/mcp.rs @@ -3,12 +3,12 @@ use std::sync::Arc; use std::time::Duration; use anyhow::{Context, Result, anyhow, bail}; -use base64::Engine as _; -use base64::engine::general_purpose::STANDARD as BASE64; use clap::Args; -use rmcp::handler::server::{router::tool::ToolRouter, wrapper::Parameters}; +use futures::future::join_all; +use rmcp::handler::server::{router::tool::ToolRouter, tool::ToolCallContext, wrapper::Parameters}; use rmcp::model::{CallToolResult, ContentBlock, Implementation, ServerCapabilities, ServerInfo}; -use rmcp::{ServerHandler, ServiceExt, schemars, tool, tool_handler, tool_router}; +use rmcp::service::RequestContext; +use rmcp::{RoleServer, ServerHandler, ServiceExt, schemars, tool, tool_handler, tool_router}; use serde::Deserialize; use serde_json::json; use tokio::sync::{Mutex, Semaphore}; @@ -43,19 +43,44 @@ pub struct CaptureRequest { dark: bool, /// Image format. Defaults to png; a recognized output extension wins. format: Option, - /// Extra settle delay in milliseconds after smart waiting. + /// JPEG/WebP quality from 0 to 100. Defaults to 90. + quality: Option, + /// Reduce image density to fit this pixel budget without changing page layout. + max_pixels: Option, + /// Finish finite animations and pause repeating ones before capture. + #[serde(default)] + freeze_animations: bool, + /// Extra delay in milliseconds before the final readiness check. #[serde(default)] wait_ms: u64, /// Wait until this CSS selector exists before capturing. wait_for: Option, /// Device scale factor overriding the viewport preset. scale: Option, - /// Per-page timeout in seconds. Defaults to 30. + /// Deadline including queue and browser startup, in seconds. Defaults to 30. timeout_seconds: Option, /// Optional image path. Relative paths resolve from the MCP server working directory. output: Option, } +#[derive(Debug, Deserialize, schemars::JsonSchema)] +pub struct BatchRequest { + /// Up to 16 captures, run concurrently. Results stay in request order. + captures: Vec, +} + +#[derive(Debug, Deserialize, schemars::JsonSchema)] +pub struct CompareRequest { + /// Path to the baseline image (PNG, JPEG, or WebP). + before: PathBuf, + /// Path to the new image (PNG, JPEG, or WebP). + after: PathBuf, + /// Ignore channel differences up to this value (0–255). Defaults to 0. + threshold: Option, + /// Optional PNG path for a visual difference map. + output: Option, +} + #[derive(Clone, Copy, Debug, Deserialize, schemars::JsonSchema)] #[serde(rename_all = "lowercase")] pub enum ImageFormat { @@ -137,6 +162,13 @@ impl CaptureRequest { .ok_or_else(|| anyhow!("unsupported output extension: .{extension}"))?; } + let quality = self.quality.unwrap_or(90); + if quality > 100 { + bail!("quality must be between 0 and 100"); + } + if self.max_pixels == Some(0) { + bail!("max_pixels must be greater than zero"); + } Ok(PreparedCapture { url, output: self.output, @@ -148,6 +180,9 @@ impl CaptureRequest { wait_for, timeout: Duration::from_secs(timeout_seconds), format, + quality, + max_pixels: self.max_pixels, + freeze_animations: self.freeze_animations, }, }) } @@ -156,7 +191,7 @@ impl CaptureRequest { struct McpState { chrome: Option, session: Mutex>>, - permits: Semaphore, + permits: Arc, } impl McpState { @@ -164,7 +199,7 @@ impl McpState { Self { chrome, session: Mutex::new(None), - permits: Semaphore::new(4), + permits: Arc::new(Semaphore::new(4)), } } @@ -185,23 +220,32 @@ impl McpState { } async fn capture(&self, url: &str, opts: &Opts) -> Result { - let _permit = self - .permits - .acquire() - .await - .map_err(|_| anyhow!("capture queue closed"))?; - let session = self.session().await?; - let result = session.capture(url, opts).await; - if result.is_err() && !session.is_healthy() { - let mut current = self.session.lock().await; - if current - .as_ref() - .is_some_and(|candidate| Arc::ptr_eq(candidate, &session)) - { - current.take(); + tokio::time::timeout(opts.timeout, async { + let _permit = self + .permits + .acquire() + .await + .map_err(|_| anyhow!("capture queue closed"))?; + let session = self.session().await?; + let result = session.capture(url, opts).await; + if result.is_err() && !session.is_healthy() { + let mut current = self.session.lock().await; + if current + .as_ref() + .is_some_and(|candidate| Arc::ptr_eq(candidate, &session)) + { + current.take(); + } } - } - result + result + }) + .await + .map_err(|_| { + anyhow!( + "capture deadline exceeded ({}s, including queue and browser startup)", + opts.timeout.as_secs() + ) + })? } async fn close(&self) { @@ -278,7 +322,7 @@ impl IrisServer { destination, ); let mut result = CallToolResult::success(vec![ - ContentBlock::image(BASE64.encode(&image.data), prepared.opts.format.mime_type()), + ContentBlock::image(image.data, prepared.opts.format.mime_type()), ContentBlock::text(summary), ]); result.structured_content = Some(structured); @@ -302,15 +346,114 @@ impl IrisServer { async fn capture(&self, Parameters(request): Parameters) -> CallToolResult { self.capture_image(request).await } + + /// Capture several pages, viewports, or elements in one request using the warm browser. + #[tool( + name = "capture_batch", + annotations( + read_only_hint = false, + destructive_hint = true, + open_world_hint = true + ) + )] + async fn capture_batch(&self, Parameters(request): Parameters) -> CallToolResult { + if request.captures.is_empty() || request.captures.len() > 16 { + return simple_error("captures must contain between 1 and 16 requests".into()); + } + let results = join_all( + request + .captures + .into_iter() + .map(|request| self.capture_image(request)), + ) + .await; + let mut content = Vec::new(); + let mut reports = Vec::with_capacity(results.len()); + let mut failed = 0; + for (index, result) in results.into_iter().enumerate() { + if result.is_error == Some(true) { + failed += 1; + } + content.push(ContentBlock::text(format!("Capture {}", index + 1))); + content.extend(result.content); + reports.push(result.structured_content); + } + let status = if failed == 0 { "ok" } else { "partial_error" }; + let mut result = CallToolResult::success(content); + result.is_error = Some(failed > 0); + result.structured_content = + Some(json!({ "status": status, "failed": failed, "captures": reports })); + result + } + + /// Compare saved screenshots by pixels, reporting changed area and an optional difference map. + #[tool( + name = "compare", + annotations( + read_only_hint = false, + destructive_hint = true, + open_world_hint = false + ) + )] + async fn compare( + &self, + Parameters(request): Parameters, + context: RequestContext, + ) -> CallToolResult { + let permit = match Arc::clone(&self.state.permits).acquire_owned().await { + Ok(permit) => permit, + Err(error) => return simple_error(format!("comparison queue closed: {error}")), + }; + let comparison = tokio::task::spawn_blocking(move || { + let _permit = permit; + crate::compare::compare( + &request.before, + &request.after, + request.threshold.unwrap_or(0), + request.output.as_deref(), + || context.ct.is_cancelled(), + ) + }) + .await; + match comparison { + Ok(Ok(report)) => { + let summary = format!( + "{} of {} pixels changed ({:.2}%)", + report.changed_pixels, + report.total_pixels, + report.changed_fraction * 100.0 + ); + let mut result = CallToolResult::success(vec![ContentBlock::text(summary)]); + result.structured_content = serde_json::to_value(report).ok(); + result + } + Ok(Err(error)) => simple_error(format!("{error:#}")), + Err(error) => simple_error(format!("image comparison failed: {error}")), + } + } } #[tool_handler(router = self.tool_router)] impl ServerHandler for IrisServer { + async fn call_tool( + &self, + request: rmcp::model::CallToolRequestParams, + context: RequestContext, + ) -> Result { + let cancelled = context.ct.clone(); + let call = ToolCallContext::new(self, request, context); + tokio::select! { + biased; + () = cancelled.cancelled() => Ok(simple_error("request cancelled".into()).into()), + result = self.tool_router.call(call) => result, + } + } + fn get_info(&self) -> ServerInfo { ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) .with_server_info(Implementation::new("iris", env!("CARGO_PKG_VERSION"))) .with_instructions( - "Iris is a camera for coding agents. Use capture for one page or element at a time; it returns the image inline and saves a file only when output is provided.", + "Iris captures screenshots without clicks or login. Use capture for one page or element, capture_batch for independent views, and compare for saved before/after images. Captures include timing, asset warnings, and the final URL. Use max_pixels to control image size, freeze_animations to finish finite animations and pause repeating ones, and output only when a durable file is needed.", ) } } @@ -384,6 +527,8 @@ async fn shutdown_signal() -> Result<()> { #[cfg(test)] mod tests { use super::*; + use base64::Engine as _; + use base64::engine::general_purpose::STANDARD as BASE64; fn request(url: &str) -> CaptureRequest { CaptureRequest { @@ -394,6 +539,9 @@ mod tests { size: None, dark: false, format: None, + quality: None, + max_pixels: None, + freeze_animations: false, wait_ms: 0, wait_for: None, scale: None, @@ -403,14 +551,7 @@ mod tests { } #[test] - fn request_defaults_and_conflicts_are_validated() { - let prepared = request("localhost:3000").prepare().unwrap(); - assert_eq!(prepared.url.as_str(), "http://localhost:3000/"); - assert_eq!(prepared.opts.mode, CaptureMode::Viewport); - assert_eq!(prepared.opts.viewport.width, 1440); - assert_eq!(prepared.opts.format, Format::Png); - assert_eq!(prepared.opts.timeout, Duration::from_secs(30)); - + fn conflicting_capture_modes_are_rejected() { let mut conflict = request("example.com"); conflict.selector = Some("main".into()); conflict.full_page = true; @@ -442,27 +583,9 @@ mod tests { assert_eq!(prepared.opts.format, Format::Webp); } - #[test] - fn server_advertises_one_focused_capture_tool() { - let server = IrisServer::new(Arc::new(McpState::new(None))); - let tools = server.tool_router.list_all(); - assert_eq!(tools.len(), 1); - let tool = &tools[0]; - assert_eq!(tool.name, "capture"); - assert_eq!( - tool.input_schema - .get("required") - .and_then(|required| required.as_array()) - .unwrap(), - &[serde_json::Value::String("url".into())] - ); - let annotations = tool.annotations.as_ref().unwrap(); - assert_eq!(annotations.read_only_hint, Some(false)); - assert_eq!(annotations.open_world_hint, Some(true)); - } - #[tokio::test] async fn capture_tool_returns_inline_pixels_metadata_and_optional_file() -> Result<()> { + let _browser_guard = crate::BROWSER_TEST_LOCK.lock().await; let temp = std::env::temp_dir().join(format!("iris-mcp-{}", std::process::id())); let output = temp.join("nested/target.png"); let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) @@ -478,7 +601,7 @@ mod tests { capture.padding = Some(10); capture.size = Some("320x240".into()); capture.scale = Some(1.0); - capture.timeout_seconds = Some(3); + capture.output = Some(output.clone()); let result = server.capture_image(capture).await; @@ -525,6 +648,50 @@ mod tests { Ok(()) } + #[tokio::test] + async fn batch_deadline_includes_time_queued_behind_earlier_captures() -> Result<()> { + let _browser_guard = crate::BROWSER_TEST_LOCK.lock().await; + let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/full-page-layout.html") + .canonicalize()?; + let url = url::Url::from_file_path(fixture) + .map_err(|_| anyhow!("fixture path is not a file URL"))?; + let state = Arc::new(McpState::new(None)); + let server = IrisServer::new(Arc::clone(&state)); + let captures = (0..5) + .map(|index| { + let mut capture = request(url.as_str()); + capture.size = Some("320x240".into()); + capture.scale = Some(1.0); + capture.timeout_seconds = Some(if index < 4 { 30 } else { 1 }); + capture.wait_ms = if index < 4 { 1500 } else { 0 }; + capture + }) + .collect(); + let result = server + .capture_batch(Parameters(BatchRequest { captures })) + .await; + state.close().await; + let report = result.structured_content.context("batch report missing")?; + let captures = report["captures"] + .as_array() + .context("capture reports missing")?; + assert!( + captures[..4] + .iter() + .all(|capture| capture["status"] == "ok"), + "{report}" + ); + assert_eq!(captures[4]["status"], "error"); + assert!( + captures[4]["error"] + .as_str() + .is_some_and(|error| error.contains("deadline exceeded")), + "{report}" + ); + Ok(()) + } + fn png_dimensions(bytes: &[u8]) -> (u32, u32) { assert_eq!(&bytes[..8], b"\x89PNG\r\n\x1a\n"); ( diff --git a/tests/cli.rs b/tests/cli.rs index a6d1222..17307c6 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -29,10 +29,10 @@ fn json_batch_reports_success_and_failure_on_stdout() { "8", "--scale", "1", - "--timeout", - "3", "--jobs", "2", + "--timeout", + "30", "--json", "-o", "shots", @@ -59,11 +59,11 @@ fn json_batch_reports_success_and_failure_on_stdout() { let success = reports .iter() .find(|report| report["status"] == "ok") - .unwrap(); + .unwrap_or_else(|| panic!("no successful capture in reports: {stdout}")); let failure = reports .iter() .find(|report| report["status"] == "error") - .unwrap(); + .unwrap_or_else(|| panic!("no failed capture in reports: {stdout}")); assert_eq!(success["mode"], "element"); assert_eq!(success["selector"], ".capture-target"); @@ -75,18 +75,12 @@ fn json_batch_reports_success_and_failure_on_stdout() { assert_eq!(failure["selector"], ".capture-target"); assert_eq!(failure["padding"], 8); assert!(failure["output"].as_str().unwrap().starts_with('/')); - assert!( - failure["error"] - .as_str() - .unwrap() - .contains("selector never appeared: .capture-target") - ); std::fs::remove_dir_all(temp).unwrap(); } #[test] -fn mcp_stdio_lists_one_tool_and_returns_an_inline_capture() { +fn mcp_stdio_returns_an_inline_capture() { let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) .join("tests/fixtures/precise-capture.html") .canonicalize() @@ -131,19 +125,6 @@ fn mcp_stdio_lists_one_tool_and_returns_an_inline_capture() { "method": "notifications/initialized" }), ); - send_message( - &mut stdin, - serde_json::json!({ - "jsonrpc": "2.0", - "id": 2, - "method": "tools/list", - "params": {} - }), - ); - let listed = receive_response(&receiver, 2, Duration::from_secs(5)); - let tools = listed["result"]["tools"].as_array().unwrap(); - assert_eq!(tools.len(), 1); - assert_eq!(tools[0]["name"], "capture"); send_message( &mut stdin, @@ -158,15 +139,14 @@ fn mcp_stdio_lists_one_tool_and_returns_an_inline_capture() { "selector": ".capture-target", "padding": 10, "size": "320x240", - "scale": 1, - "timeout_seconds": 3 + "scale": 1 } } }), ); - let captured = receive_response(&receiver, 3, Duration::from_secs(15)); + let captured = receive_response(&receiver, 3, Duration::from_secs(35)); let result = &captured["result"]; - assert_eq!(result["isError"], false); + assert_eq!(result["isError"], false, "{captured}"); assert_eq!(result["structuredContent"]["status"], "ok"); assert_eq!(result["structuredContent"]["css_width"], 141); assert_eq!(result["structuredContent"]["css_height"], 81); @@ -270,13 +250,13 @@ fn mcp_sigterm_closes_chrome_and_removes_its_profile() { "selector": ".capture-target", "size": "320x240", "scale": 1, - "timeout_seconds": 3 + "timeout_seconds": 30 } } }), ); - let captured = receive_response(&receiver, 2, Duration::from_secs(15)); - assert_eq!(captured["result"]["isError"], false); + let captured = receive_response(&receiver, 2, Duration::from_secs(35)); + assert_eq!(captured["result"]["isError"], false, "{captured}"); let signal = Command::new("kill") .args(["-TERM", &child.id().to_string()]) diff --git a/tests/fixtures/full-page-layout.html b/tests/fixtures/full-page-layout.html new file mode 100644 index 0000000..b446cc4 --- /dev/null +++ b/tests/fixtures/full-page-layout.html @@ -0,0 +1,10 @@ + + +Viewport-relative full-page layout + +
diff --git a/tools/ast-grep/anti-slop/rules/rust-no-dyn-any.yml b/tools/ast-grep/anti-slop/rules/rust-no-dyn-any.yml new file mode 100644 index 0000000..3857462 --- /dev/null +++ b/tools/ast-grep/anti-slop/rules/rust-no-dyn-any.yml @@ -0,0 +1,16 @@ +id: anti-slop-rust-no-dyn-any +language: rust +severity: error +message: Do not erase types behind `dyn Any` or downcast out of it. +note: | + `dyn Any` throws away the type system the same way `unknown` does in TypeScript. + Give the value a real enum or trait with the methods callers need instead of + boxing it as `Any` and downcasting at every call site. +rule: + any: + - pattern: + context: "fn f(x: Box) {}" + selector: dynamic_type + - pattern: $X.downcast_ref::<$T>() + - pattern: $X.downcast::<$T>() + - pattern: $X.downcast_mut::<$T>() diff --git a/tools/ast-grep/anti-slop/rules/rust-no-json-value-in-signatures.yml b/tools/ast-grep/anti-slop/rules/rust-no-json-value-in-signatures.yml new file mode 100644 index 0000000..63b5480 --- /dev/null +++ b/tools/ast-grep/anti-slop/rules/rust-no-json-value-in-signatures.yml @@ -0,0 +1,26 @@ +id: anti-slop-rust-no-json-value-in-signatures +language: rust +severity: error +message: Do not accept or return `serde_json::Value` at a function boundary. +note: | + `Value` at a signature is the same untyped hole as TypeScript's `unknown` + parameter or return type. Parse into a concrete struct at the boundary + (`#[derive(Deserialize)]`) and keep `Value` internal, if used at all. +rule: + any: + - all: + - any: + - kind: type_identifier + - kind: scoped_type_identifier + - regex: '(^|::)Value$' + - inside: + kind: parameter + field: type + - all: + - any: + - kind: type_identifier + - kind: scoped_type_identifier + - regex: '(^|::)Value$' + - inside: + kind: function_item + field: return_type diff --git a/tools/ast-grep/anti-slop/rules/rust-no-leftover-panic-markers.yml b/tools/ast-grep/anti-slop/rules/rust-no-leftover-panic-markers.yml new file mode 100644 index 0000000..bbe3849 --- /dev/null +++ b/tools/ast-grep/anti-slop/rules/rust-no-leftover-panic-markers.yml @@ -0,0 +1,36 @@ +id: anti-slop-rust-no-leftover-panic-markers +language: rust +severity: error +message: Remove this leftover `todo!`/`unimplemented!`/`dbg!` marker before shipping. +note: | + These are development scaffolding, not a finished implementation. Finish + the branch, return a real `Result`/`Option`, or use `tracing`/`log` instead + of `dbg!`. Test code is exempt. +rule: + all: + - kind: macro_invocation + - has: + field: macro + regex: '^(todo|unimplemented|dbg)$' + - not: + inside: + stopBy: end + any: + - all: + - kind: function_item + - follows: + kind: attribute_item + stopBy: neighbor + has: + kind: attribute + has: + kind: identifier + regex: '^test$' + - all: + - kind: mod_item + - follows: + kind: attribute_item + stopBy: neighbor + has: + kind: attribute + regex: 'cfg\s*\(\s*test\s*\)' diff --git a/tools/ast-grep/anti-slop/rules/rust-no-numeric-as-cast.yml b/tools/ast-grep/anti-slop/rules/rust-no-numeric-as-cast.yml new file mode 100644 index 0000000..fea6c31 --- /dev/null +++ b/tools/ast-grep/anti-slop/rules/rust-no-numeric-as-cast.yml @@ -0,0 +1,14 @@ +id: anti-slop-rust-no-numeric-as-cast +language: rust +severity: error +message: Do not use `as` to cast between numeric types; it truncates silently. +note: | + `as` between numeric types wraps or truncates without telling you. Use + `TryFrom`/`try_into` and handle the error, or a checked/saturating + conversion when truncation is genuinely intended and say so. +rule: + kind: type_cast_expression + has: + field: type + kind: primitive_type + regex: '^(u8|u16|u32|u64|u128|usize|i8|i16|i32|i64|i128|isize|f32|f64)$' diff --git a/tools/ast-grep/anti-slop/rules/rust-no-shape-in-symbol-names.yml b/tools/ast-grep/anti-slop/rules/rust-no-shape-in-symbol-names.yml new file mode 100644 index 0000000..cf98a2c --- /dev/null +++ b/tools/ast-grep/anti-slop/rules/rust-no-shape-in-symbol-names.yml @@ -0,0 +1,31 @@ +id: anti-slop-rust-no-shape-in-symbol-names +language: rust +severity: error +message: Rename this symbol for its domain role instead of its container shape. +note: | + A `_map`/`_vec`/`_string`/`_obj`/`_arr` suffix describes the container, not + what the value means. Rename for ownership (`users`, `config`, `pending`); + the type annotation already says it's a `HashMap` or `Vec`. Scoped to + declaration sites (`let` bindings, parameters, struct fields) rather than + every identifier, because Rust's own stdlib uses these suffixes as verbs + (`to_string`, `to_vec`, `flat_map`, `filter_map`), not as shape nouns. +rule: + any: + - all: + - kind: identifier + - regex: '_(map|vec|string|obj|arr)$' + - inside: + kind: let_declaration + field: pattern + - all: + - kind: identifier + - regex: '_(map|vec|string|obj|arr)$' + - inside: + kind: parameter + field: pattern + - all: + - kind: field_identifier + - regex: '_(map|vec|string|obj|arr)$' + - inside: + kind: field_declaration + field: name diff --git a/tools/ast-grep/anti-slop/rules/rust-no-transmute.yml b/tools/ast-grep/anti-slop/rules/rust-no-transmute.yml new file mode 100644 index 0000000..efbc616 --- /dev/null +++ b/tools/ast-grep/anti-slop/rules/rust-no-transmute.yml @@ -0,0 +1,13 @@ +id: anti-slop-rust-no-transmute +language: rust +severity: error +message: Do not use `mem::transmute`/`transmute_copy` to reinterpret bytes. +note: | + `transmute` bypasses every layout and validity check the compiler has. + Use a safe conversion (`From`/`TryFrom`, `to_ne_bytes`/`from_ne_bytes`, + `as` between numeric types, or a checked cast crate) instead. +rule: + kind: call_expression + has: + field: function + regex: '(^|::)transmute(_copy)?$' diff --git a/tools/ast-grep/anti-slop/rules/rust-no-unchecked-unwrap.yml b/tools/ast-grep/anti-slop/rules/rust-no-unchecked-unwrap.yml new file mode 100644 index 0000000..36b1d18 --- /dev/null +++ b/tools/ast-grep/anti-slop/rules/rust-no-unchecked-unwrap.yml @@ -0,0 +1,65 @@ +id: anti-slop-rust-no-unchecked-unwrap +language: rust +severity: error +message: This `.unwrap()`/`.expect(..)` has no justification for why it cannot panic. +note: | + State why the panic path is impossible or acceptable with a `// SAFETY:` + (invariant holds) or `// PANIC:` (crash is the intended behavior) comment + directly above the statement. Otherwise propagate the error with `?`. +rule: + any: + - all: + - kind: let_declaration + - has: &unwrap_or_expect_call + stopBy: end + kind: call_expression + has: + field: function + kind: field_expression + has: + field: field + regex: '^(unwrap|expect)$' + - not: &has_safety_or_panic_comment + follows: + kind: line_comment + stopBy: neighbor + regex: '(SAFETY|PANIC)\s*:' + - not: &inside_test_code + inside: + stopBy: end + any: + - all: + - kind: function_item + - follows: + kind: attribute_item + stopBy: neighbor + has: + kind: attribute + has: + kind: identifier + regex: '^test$' + - all: + - kind: mod_item + - follows: + kind: attribute_item + stopBy: neighbor + has: + kind: attribute + regex: 'cfg\s*\(\s*test\s*\)' + - all: + - kind: call_expression + - inside: + kind: block + - has: + field: function + kind: field_expression + has: + field: field + regex: '^(unwrap|expect)$' + - not: *has_safety_or_panic_comment + - not: *inside_test_code + - all: + - kind: expression_statement + - has: *unwrap_or_expect_call + - not: *has_safety_or_panic_comment + - not: *inside_test_code diff --git a/tools/ast-grep/anti-slop/rules/rust-no-unsafe-dictionary-type.yml b/tools/ast-grep/anti-slop/rules/rust-no-unsafe-dictionary-type.yml new file mode 100644 index 0000000..6f07298 --- /dev/null +++ b/tools/ast-grep/anti-slop/rules/rust-no-unsafe-dictionary-type.yml @@ -0,0 +1,25 @@ +id: anti-slop-rust-no-unsafe-dictionary-type +language: rust +severity: error +message: Do not use a map keyed to `Value` or `Box` as a value type. +note: | + Same problem as TypeScript's `Record`: the map's value type + claims nothing. Define a struct for the fields you actually read, or an enum + for the closed set of shapes the map can hold. +rule: + all: + - kind: generic_type + - has: + field: type + regex: '(^|::)(HashMap|BTreeMap)$' + - has: + field: type_arguments + has: + stopBy: end + any: + - kind: type_identifier + regex: '^Value$' + - kind: scoped_type_identifier + regex: '(^|::)Value$' + - kind: dynamic_type + regex: 'Any' diff --git a/tools/ast-grep/anti-slop/rules/rust-require-safety-comment-for-unsafe.yml b/tools/ast-grep/anti-slop/rules/rust-require-safety-comment-for-unsafe.yml new file mode 100644 index 0000000..39c9b32 --- /dev/null +++ b/tools/ast-grep/anti-slop/rules/rust-require-safety-comment-for-unsafe.yml @@ -0,0 +1,29 @@ +id: anti-slop-rust-require-safety-comment-for-unsafe +language: rust +severity: error +message: This `unsafe` block has no `// SAFETY:` justification directly above it. +note: | + State the invariant the compiler cannot check immediately before the block + (`// SAFETY: `), the same way a TypeScript + type assertion needs a `SAFETY:` comment to explain what the checker can't see. +rule: + any: + - all: + - kind: expression_statement + - has: + kind: unsafe_block + - not: + follows: + kind: line_comment + stopBy: neighbor + regex: 'SAFETY\s*:' + - all: + - kind: let_declaration + - has: + field: value + kind: unsafe_block + - not: + follows: + kind: line_comment + stopBy: neighbor + regex: 'SAFETY\s*:' diff --git a/tools/ast-grep/anti-slop/sgconfig.yml b/tools/ast-grep/anti-slop/sgconfig.yml new file mode 100644 index 0000000..2dd8e53 --- /dev/null +++ b/tools/ast-grep/anti-slop/sgconfig.yml @@ -0,0 +1,2 @@ +ruleDirs: + - rules