Skip to main content

mermaid_cli/providers/tool/computer_use/
driver.rs

1//! `ComputerUseDriver` — the shared backend-dispatch layer for the
2//! seven computer-use tools.
3//!
4//! The driver wraps three things:
5//!
6//!   1. A `Backend` discriminant (`X11`, `Wayland`, `MacOS`, …). Tools
7//!      match on it to pick the right subprocess dispatch.
8//!   2. A bounded `ScreenshotRegistry`. Every capture gets a stable
9//!      `id`; the model includes that id on later `click(x, y,
10//!      screenshot_id)` so coordinate translation uses the right
11//!      scale+offset even if the newest screenshot has shifted the
12//!      "latest" entry.
13//!   3. `ensure_alive()` — a cheap re-probe called at the top of every
14//!      tool's `execute`. Catches the case where the display went
15//!      away between registration and invocation (detached SSH,
16//!      closed lid).
17//!
18//! Subprocess dispatch uses `tokio::process::Command` so each external
19//! binary can race against `ctx.token.cancelled()`. `kill_on_drop(true)`
20//! reaps children whose parent future gets cancelled.
21
22use std::collections::VecDeque;
23use std::path::PathBuf;
24use std::sync::Mutex;
25use std::sync::atomic::{AtomicU64, Ordering};
26
27use anyhow::{Context, Result};
28use base64::{Engine as _, engine::general_purpose};
29use tokio::process::Command;
30use tokio_util::sync::CancellationToken;
31
32use crate::constants::{SCREENSHOT_MAX_WIDTH, SCREENSHOT_REGISTRY_CAPACITY};
33
34use super::Backend;
35
36/// Per-capture metadata retained so subsequent clicks can translate
37/// model-space coords back to screen-space.
38#[derive(Debug, Clone)]
39pub struct ScreenshotMetadata {
40    pub id: u64,
41    pub scale_factor: f64,
42    pub offset_x: i32,
43    pub offset_y: i32,
44    /// Displayed (post-downscale) pixel dimensions — the exact frame the model
45    /// reasoned over. `scale_coords` clamps model-supplied coords to
46    /// `[0, width) × [0, height)` before translating (#96); a value of `0`
47    /// means the PNG header was unreadable, so the upper clamp is skipped.
48    pub width: u32,
49    pub height: u32,
50    /// Human-readable capture kind (`"fullscreen"`, `"focused window"`,
51    /// …). Surfaced in error messages if the model references an
52    /// evicted id.
53    pub kind: String,
54}
55
56/// Bounded ring buffer of recent screenshot metadata. Capacity from
57/// `constants::SCREENSHOT_REGISTRY_CAPACITY` (= 16). When full, push
58/// evicts the oldest; referencing an evicted id fails cleanly with
59/// "take a fresh screenshot" rather than silently clicking at wrong
60/// coordinates.
61#[derive(Debug, Default)]
62pub struct ScreenshotRegistry {
63    entries: VecDeque<ScreenshotMetadata>,
64}
65
66impl ScreenshotRegistry {
67    pub fn new() -> Self {
68        Self {
69            entries: VecDeque::new(),
70        }
71    }
72
73    pub fn push(&mut self, meta: ScreenshotMetadata) {
74        if self.entries.len() >= SCREENSHOT_REGISTRY_CAPACITY {
75            self.entries.pop_front();
76        }
77        self.entries.push_back(meta);
78    }
79
80    pub fn get(&self, id: u64) -> Option<&ScreenshotMetadata> {
81        self.entries.iter().find(|m| m.id == id)
82    }
83
84    pub fn latest(&self) -> Option<&ScreenshotMetadata> {
85        self.entries.back()
86    }
87
88    pub fn len(&self) -> usize {
89        self.entries.len()
90    }
91
92    pub fn is_empty(&self) -> bool {
93        self.entries.is_empty()
94    }
95}
96
97/// What the screenshot tool accepts: which slice of the display to
98/// capture.
99#[derive(Debug, Clone)]
100pub enum ScreenshotSpec {
101    Fullscreen,
102    Focused,
103    Monitor(String),
104    /// `(x, y, width, height)` in screen pixels.
105    Region(i32, i32, u32, u32),
106    Window(String),
107}
108
109/// Result of a capture: encoded bytes + registry id + a human-readable
110/// summary for the tool's `output` field.
111#[derive(Debug)]
112pub struct CaptureResult {
113    pub id: u64,
114    pub base64_png: String,
115    pub raw_bytes: Vec<u8>,
116    pub width: u32,
117    pub height: u32,
118    pub scale_factor: f64,
119    pub offset_x: i32,
120    pub offset_y: i32,
121    pub summary: String,
122}
123
124/// Shared driver all seven computer-use tools hold an `Arc<>` to.
125pub struct ComputerUseDriver {
126    backend: Backend,
127    registry: Mutex<ScreenshotRegistry>,
128    /// Monotonic counter for temp file uniqueness. Distinct from the
129    /// registry id counter so filenames don't collide across runs
130    /// that share temp dir.
131    file_counter: AtomicU64,
132    /// Monotonic counter for registry ids — stable across process
133    /// lifetime, survives evictions.
134    id_counter: AtomicU64,
135}
136
137impl ComputerUseDriver {
138    pub fn new(backend: Backend) -> Self {
139        Self {
140            backend,
141            registry: Mutex::new(ScreenshotRegistry::new()),
142            file_counter: AtomicU64::new(0),
143            id_counter: AtomicU64::new(0),
144        }
145    }
146
147    pub fn backend(&self) -> Backend {
148        self.backend
149    }
150
151    /// Cheap mid-call liveness check. Tools call this first inside
152    /// `execute()`; if the display went away after registration,
153    /// they return a clean error instead of hanging on subprocess
154    /// dispatch.
155    pub fn ensure_alive(&self) -> Result<(), String> {
156        if super::display_is_reachable(self.backend) {
157            Ok(())
158        } else {
159            Err(format!(
160                "Display unreachable (backend={:?}). Was the session \
161                 detached, or did `DISPLAY` change?",
162                self.backend
163            ))
164        }
165    }
166
167    /// Async form of [`ensure_alive`]. The X11 display probe spawns a
168    /// `xdpyinfo`/`xdotool` subprocess and blocks on its exit; on the async
169    /// tool path that would block a worker, so run it on the blocking pool (#34).
170    pub async fn ensure_alive_async(&self) -> Result<(), String> {
171        let backend = self.backend;
172        match tokio::task::spawn_blocking(move || super::display_is_reachable(backend)).await {
173            Ok(true) => Ok(()),
174            Ok(false) => Err(format!(
175                "Display unreachable (backend={:?}). Was the session \
176                 detached, or did `DISPLAY` change?",
177                self.backend
178            )),
179            Err(_) => Err("display liveness probe failed to run".to_string()),
180        }
181    }
182
183    /// Translate model-space coords to screen-space using the metadata
184    /// registered for `screenshot_id` (or the latest if None).
185    pub fn scale_coords(
186        &self,
187        x: i32,
188        y: i32,
189        screenshot_id: Option<u64>,
190    ) -> Result<(i32, i32), String> {
191        let reg = self.registry.lock().map_err(|e| e.to_string())?;
192        let meta = match screenshot_id {
193            Some(id) => reg.get(id).cloned().ok_or_else(|| {
194                format!(
195                    "Screenshot id {} not found in registry (likely evicted — capacity {}). \
196                     Take a fresh screenshot and retry with the new id.",
197                    id, SCREENSHOT_REGISTRY_CAPACITY
198                )
199            })?,
200            None => reg.latest().cloned().ok_or_else(|| {
201                "No screenshots registered yet — call `screenshot` before \
202                 `click` / `mouse_move`."
203                    .to_string()
204            })?,
205        };
206        // #96: clamp the model-supplied coords into the frame's own pixel bounds
207        // before translating (skip the upper clamp when a dim is 0 = PNG header
208        // unreadable, so `clamp`'s min <= max precondition always holds).
209        let cx = if meta.width > 0 {
210            x.clamp(0, meta.width as i32 - 1)
211        } else {
212            x.max(0)
213        };
214        let cy = if meta.height > 0 {
215            y.clamp(0, meta.height as i32 - 1)
216        } else {
217            y.max(0)
218        };
219        // #32: the f64 -> i32 cast already saturates (NaN -> 0, ±huge ->
220        // i32::MIN/MAX); the `+ offset` i32 add is the real defect — it panics on
221        // overflow in debug and wraps to a negative coord in release.
222        // saturating_add mirrors the saturating_sub at driver.rs:452. The clamp
223        // above already bounds the result inside the region; this is
224        // defense-in-depth for extreme offsets.
225        Ok((
226            ((cx as f64 * meta.scale_factor) as i32).saturating_add(meta.offset_x),
227            ((cy as f64 * meta.scale_factor) as i32).saturating_add(meta.offset_y),
228        ))
229    }
230
231    /// Allocate a fresh registry id and record metadata.
232    pub fn register_screenshot(
233        &self,
234        scale_factor: f64,
235        offset_x: i32,
236        offset_y: i32,
237        width: u32,
238        height: u32,
239        kind: String,
240    ) -> u64 {
241        let id = self.id_counter.fetch_add(1, Ordering::Relaxed);
242        if let Ok(mut reg) = self.registry.lock() {
243            reg.push(ScreenshotMetadata {
244                id,
245                scale_factor,
246                offset_x,
247                offset_y,
248                width,
249                height,
250                kind,
251            });
252        }
253        id
254    }
255
256    /// Capture and return the encoded result. Respects cancellation
257    /// via `token.cancelled()` races in the subprocess wait.
258    pub async fn capture(
259        &self,
260        spec: ScreenshotSpec,
261        token: &CancellationToken,
262    ) -> Result<CaptureResult> {
263        self.ensure_alive_async()
264            .await
265            .map_err(|error| anyhow::anyhow!(error))?;
266
267        let seq = self.file_counter.fetch_add(1, Ordering::Relaxed);
268        // Write screenshots into the 0700 per-user scratch dir, not a
269        // world-readable fixed path in shared /tmp where another local user
270        // could read the captured frame (#33).
271        let temp_path =
272            crate::utils::private_temp_dir()?.join(format!("mermaid-screenshot-{}.png", seq));
273        let temp_str = temp_path.to_string_lossy().to_string();
274        let _guard = TempFileGuard(temp_path.clone());
275
276        let (offset_x, offset_y, kind) =
277            dispatch_capture(self.backend, &spec, &temp_str, token).await?;
278
279        // F59: dispatch_capture already races cancellation internally, but the
280        // downscale (a slow ImageMagick/ffmpeg subprocess capped at
281        // SCREENSHOT_DOWNSCALE_TIMEOUT_SECS) and the trailing read did not — an Esc
282        // mid-downscale would block until that timeout. Race both against the token
283        // so Esc aborts promptly. On cancel the dropped downscale future runs its
284        // own scaled-file guard, and `_guard` above removes the original temp file.
285        let scale_factor =
286            cancellable(token, downscale_if_needed(&temp_str, SCREENSHOT_MAX_WIDTH)).await?;
287
288        let raw_bytes = cancellable(token, async {
289            tokio::fs::read(&temp_path)
290                .await
291                .context("reading captured screenshot")
292        })
293        .await?;
294        let width = read_png_width(&raw_bytes).unwrap_or(0);
295        let height = read_png_height(&raw_bytes).unwrap_or(0);
296
297        // Register AFTER dims are known so `scale_coords` can clamp model coords
298        // to this frame's pixel bounds (#96).
299        let id = self.register_screenshot(
300            scale_factor,
301            offset_x,
302            offset_y,
303            width,
304            height,
305            kind.clone(),
306        );
307
308        let base64_png = general_purpose::STANDARD.encode(&raw_bytes);
309
310        let offset_info = if offset_x != 0 || offset_y != 0 {
311            format!(", offset: +{}+{}", offset_x, offset_y)
312        } else {
313            String::new()
314        };
315        let summary = format!(
316            "Screenshot captured (id: {}, {}, {}x{}, scale: {:.2}x{})",
317            id, kind, width, height, scale_factor, offset_info
318        );
319
320        Ok(CaptureResult {
321            id,
322            base64_png,
323            raw_bytes,
324            width,
325            height,
326            scale_factor,
327            offset_x,
328            offset_y,
329            summary,
330        })
331    }
332
333    /// Convenience for click/type/key tools: capture the focused
334    /// window and return `(summary, base64_png)` for inclusion in
335    /// the tool's auto-screenshot. Best-effort — on error returns
336    /// `None` and the caller can fall back to a screenshot-less
337    /// outcome.
338    pub async fn capture_focused_for_autoshot(
339        &self,
340        token: &CancellationToken,
341    ) -> Option<(String, String)> {
342        let cap = self.capture(ScreenshotSpec::Focused, token).await.ok()?;
343        Some((cap.summary, cap.base64_png))
344    }
345
346    /// X11-only: verify the cursor actually landed where xdotool was
347    /// told to move it. Returns `Some(warning)` if the cursor ended
348    /// up more than `CURSOR_LANDED_TOLERANCE_PX` away (focus change,
349    /// window moved, WM rejected the move). `None` if within
350    /// tolerance or the probe itself failed (best-effort — never
351    /// blocks the click).
352    pub async fn check_cursor_landed(&self, sx: i32, sy: i32) -> Option<String> {
353        if !matches!(self.backend, Backend::X11) {
354            return None;
355        }
356        let out = run_cmd_stdout(Command::new("xdotool").arg("getmouselocation"))
357            .await
358            .ok()?;
359        let mut actual_x: Option<i32> = None;
360        let mut actual_y: Option<i32> = None;
361        for tok in out.split_whitespace() {
362            if let Some(v) = tok.strip_prefix("X:") {
363                actual_x = v.parse().ok();
364            } else if let Some(v) = tok.strip_prefix("Y:") {
365                actual_y = v.parse().ok();
366            }
367        }
368        let (ax, ay) = (actual_x?, actual_y?);
369        if (ax - sx).abs() > CURSOR_LANDED_TOLERANCE_PX
370            || (ay - sy).abs() > CURSOR_LANDED_TOLERANCE_PX
371        {
372            Some(format!(
373                "WARNING: cursor at ({}, {}), expected ({}, {}). Window may have moved \
374                 or focus changed before the click landed.",
375                ax, ay, sx, sy
376            ))
377        } else {
378            None
379        }
380    }
381}
382
383/// HiDPI fractional scaling can put the cursor a pixel or two off the
384/// exact target; >5px means something other than rounding is wrong.
385const CURSOR_LANDED_TOLERANCE_PX: i32 = 5;
386
387// ───── action dispatch (shared by click / type / key / scroll / move / list) ──
388
389impl ComputerUseDriver {
390    /// Click at the given SCREEN coordinates (already scaled by
391    /// `scale_coords`). `button` is `"left" | "middle" | "right"`.
392    pub async fn click(
393        &self,
394        sx: i32,
395        sy: i32,
396        button: &str,
397        token: &CancellationToken,
398    ) -> Result<()> {
399        let code = match button {
400            "middle" => "2",
401            "right" => "3",
402            _ => "1",
403        };
404        match self.backend {
405            Backend::X11 => {
406                run_cmd_cancellable(
407                    Command::new("xdotool").args([
408                        "mousemove",
409                        "--sync",
410                        &sx.to_string(),
411                        &sy.to_string(),
412                        "click",
413                        "--clearmodifiers",
414                        code,
415                    ]),
416                    token,
417                )
418                .await
419            },
420            Backend::Wayland => {
421                if !super::has_command("ydotool") {
422                    anyhow::bail!("ydotool required for Wayland mouse control")
423                }
424                run_cmd_cancellable(
425                    Command::new("ydotool").args([
426                        "mousemove",
427                        "--absolute",
428                        "-x",
429                        &sx.to_string(),
430                        "-y",
431                        &sy.to_string(),
432                    ]),
433                    token,
434                )
435                .await?;
436                run_cmd_cancellable(
437                    Command::new("ydotool").args(["click", &format!("0x{}", code)]),
438                    token,
439                )
440                .await
441            },
442            _ => anyhow::bail!("click not supported on this platform"),
443        }
444    }
445
446    /// Type text at the current focus. Per-keystroke delay from
447    /// `TYPE_KEY_DELAY_MS` — empirically needed for slow Electron /
448    /// web targets that drop characters at lower rates.
449    pub async fn type_text(&self, text: &str, token: &CancellationToken) -> Result<()> {
450        let delay = crate::constants::TYPE_KEY_DELAY_MS.to_string();
451        match self.backend {
452            Backend::X11 => {
453                run_cmd_cancellable(
454                    Command::new("xdotool").args([
455                        "type",
456                        "--clearmodifiers",
457                        "--delay",
458                        &delay,
459                        text,
460                    ]),
461                    token,
462                )
463                .await
464            },
465            Backend::Wayland => {
466                if super::has_command("wtype") {
467                    run_cmd_cancellable(Command::new("wtype").arg(text), token).await
468                } else if super::has_command("ydotool") {
469                    run_cmd_cancellable(
470                        Command::new("ydotool").args(["type", "--delay", &delay, text]),
471                        token,
472                    )
473                    .await
474                } else {
475                    anyhow::bail!("wtype or ydotool required for Wayland text input")
476                }
477            },
478            _ => anyhow::bail!("type_text not supported on this platform"),
479        }
480    }
481
482    /// Press a key (or key combination like `"ctrl+shift+t"`).
483    pub async fn press_key(&self, key: &str, token: &CancellationToken) -> Result<()> {
484        match self.backend {
485            Backend::X11 => {
486                run_cmd_cancellable(Command::new("xdotool").args(["key", key]), token).await
487            },
488            Backend::Wayland => {
489                if super::has_command("wtype") {
490                    // wtype: -M/-m modifiers around -k final key.
491                    let parts: Vec<&str> = key.split('+').collect();
492                    let mut args: Vec<String> = Vec::new();
493                    for (i, part) in parts.iter().enumerate() {
494                        if i < parts.len() - 1 {
495                            args.push("-M".to_string());
496                            args.push(part.to_string());
497                        } else {
498                            args.push("-k".to_string());
499                            args.push(part.to_string());
500                        }
501                    }
502                    for part in parts.iter().take(parts.len().saturating_sub(1)) {
503                        args.push("-m".to_string());
504                        args.push(part.to_string());
505                    }
506                    run_cmd_cancellable(Command::new("wtype").args(&args), token).await
507                } else if super::has_command("ydotool") {
508                    run_cmd_cancellable(Command::new("ydotool").args(["key", key]), token).await
509                } else {
510                    anyhow::bail!("wtype or ydotool required for Wayland key input")
511                }
512            },
513            _ => anyhow::bail!("press_key not supported on this platform"),
514        }
515    }
516
517    /// Scroll `amount` ticks in `direction` ("up" / "down").
518    pub async fn scroll(
519        &self,
520        direction: &str,
521        amount: i32,
522        token: &CancellationToken,
523    ) -> Result<()> {
524        match self.backend {
525            Backend::X11 => {
526                // xdotool: button 4 = scroll up, 5 = scroll down.
527                let button = if direction == "up" { "4" } else { "5" };
528                let mut args: Vec<String> = Vec::new();
529                for _ in 0..amount {
530                    args.push("click".to_string());
531                    args.push(button.to_string());
532                }
533                run_cmd_cancellable(Command::new("xdotool").args(&args), token).await
534            },
535            Backend::Wayland => {
536                if !super::has_command("ydotool") {
537                    anyhow::bail!("ydotool required for Wayland scroll")
538                }
539                let wheel_amount = if direction == "up" { -amount } else { amount };
540                run_cmd_cancellable(
541                    Command::new("ydotool").args([
542                        "mousemove",
543                        "--wheel",
544                        &wheel_amount.to_string(),
545                    ]),
546                    token,
547                )
548                .await
549            },
550            _ => anyhow::bail!("scroll not supported on this platform"),
551        }
552    }
553
554    /// Move the mouse cursor to SCREEN coords (already scaled).
555    pub async fn mouse_move(&self, sx: i32, sy: i32, token: &CancellationToken) -> Result<()> {
556        match self.backend {
557            Backend::X11 => {
558                run_cmd_cancellable(
559                    Command::new("xdotool").args([
560                        "mousemove",
561                        "--sync",
562                        &sx.to_string(),
563                        &sy.to_string(),
564                    ]),
565                    token,
566                )
567                .await
568            },
569            Backend::Wayland => {
570                if !super::has_command("ydotool") {
571                    anyhow::bail!("ydotool required for Wayland mouse control")
572                }
573                run_cmd_cancellable(
574                    Command::new("ydotool").args([
575                        "mousemove",
576                        "--absolute",
577                        "-x",
578                        &sx.to_string(),
579                        "-y",
580                        &sy.to_string(),
581                    ]),
582                    token,
583                )
584                .await
585            },
586            _ => anyhow::bail!("mouse_move not supported on this platform"),
587        }
588    }
589
590    /// List visible window titles. X11 only; Wayland has no portable
591    /// enumeration primitive.
592    pub async fn list_windows(&self, _token: &CancellationToken) -> Result<Vec<String>> {
593        if !matches!(self.backend, Backend::X11) {
594            anyhow::bail!(
595                "list_windows requires X11. Wayland has no portable window-enumeration \
596                 primitive. Run mermaid from an X11 session."
597            );
598        }
599        let wids =
600            run_cmd_stdout(Command::new("xdotool").args(["search", "--onlyvisible", "--name", ""]))
601                .await?;
602        let mut windows = Vec::new();
603        for wid in wids.lines() {
604            let wid = wid.trim();
605            if wid.is_empty() {
606                continue;
607            }
608            if let Ok(name) =
609                run_cmd_stdout(Command::new("xdotool").args(["getwindowname", wid])).await
610            {
611                let name = name.trim().to_string();
612                if !name.is_empty() && !windows.contains(&name) {
613                    windows.push(name);
614                }
615            }
616        }
617        Ok(windows)
618    }
619}
620
621// ───── RAII temp-file cleanup ──────────────────────────────────────
622
623struct TempFileGuard(PathBuf);
624
625impl Drop for TempFileGuard {
626    fn drop(&mut self) {
627        let _ = std::fs::remove_file(&self.0);
628    }
629}
630
631// ───── subprocess dispatch ─────────────────────────────────────────
632
633async fn dispatch_capture(
634    backend: Backend,
635    spec: &ScreenshotSpec,
636    out_path: &str,
637    token: &CancellationToken,
638) -> Result<(i32, i32, String)> {
639    // Returns (offset_x, offset_y, kind_label). Each branch `select!`s
640    // on `token.cancelled()` so Ctrl+C during a slow capture aborts
641    // the subprocess cleanly.
642    match (backend, spec) {
643        (Backend::X11, ScreenshotSpec::Fullscreen) => {
644            run_cmd_cancellable(Command::new("scrot").args(["-o", out_path]), token).await?;
645            Ok((0, 0, "fullscreen".to_string()))
646        },
647        (Backend::Wayland, ScreenshotSpec::Fullscreen) => {
648            run_cmd_cancellable(Command::new("grim").arg(out_path), token).await?;
649            Ok((0, 0, "fullscreen".to_string()))
650        },
651        (Backend::MacOS, ScreenshotSpec::Fullscreen) => {
652            run_cmd_cancellable(Command::new("screencapture").args(["-x", out_path]), token)
653                .await?;
654            Ok((0, 0, "fullscreen".to_string()))
655        },
656        (Backend::X11, ScreenshotSpec::Focused) => {
657            let (wx, wy) = get_focused_window_geometry_x11()
658                .await
659                .map(|(x, y, _, _)| (x, y))
660                .unwrap_or((0, 0));
661            run_cmd_cancellable(Command::new("scrot").args(["-u", "-o", out_path]), token).await?;
662            Ok((wx, wy, "focused window".to_string()))
663        },
664        (Backend::Wayland, ScreenshotSpec::Focused) => anyhow::bail!(
665            "Mode 'focused' not supported on Wayland (grim has no focused-window \
666             primitive). Use mode: 'fullscreen' or mode: 'monitor' with a specific \
667             output name."
668        ),
669        (Backend::MacOS, ScreenshotSpec::Focused) => {
670            // #100: `screencapture -W` captures only the focused window but
671            // reports no origin, so a click computed from the frame mis-targets
672            // whenever the window isn't at screen (0,0). Capture the full main
673            // display instead — the (0,0) offset is then genuinely correct and
674            // coords translate exactly. We deliberately avoid an AppleScript
675            // window-position probe: it returns points, but offsets here are in
676            // device pixels, so it'd be 2x off on a Retina display (a silent bug
677            // we can't catch on the Linux CI).
678            run_cmd_cancellable(Command::new("screencapture").args(["-x", out_path]), token)
679                .await?;
680            Ok((0, 0, "focused window (full display on macOS)".to_string()))
681        },
682        (Backend::X11, ScreenshotSpec::Region(x, y, w, h)) => {
683            run_cmd_cancellable(
684                Command::new("scrot").args([
685                    "-a",
686                    &format!("{},{},{},{}", x, y, w, h),
687                    "-o",
688                    out_path,
689                ]),
690                token,
691            )
692            .await?;
693            Ok((*x, *y, format!("region {}x{}+{}+{}", w, h, x, y)))
694        },
695        (Backend::Wayland, ScreenshotSpec::Region(x, y, w, h)) => {
696            run_cmd_cancellable(
697                Command::new("grim").args(["-g", &format!("{},{} {}x{}", x, y, w, h), out_path]),
698                token,
699            )
700            .await?;
701            Ok((*x, *y, format!("region {}x{}+{}+{}", w, h, x, y)))
702        },
703        (Backend::X11, ScreenshotSpec::Monitor(name)) => {
704            let (mx, my, mw, mh) = parse_monitor_geometry_x11(name).await.ok_or_else(|| {
705                anyhow::anyhow!(
706                    "Monitor '{}' not found. Run `xrandr --query` to list outputs.",
707                    name
708                )
709            })?;
710            run_cmd_cancellable(
711                Command::new("scrot").args([
712                    "-a",
713                    &format!("{},{},{},{}", mx, my, mw, mh),
714                    "-o",
715                    out_path,
716                ]),
717                token,
718            )
719            .await?;
720            Ok((mx, my, format!("monitor {}", name)))
721        },
722        (Backend::Wayland, ScreenshotSpec::Monitor(name)) => {
723            run_cmd_cancellable(Command::new("grim").args(["-o", name, out_path]), token).await?;
724            Ok((0, 0, format!("monitor {}", name)))
725        },
726        (Backend::X11, ScreenshotSpec::Window(title)) => {
727            // Search for window by name, activate it, sync, then
728            // capture the focused window.
729            let wid = run_cmd_stdout(Command::new("xdotool").args(["search", "--name", title]))
730                .await?
731                .lines()
732                .next()
733                .map(str::trim)
734                .filter(|s| !s.is_empty())
735                .map(str::to_string)
736                .ok_or_else(|| {
737                    anyhow::anyhow!(
738                        "No window found matching '{}'. Use list_windows to see available \
739                         windows.",
740                        title
741                    )
742                })?;
743            run_cmd_cancellable(
744                Command::new("xdotool").args(["windowactivate", "--sync", &wid]),
745                token,
746            )
747            .await?;
748            tokio::time::sleep(std::time::Duration::from_millis(
749                crate::constants::WINDOW_FOCUS_DELAY_MS,
750            ))
751            .await;
752            let (wx, wy) = get_window_geometry_x11(&wid)
753                .await
754                .map(|(x, y, _, _)| (x, y))
755                .unwrap_or((0, 0));
756            run_cmd_cancellable(Command::new("scrot").args(["-u", "-o", out_path]), token).await?;
757            Ok((wx, wy, format!("window \"{}\"", title)))
758        },
759        (Backend::Wayland, ScreenshotSpec::Window(_)) => anyhow::bail!(
760            "Mode 'window' not supported on Wayland (grim has no window-by-name capture). \
761             Use mode: 'fullscreen' or mode: 'monitor' with a specific output name."
762        ),
763        (Backend::MacOS, _) => anyhow::bail!(
764            "This screenshot mode is not yet ported to macOS. Use mode: 'fullscreen' for now."
765        ),
766        (Backend::Windows, _) | (Backend::Unsupported, _) => {
767            anyhow::bail!("Unsupported platform for computer-use capture")
768        },
769    }
770}
771
772/// Race an arbitrary in-process future against `token.cancelled()`, biased toward
773/// cancellation so an already-signalled Esc wins before the (possibly slow) future
774/// is polled. Mirrors `run_cmd_cancellable`'s cancel arm but for a plain future —
775/// `capture` uses it so the downscale + trailing read abort promptly on Esc rather
776/// than blocking up to `SCREENSHOT_DOWNSCALE_TIMEOUT_SECS` (F59). On cancel the
777/// losing branch's future is dropped, so its own temp guards (e.g. the scaled-file
778/// guard inside `downscale_if_needed`) run as it unwinds.
779async fn cancellable<F, T>(token: &CancellationToken, fut: F) -> Result<T>
780where
781    F: std::future::Future<Output = Result<T>>,
782{
783    tokio::select! {
784        biased;
785        _ = token.cancelled() => anyhow::bail!("cancelled"),
786        r = fut => r,
787    }
788}
789
790/// Run a `Command` to completion, racing it against cancellation AND a
791/// wall-clock timeout. Relies on `kill_on_drop(true)` reaping the child when the
792/// future is dropped on cancel or timeout.
793pub(crate) async fn run_cmd_cancellable(
794    cmd: &mut Command,
795    token: &CancellationToken,
796) -> Result<()> {
797    run_cmd_cancellable_with_timeout(
798        cmd,
799        token,
800        std::time::Duration::from_secs(crate::constants::COMPUTER_USE_CMD_TIMEOUT_SECS),
801    )
802    .await
803}
804
805/// `run_cmd_cancellable` with an explicit timeout (extracted so the bound is
806/// testable with a tiny duration). Without the timeout a wedged backend (a dead
807/// `ydotoold` socket, a hung X server, a blocking permission dialog) would hang
808/// the tool until the user pressed Esc (#127).
809async fn run_cmd_cancellable_with_timeout(
810    cmd: &mut Command,
811    token: &CancellationToken,
812    timeout: std::time::Duration,
813) -> Result<()> {
814    cmd.kill_on_drop(true);
815    tokio::select! {
816        biased;
817        _ = token.cancelled() => anyhow::bail!("cancelled"),
818        _ = tokio::time::sleep(timeout) => {
819            anyhow::bail!("subprocess timed out after {:?}", timeout)
820        }
821        res = cmd.output() => {
822            let out = res.context("subprocess spawn")?;
823            if !out.status.success() {
824                anyhow::bail!(
825                    "subprocess failed: {}",
826                    String::from_utf8_lossy(&out.stderr).trim()
827                );
828            }
829            Ok(())
830        }
831    }
832}
833
834async fn run_cmd_stdout(cmd: &mut Command) -> Result<String> {
835    run_cmd_stdout_with_timeout(
836        cmd,
837        std::time::Duration::from_secs(crate::constants::COMPUTER_USE_CMD_TIMEOUT_SECS),
838    )
839    .await
840}
841
842/// `run_cmd_stdout` with an explicit cap (extracted so the timeout is testable
843/// with a tiny duration). `kill_on_drop(true)` reaps the child if the timeout
844/// fires and the `output()` future is dropped (#97).
845async fn run_cmd_stdout_with_timeout(
846    cmd: &mut Command,
847    timeout: std::time::Duration,
848) -> Result<String> {
849    cmd.kill_on_drop(true);
850    let out = match tokio::time::timeout(timeout, cmd.output()).await {
851        Ok(res) => res.context("subprocess spawn")?,
852        Err(_) => anyhow::bail!("subprocess timed out after {:?}", timeout),
853    };
854    if !out.status.success() {
855        anyhow::bail!(
856            "subprocess failed: {}",
857            String::from_utf8_lossy(&out.stderr).trim()
858        );
859    }
860    Ok(String::from_utf8_lossy(&out.stdout).to_string())
861}
862
863// ───── geometry helpers (X11 only; Wayland has no equivalent) ──────
864
865async fn get_focused_window_geometry_x11() -> Option<(i32, i32, u32, u32)> {
866    let wid = run_cmd_stdout(Command::new("xdotool").arg("getactivewindow"))
867        .await
868        .ok()?;
869    let wid = wid.trim();
870    if wid.is_empty() {
871        return None;
872    }
873    get_window_geometry_x11(wid).await
874}
875
876async fn get_window_geometry_x11(wid: &str) -> Option<(i32, i32, u32, u32)> {
877    let out = run_cmd_stdout(Command::new("xdotool").args(["getwindowgeometry", "--shell", wid]))
878        .await
879        .ok()?;
880    let mut x = None;
881    let mut y = None;
882    let mut width = None;
883    let mut height = None;
884    for line in out.lines() {
885        if let Some(v) = line.strip_prefix("X=") {
886            x = v.parse().ok();
887        } else if let Some(v) = line.strip_prefix("Y=") {
888            y = v.parse().ok();
889        } else if let Some(v) = line.strip_prefix("WIDTH=") {
890            width = v.parse().ok();
891        } else if let Some(v) = line.strip_prefix("HEIGHT=") {
892            height = v.parse().ok();
893        }
894    }
895    Some((x?, y?, width?, height?))
896}
897
898async fn parse_monitor_geometry_x11(name: &str) -> Option<(i32, i32, u32, u32)> {
899    let out = run_cmd_stdout(Command::new("xrandr").arg("--query"))
900        .await
901        .ok()?;
902    out.lines()
903        .find_map(|line| parse_xrandr_monitor_line(line, name))
904}
905
906/// Parse one `xrandr --query` line, returning `(x, y, width, height)` when it is
907/// the connected output named `name` and carries a `WxH+X+Y` geometry token.
908/// Pure (no subprocess) so the slice-bounds handling is unit-testable.
909fn parse_xrandr_monitor_line(line: &str, name: &str) -> Option<(i32, i32, u32, u32)> {
910    if !line.contains(" connected") {
911        return None;
912    }
913    let parts: Vec<&str> = line.split_whitespace().collect();
914    if parts.first() != Some(&name) {
915        return None;
916    }
917    // F60: `&parts[2..]` panicked ("range start 2 out of range for slice of
918    // length 1") when an xrandr line split into a single token and the
919    // model-supplied `name` equalled it (e.g. name="connected"): the
920    // `parts.first() == Some(&name)` guard passed, then slice start 2 was past the
921    // length-1 slice. `get(2..).unwrap_or(&[])` yields an empty slice for any line
922    // with fewer than 3 tokens, so a short/odd line can never panic.
923    for part in parts.get(2..).unwrap_or(&[]) {
924        if let Some((res, offsets)) = part.split_once('+')
925            && let Some((w, h)) = res.split_once('x')
926        {
927            let width = w.parse::<u32>().ok()?;
928            let height = h.parse::<u32>().ok()?;
929            let mut off = offsets.splitn(2, '+');
930            let x = off.next()?.parse::<i32>().ok()?;
931            let y = off.next()?.parse::<i32>().ok()?;
932            return Some((x, y, width, height));
933        }
934    }
935    None
936}
937
938// ───── PNG inspection (no image crate dep) ─────────────────────────
939
940fn read_png_width(bytes: &[u8]) -> Option<u32> {
941    if bytes.len() > 24 && &bytes[0..8] == b"\x89PNG\r\n\x1a\n" {
942        Some(u32::from_be_bytes([
943            bytes[16], bytes[17], bytes[18], bytes[19],
944        ]))
945    } else {
946        None
947    }
948}
949
950fn read_png_height(bytes: &[u8]) -> Option<u32> {
951    if bytes.len() > 28 && &bytes[0..8] == b"\x89PNG\r\n\x1a\n" {
952        Some(u32::from_be_bytes([
953            bytes[20], bytes[21], bytes[22], bytes[23],
954        ]))
955    } else {
956        None
957    }
958}
959
960/// Downscale the PNG at `path` to at most `max_width` pixels wide,
961/// using ImageMagick `convert` or ffmpeg as a fallback. Returns the
962/// scale factor (original_width / max_width; 1.0 if no scaling was
963/// needed).
964async fn downscale_if_needed(path: &str, max_width: u32) -> Result<f64> {
965    let bytes = tokio::fs::read(path).await?;
966    let original_width = read_png_width(&bytes).unwrap_or(1920);
967    if original_width <= max_width {
968        return Ok(1.0);
969    }
970    let scale_factor = original_width as f64 / max_width as f64;
971    let scaled = format!("{}.scaled.png", path);
972    // F58: the caller's TempFileGuard only tracks the original temp file, not this
973    // sibling. Guard the scaled file here so every exit path removes it — most
974    // importantly a successful encode followed by a failed `rename(&scaled, path)`,
975    // whose `?` early return previously leaked `{path}.scaled.png`. On the success
976    // path the rename has already moved the file, so this guard's remove no-ops.
977    let _scaled_guard = TempFileGuard(PathBuf::from(&scaled));
978    // #97: time-box the encoders + kill_on_drop. The double `Ok(Ok(..))` means a
979    // timeout (outer Err) OR a spawn error (inner Err) falls through to the next
980    // encoder and finally to the full-resolution fallback below — preserving the
981    // existing graceful degradation rather than hanging the agent loop.
982    let downscale_timeout =
983        std::time::Duration::from_secs(crate::constants::SCREENSHOT_DOWNSCALE_TIMEOUT_SECS);
984
985    let convert = tokio::time::timeout(
986        downscale_timeout,
987        Command::new("convert")
988            .args([path, "-resize", &format!("{}x", max_width), &scaled])
989            .kill_on_drop(true)
990            .output(),
991    )
992    .await;
993    if let Ok(Ok(o)) = convert
994        && o.status.success()
995    {
996        tokio::fs::rename(&scaled, path).await?;
997        return Ok(scale_factor);
998    }
999
1000    let ffmpeg = tokio::time::timeout(
1001        downscale_timeout,
1002        Command::new("ffmpeg")
1003            .args([
1004                "-y",
1005                "-i",
1006                path,
1007                "-vf",
1008                &format!("scale={}:-1", max_width),
1009                &scaled,
1010            ])
1011            .kill_on_drop(true)
1012            .output(),
1013    )
1014    .await;
1015    if let Ok(Ok(o)) = ffmpeg
1016        && o.status.success()
1017    {
1018        tokio::fs::rename(&scaled, path).await?;
1019        return Ok(scale_factor);
1020    }
1021
1022    // `_scaled_guard` (declared above with `scaled`) removes any partial
1023    // `.scaled.png` on drop — covering this no-encoder fallback as well as a failed
1024    // `rename(&scaled, path)` early return (F58).
1025    tracing::warn!(
1026        original_width,
1027        "neither ImageMagick nor ffmpeg available; sending full-resolution screenshot"
1028    );
1029    Ok(1.0)
1030}
1031
1032#[cfg(test)]
1033mod tests {
1034    use super::*;
1035
1036    #[test]
1037    fn registry_lru_evicts_oldest_past_capacity() {
1038        let mut r = ScreenshotRegistry::new();
1039        for i in 0..(SCREENSHOT_REGISTRY_CAPACITY as u64 + 3) {
1040            r.push(ScreenshotMetadata {
1041                id: i,
1042                scale_factor: 1.0,
1043                offset_x: 0,
1044                offset_y: 0,
1045                width: 0,
1046                height: 0,
1047                kind: "fullscreen".to_string(),
1048            });
1049        }
1050        assert_eq!(r.len(), SCREENSHOT_REGISTRY_CAPACITY);
1051        // First 3 should have been evicted.
1052        assert!(r.get(0).is_none());
1053        assert!(r.get(1).is_none());
1054        assert!(r.get(2).is_none());
1055        // Latest remains.
1056        assert_eq!(
1057            r.latest().unwrap().id,
1058            SCREENSHOT_REGISTRY_CAPACITY as u64 + 2
1059        );
1060    }
1061
1062    #[test]
1063    fn scale_coords_applies_scale_and_offset() {
1064        let d = ComputerUseDriver::new(Backend::X11);
1065        let id = d.register_screenshot(2.0, 100, 50, 640, 480, "fullscreen".to_string());
1066        let (sx, sy) = d.scale_coords(10, 20, Some(id)).unwrap();
1067        assert_eq!(sx, 100 + 20);
1068        assert_eq!(sy, 50 + 40);
1069    }
1070
1071    #[test]
1072    fn scale_coords_saturates_on_offset_overflow() {
1073        // #32: a huge model x + positive offset must not panic (debug) or wrap to
1074        // a negative coord (release). width/height = 0 disables the #96 upper
1075        // clamp so the value actually reaches the offset add.
1076        let d = ComputerUseDriver::new(Backend::X11);
1077        let id = d.register_screenshot(1.0, 100, 100, 0, 0, "fullscreen".to_string());
1078        let (sx, sy) = d.scale_coords(i32::MAX, i32::MAX, Some(id)).unwrap();
1079        assert_eq!(sx, i32::MAX);
1080        assert_eq!(sy, i32::MAX);
1081    }
1082
1083    #[test]
1084    fn scale_coords_clamps_negative_into_region() {
1085        // #96: negative model coords clamp to the region's top-left origin.
1086        let d = ComputerUseDriver::new(Backend::X11);
1087        let id = d.register_screenshot(2.0, 100, 50, 640, 480, "region".to_string());
1088        assert_eq!(d.scale_coords(-9999, -1, Some(id)).unwrap(), (100, 50));
1089    }
1090
1091    #[test]
1092    fn scale_coords_clamps_over_max_into_region() {
1093        // #96: coords past the frame clamp to the last in-frame pixel and stay
1094        // inside [offset, offset + region_dim).
1095        let d = ComputerUseDriver::new(Backend::X11);
1096        let id = d.register_screenshot(2.0, 100, 50, 640, 480, "region".to_string());
1097        let (sx, sy) = d.scale_coords(100_000, 100_000, Some(id)).unwrap();
1098        // x: clamp 100000 -> 639; 639*2 + 100 = 1378. y: 479*2 + 50 = 1008.
1099        assert_eq!((sx, sy), (1378, 1008));
1100        // region real size = 1280x960; in-frame bounds [100,1380) x [50,1010).
1101        assert!(sx < 100 + 1280 && sy < 50 + 960);
1102    }
1103
1104    #[test]
1105    fn scale_coords_errors_on_evicted_id() {
1106        let d = ComputerUseDriver::new(Backend::X11);
1107        for _ in 0..(SCREENSHOT_REGISTRY_CAPACITY + 1) {
1108            d.register_screenshot(1.0, 0, 0, 0, 0, "fullscreen".to_string());
1109        }
1110        // id 0 is evicted now.
1111        let err = d.scale_coords(0, 0, Some(0)).unwrap_err();
1112        assert!(
1113            err.contains("evicted"),
1114            "expected eviction message, got: {}",
1115            err
1116        );
1117    }
1118
1119    #[test]
1120    fn scale_coords_errors_with_no_screenshots_yet() {
1121        let d = ComputerUseDriver::new(Backend::X11);
1122        let err = d.scale_coords(10, 20, None).unwrap_err();
1123        assert!(err.contains("No screenshots"));
1124    }
1125
1126    #[test]
1127    fn ensure_alive_fails_on_unsupported_backend() {
1128        let d = ComputerUseDriver::new(Backend::Unsupported);
1129        assert!(d.ensure_alive().is_err());
1130    }
1131
1132    #[cfg(unix)]
1133    #[tokio::test]
1134    async fn run_cmd_stdout_times_out_on_slow_command() {
1135        // #97: a wedged probe must not hang — the seam fires the timeout, drops
1136        // the future (kill_on_drop reaps the child), and bails.
1137        let mut cmd = Command::new("sleep");
1138        cmd.arg("5");
1139        let err = run_cmd_stdout_with_timeout(&mut cmd, std::time::Duration::from_millis(50))
1140            .await
1141            .unwrap_err();
1142        assert!(err.to_string().contains("timed out"), "got: {err}");
1143    }
1144
1145    #[cfg(unix)]
1146    #[tokio::test]
1147    async fn run_cmd_stdout_returns_output_for_fast_command() {
1148        let mut cmd = Command::new("echo");
1149        cmd.arg("hi");
1150        assert_eq!(run_cmd_stdout(&mut cmd).await.unwrap().trim(), "hi");
1151    }
1152
1153    #[cfg(unix)]
1154    #[tokio::test]
1155    async fn run_cmd_cancellable_times_out_on_wedged_backend() {
1156        // #127: a wedged capture/input backend must not hang the tool until Esc;
1157        // the timeout fires, drops the future (kill_on_drop reaps it), and bails.
1158        let token = tokio_util::sync::CancellationToken::new();
1159        let mut cmd = Command::new("sleep");
1160        cmd.arg("5");
1161        let err = run_cmd_cancellable_with_timeout(
1162            &mut cmd,
1163            &token,
1164            std::time::Duration::from_millis(50),
1165        )
1166        .await
1167        .unwrap_err();
1168        assert!(err.to_string().contains("timed out"), "got: {err}");
1169    }
1170
1171    // ── F60: xrandr line parsing must never panic on a short/odd line ──
1172
1173    #[test]
1174    fn parse_xrandr_monitor_line_short_line_does_not_panic() {
1175        // F60: a line that splits into the single token "connected", with the
1176        // model-supplied monitor name also "connected", passed the
1177        // `parts.first() == Some(&name)` guard and then panicked on `&parts[2..]`
1178        // ("range start 2 out of range for slice of length 1"). Bounds-checked
1179        // slicing must yield None instead of panicking.
1180        assert_eq!(parse_xrandr_monitor_line(" connected", "connected"), None);
1181        // Two tokens (still < 3), name matches: no geometry token, no panic.
1182        assert_eq!(
1183            parse_xrandr_monitor_line("HDMI-1 connected", "HDMI-1"),
1184            None
1185        );
1186    }
1187
1188    #[test]
1189    fn parse_xrandr_monitor_line_parses_geometry_and_skips_others() {
1190        let line = "HDMI-1 connected primary 2560x1440+1920+0 \
1191                    (normal left inverted right) 597mm x 336mm";
1192        assert_eq!(
1193            parse_xrandr_monitor_line(line, "HDMI-1"),
1194            Some((1920, 0, 2560, 1440))
1195        );
1196        // Non-matching name -> None.
1197        assert_eq!(parse_xrandr_monitor_line(line, "DP-2"), None);
1198        // A "disconnected" line is ignored even though it contains "connected" and
1199        // its first token matches the requested name.
1200        assert_eq!(
1201            parse_xrandr_monitor_line("DP-3 disconnected (normal left inverted right)", "DP-3"),
1202            None
1203        );
1204    }
1205
1206    // ── F58: the scaled temp sibling must never leak ──
1207
1208    #[test]
1209    fn temp_file_guard_removes_scaled_sibling_on_drop() {
1210        // F58: a successful encode followed by a failed `rename(&scaled, path)` used
1211        // to leak `{path}.scaled.png` because capture's TempFileGuard only tracks
1212        // the original temp file. `downscale_if_needed` now wraps the scaled sibling
1213        // in its own TempFileGuard; this asserts that guard removes the file on drop
1214        // — the mechanism that closes the rename-failure (and fallback) leak paths.
1215        let scaled = std::env::temp_dir().join(format!(
1216            "mermaid-f58-guard-{}.png.scaled.png",
1217            std::process::id()
1218        ));
1219        std::fs::write(&scaled, b"x").unwrap();
1220        assert!(scaled.exists());
1221        {
1222            let _guard = TempFileGuard(scaled.clone());
1223        }
1224        assert!(
1225            !scaled.exists(),
1226            "scaled sibling must be removed when its guard drops"
1227        );
1228    }
1229
1230    #[tokio::test]
1231    async fn downscale_skips_and_leaves_no_scaled_sibling_when_within_max() {
1232        // A capture already within max_width early-returns scale 1.0 and must create
1233        // no `.scaled.png` sibling. Encoder-independent (no convert/ffmpeg needed).
1234        let path =
1235            std::env::temp_dir().join(format!("mermaid-f58-skip-{}.png", std::process::id()));
1236        let path_str = path.to_string_lossy().to_string();
1237        let _cleanup = TempFileGuard(path.clone());
1238        // Minimal PNG header advertising width=16 in the IHDR (bytes 16..20).
1239        let mut png = vec![0x89, b'P', b'N', b'G', b'\r', b'\n', 0x1a, b'\n'];
1240        png.extend_from_slice(&[0, 0, 0, 13]); // IHDR chunk length
1241        png.extend_from_slice(b"IHDR");
1242        png.extend_from_slice(&16u32.to_be_bytes()); // width
1243        png.extend_from_slice(&16u32.to_be_bytes()); // height
1244        png.extend_from_slice(&[8, 6, 0, 0, 0]); // bit depth/colour/...
1245        std::fs::write(&path, &png).unwrap();
1246
1247        let scale = downscale_if_needed(&path_str, 1920).await.unwrap();
1248        assert_eq!(scale, 1.0);
1249        assert!(
1250            !std::path::Path::new(&format!("{}.scaled.png", path_str)).exists(),
1251            "no scaled sibling for an already-small capture"
1252        );
1253    }
1254
1255    // ── F59: the downscale/read race must abort promptly on cancel ──
1256
1257    #[tokio::test]
1258    async fn cancellable_returns_cancelled_when_token_already_cancelled() {
1259        // F59: with the token already cancelled, a slow future (stand-in for the
1260        // ImageMagick/ffmpeg downscale) must abort at once rather than run to
1261        // completion or its timeout.
1262        let token = tokio_util::sync::CancellationToken::new();
1263        token.cancel();
1264        let slow = async {
1265            tokio::time::sleep(std::time::Duration::from_secs(30)).await;
1266            Ok::<(), anyhow::Error>(())
1267        };
1268        let err = cancellable(&token, slow).await.unwrap_err();
1269        assert!(err.to_string().contains("cancelled"), "got: {err}");
1270    }
1271
1272    #[tokio::test]
1273    async fn cancellable_passes_through_result_when_not_cancelled() {
1274        let token = tokio_util::sync::CancellationToken::new();
1275        let v = cancellable(&token, async { Ok::<u32, anyhow::Error>(7) })
1276            .await
1277            .unwrap();
1278        assert_eq!(v, 7);
1279    }
1280
1281    #[tokio::test]
1282    async fn cancellable_aborts_inflight_future_on_cancel() {
1283        // Cancel arrives while the future is parked: the biased select wakes on the
1284        // token and returns "cancelled" without waiting out the slow future.
1285        let token = tokio_util::sync::CancellationToken::new();
1286        let t2 = token.clone();
1287        tokio::spawn(async move {
1288            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
1289            t2.cancel();
1290        });
1291        let slow = async {
1292            tokio::time::sleep(std::time::Duration::from_secs(30)).await;
1293            Ok::<(), anyhow::Error>(())
1294        };
1295        let start = std::time::Instant::now();
1296        let err = cancellable(&token, slow).await.unwrap_err();
1297        assert!(err.to_string().contains("cancelled"), "got: {err}");
1298        assert!(
1299            start.elapsed() < std::time::Duration::from_secs(5),
1300            "must abort promptly, not wait out the slow future"
1301        );
1302    }
1303}