1use 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#[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 pub width: u32,
49 pub height: u32,
50 pub kind: String,
54}
55
56#[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#[derive(Debug, Clone)]
100pub enum ScreenshotSpec {
101 Fullscreen,
102 Focused,
103 Monitor(String),
104 Region(i32, i32, u32, u32),
106 Window(String),
107}
108
109#[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
124pub struct ComputerUseDriver {
126 backend: Backend,
127 registry: Mutex<ScreenshotRegistry>,
128 file_counter: AtomicU64,
132 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 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 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 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 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 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 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 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 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 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 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 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 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
383const CURSOR_LANDED_TOLERANCE_PX: i32 = 5;
386
387impl ComputerUseDriver {
390 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 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 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 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 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 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 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 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
621struct TempFileGuard(PathBuf);
624
625impl Drop for TempFileGuard {
626 fn drop(&mut self) {
627 let _ = std::fs::remove_file(&self.0);
628 }
629}
630
631async fn dispatch_capture(
634 backend: Backend,
635 spec: &ScreenshotSpec,
636 out_path: &str,
637 token: &CancellationToken,
638) -> Result<(i32, i32, String)> {
639 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 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 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
772async 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
790pub(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
805async 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
842async 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
863async 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
906fn 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 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
938fn 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
960async 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 let _scaled_guard = TempFileGuard(PathBuf::from(&scaled));
978 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 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 assert!(r.get(0).is_none());
1053 assert!(r.get(1).is_none());
1054 assert!(r.get(2).is_none());
1055 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 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 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 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 assert_eq!((sx, sy), (1378, 1008));
1100 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 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 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 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 #[test]
1174 fn parse_xrandr_monitor_line_short_line_does_not_panic() {
1175 assert_eq!(parse_xrandr_monitor_line(" connected", "connected"), None);
1181 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 assert_eq!(parse_xrandr_monitor_line(line, "DP-2"), None);
1198 assert_eq!(
1201 parse_xrandr_monitor_line("DP-3 disconnected (normal left inverted right)", "DP-3"),
1202 None
1203 );
1204 }
1205
1206 #[test]
1209 fn temp_file_guard_removes_scaled_sibling_on_drop() {
1210 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 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 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]); png.extend_from_slice(b"IHDR");
1242 png.extend_from_slice(&16u32.to_be_bytes()); png.extend_from_slice(&16u32.to_be_bytes()); png.extend_from_slice(&[8, 6, 0, 0, 0]); 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 #[tokio::test]
1258 async fn cancellable_returns_cancelled_when_token_already_cancelled() {
1259 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 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}