feat: Support for screenshot annotate on rust (#706)

Co-authored-by: hidezhao <hidezhao@tencent.com>

cli/src/native/actions.rs

@@ -1295,8 +1295,20 @@ async fn handle_snapshot(cmd: &Value, state: &mut DaemonState) -> Result<Value,
 }
 
 async fn handle_screenshot(cmd: &Value, state: &mut DaemonState) -> Result<Value, String> {
+    let annotate = cmd
+        .get("annotate")
+        .and_then(|v| v.as_bool())
+        .unwrap_or(false);
+
     if let Some(ref wb) = state.webdriver_backend {
         if state.browser.is_none() {
+            if annotate {
+                return Err(
+                    "Annotated screenshots are not yet implemented on the WebDriver backend"
+                        .to_string(),
+                );
+            }
+
             let base64_data = wb.screenshot().await?;
             let path = cmd.get("path").and_then(|v| v.as_str());
             if let Some(p) = path {
@@ -1349,12 +1361,33 @@ async fn handle_screenshot(cmd: &Value, state: &mut DaemonState) -> Result<Value
             .get("quality")
             .and_then(|v| v.as_i64())
             .map(|q| q as i32),
+        annotate,
     };
 
-    let (path, _base64) =
+    if annotate {
+        state.ref_map.clear();
+        let _ = snapshot::take_snapshot(
+            &mgr.client,
+            &session_id,
+            &SnapshotOptions {
+                interactive: true,
+                ..SnapshotOptions::default()
+            },
+            &mut state.ref_map,
+        )
+        .await?;
+    }
+
+    let result =
         screenshot::take_screenshot(&mgr.client, &session_id, &state.ref_map, &options).await?;
 
-    Ok(json!({ "path": path }))
+    let mut response = json!({ "path": result.path });
+    if !result.annotations.is_empty() {
+        response["annotations"] = serde_json::to_value(&result.annotations)
+            .map_err(|e| format!("Failed to serialize annotations: {}", e))?;
+    }
+
+    Ok(response)
 }
 
 async fn handle_click(cmd: &Value, state: &mut DaemonState) -> Result<Value, String> {
@@ -4103,13 +4136,14 @@ async fn handle_diff_screenshot(cmd: &Value, state: &DaemonState) -> Result<Valu
             .unwrap_or(false),
         format: "png".to_string(),
         quality: None,
+        annotate: false,
     };
 
-    let (_path, base64_data) =
+    let result =
         screenshot::take_screenshot(&mgr.client, &session_id, &state.ref_map, &options).await?;
 
     let current_bytes =
-        base64::Engine::decode(&base64::engine::general_purpose::STANDARD, &base64_data)
+        base64::Engine::decode(&base64::engine::general_purpose::STANDARD, &result.base64)
             .map_err(|e| format!("Failed to decode screenshot: {}", e))?;
 
     let baseline_bytes =

cli/src/native/browser.rs

@@ -1,7 +1,7 @@
 use serde_json::{json, Value};
 use std::collections::HashSet;
 use std::sync::Arc;
-use tokio::sync::Mutex;
+use tokio::sync::{broadcast, Mutex};
 
 use super::cdp::chrome::{
     auto_connect_cdp, discover_cdp_url, launch_chrome, ChromeProcess, LaunchOptions,
@@ -412,6 +412,7 @@ impl BrowserManager {
 
     pub async fn navigate(&mut self, url: &str, wait_until: WaitUntil) -> Result<Value, String> {
         let session_id = self.active_session_id()?.to_string();
+        let mut lifecycle_rx = self.client.subscribe();
 
         let nav_result: PageNavigateResult = self
             .client
@@ -429,7 +430,8 @@ impl BrowserManager {
             return Err(format!("Navigation failed: {}", error_text));
         }
 
-        self.wait_for_lifecycle(wait_until, &session_id).await?;
+        self.wait_for_lifecycle(wait_until, &session_id, &mut lifecycle_rx)
+            .await?;
 
         let page_url = self.get_url().await.unwrap_or_else(|_| url.to_string());
         let title = self.get_title().await.unwrap_or_default();
@@ -446,14 +448,14 @@ impl BrowserManager {
         &self,
         wait_until: WaitUntil,
         session_id: &str,
+        rx: &mut broadcast::Receiver<CdpEvent>,
     ) -> Result<(), String> {
         let event_name = match wait_until {
             WaitUntil::Load => "Page.loadEventFired",
             WaitUntil::DomContentLoaded => "Page.domContentEventFired",
-            WaitUntil::NetworkIdle => return self.wait_for_network_idle(session_id).await,
+            WaitUntil::NetworkIdle => return self.wait_for_network_idle(session_id, rx).await,
         };
 
-        let mut rx = self.client.subscribe();
         let timeout = tokio::time::Duration::from_millis(self.default_timeout_ms);
 
         tokio::time::timeout(timeout, async {
@@ -468,8 +470,11 @@ impl BrowserManager {
         .map_err(|_| format!("Timeout waiting for {}", event_name))?
     }
 
-    async fn wait_for_network_idle(&self, session_id: &str) -> Result<(), String> {
-        let mut rx = self.client.subscribe();
+    async fn wait_for_network_idle(
+        &self,
+        session_id: &str,
+        rx: &mut broadcast::Receiver<CdpEvent>,
+    ) -> Result<(), String> {
         let pending = Arc::new(Mutex::new(HashSet::<String>::new()));
         let timeout = tokio::time::Duration::from_millis(self.default_timeout_ms);
 
@@ -588,7 +593,9 @@ impl BrowserManager {
         wait_until: WaitUntil,
         session_id: &str,
     ) -> Result<(), String> {
-        self.wait_for_lifecycle(wait_until, session_id).await
+        let mut rx = self.client.subscribe();
+        self.wait_for_lifecycle(wait_until, session_id, &mut rx)
+            .await
     }
 
     pub async fn close(&mut self) -> Result<(), String> {

cli/src/native/e2e_tests.rs

@@ -197,6 +197,74 @@ async fn e2e_screenshot() {
     assert!(std::path::Path::new(&tmp_path).exists());
     let _ = std::fs::remove_file(&tmp_path);
 
+    let resp = execute_command(
+        &json!({
+            "id": "5",
+            "action": "setcontent",
+            "html": r##"
+                <html><body>
+                  <button onclick="document.getElementById('result').textContent = 'clicked'">Submit</button>
+                  <a href="#">Home</a>
+                  <div id="result"></div>
+                </body></html>
+            "##,
+        }),
+        &mut state,
+    )
+    .await;
+    assert_success(&resp);
+
+    let resp = execute_command(
+        &json!({ "id": "6", "action": "screenshot", "annotate": true }),
+        &mut state,
+    )
+    .await;
+    assert_success(&resp);
+    let annotations = get_data(&resp)["annotations"]
+        .as_array()
+        .expect("Annotated screenshot should return annotations");
+    assert!(
+        !annotations.is_empty(),
+        "Annotated screenshot should have at least one annotation"
+    );
+
+    let submit_ref = annotations
+        .iter()
+        .find(|ann| ann.get("name").and_then(|v| v.as_str()) == Some("Submit"))
+        .and_then(|ann| ann.get("ref").and_then(|v| v.as_str()))
+        .expect("Expected a Submit annotation");
+
+    let resp = execute_command(
+        &json!({
+            "id": "7",
+            "action": "evaluate",
+            "script": "document.getElementById('__agent_browser_annotations__') === null"
+        }),
+        &mut state,
+    )
+    .await;
+    assert_success(&resp);
+    assert_eq!(get_data(&resp)["result"], true);
+
+    let resp = execute_command(
+        &json!({ "id": "8", "action": "click", "selector": format!("@{}", submit_ref) }),
+        &mut state,
+    )
+    .await;
+    assert_success(&resp);
+
+    let resp = execute_command(
+        &json!({
+            "id": "9",
+            "action": "evaluate",
+            "script": "document.getElementById('result').textContent"
+        }),
+        &mut state,
+    )
+    .await;
+    assert_success(&resp);
+    assert_eq!(get_data(&resp)["result"], "clicked");
+
     let resp = execute_command(&json!({ "id": "99", "action": "close" }), &mut state).await;
     assert_success(&resp);
 }

cli/src/native/element.rs

@@ -47,10 +47,47 @@ impl RefMap {
         );
     }
 
+    pub fn add_selector(
+        &mut self,
+        ref_id: String,
+        selector: String,
+        role: &str,
+        name: &str,
+        nth: Option<usize>,
+    ) {
+        self.map.insert(
+            ref_id,
+            RefEntry {
+                backend_node_id: None,
+                role: role.to_string(),
+                name: name.to_string(),
+                nth,
+                selector: Some(selector),
+            },
+        );
+    }
+
     pub fn get(&self, ref_id: &str) -> Option<&RefEntry> {
         self.map.get(ref_id)
     }
 
+    pub fn entries_sorted(&self) -> Vec<(String, RefEntry)> {
+        let mut entries = self
+            .map
+            .iter()
+            .map(|(ref_id, entry)| (ref_id.clone(), entry.clone()))
+            .collect::<Vec<_>>();
+
+        entries.sort_by_key(|(ref_id, _)| {
+            ref_id
+                .strip_prefix('e')
+                .and_then(|n| n.parse::<usize>().ok())
+                .unwrap_or(usize::MAX)
+        });
+
+        entries
+    }
+
     pub fn clear(&mut self) {
         self.map.clear();
         self.next_ref = 1;

cli/src/native/screenshot.rs

@@ -1,3 +1,4 @@
+use serde::Serialize;
 use serde_json::Value;
 use std::path::PathBuf;
 
@@ -5,12 +6,57 @@ use super::cdp::client::CdpClient;
 use super::cdp::types::*;
 use super::element::RefMap;
 
+const ANNOTATION_OVERLAY_ID: &str = "__agent_browser_annotations__";
+
+#[derive(Debug, Clone)]
+struct Rect {
+    x: f64,
+    y: f64,
+    width: f64,
+    height: f64,
+}
+
+#[derive(Debug, Clone)]
+struct RawAnnotation {
+    ref_id: String,
+    number: u64,
+    role: String,
+    name: Option<String>,
+    rect: Rect,
+}
+
+#[derive(Debug, Clone, Serialize)]
+pub struct AnnotationBox {
+    pub x: i64,
+    pub y: i64,
+    pub width: i64,
+    pub height: i64,
+}
+
+#[derive(Debug, Clone)]
+pub struct ScreenshotAnnotation {
+    pub ref_id: String,
+    pub number: u64,
+    pub role: String,
+    pub name: Option<String>,
+    pub box_: AnnotationBox,
+}
+
+#[derive(Debug, Clone)]
+pub struct ScreenshotResult {
+    pub path: String,
+    pub base64: String,
+    pub annotations: Vec<ScreenshotAnnotation>,
+}
+
+#[derive(Debug, Clone)]
 pub struct ScreenshotOptions {
     pub selector: Option<String>,
     pub path: Option<String>,
     pub full_page: bool,
     pub format: String,
     pub quality: Option<i32>,
+    pub annotate: bool,
 }
 
 impl Default for ScreenshotOptions {
@@ -21,16 +67,99 @@ impl Default for ScreenshotOptions {
             full_page: false,
             format: "png".to_string(),
             quality: None,
+            annotate: false,
         }
     }
 }
 
+impl Serialize for ScreenshotAnnotation {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        use serde::ser::SerializeStruct;
+
+        let mut state = serializer.serialize_struct("ScreenshotAnnotation", 5)?;
+        state.serialize_field("ref", &self.ref_id)?;
+        state.serialize_field("number", &self.number)?;
+        state.serialize_field("role", &self.role)?;
+        if let Some(name) = &self.name {
+            state.serialize_field("name", name)?;
+        }
+        state.serialize_field("box", &self.box_)?;
+        state.end()
+    }
+}
+
+/// Captures a screenshot via CDP and optionally overlays numbered annotations
+/// that mirror the Node.js screenshot `annotate` mode.
 pub async fn take_screenshot(
     client: &CdpClient,
     session_id: &str,
     ref_map: &RefMap,
     options: &ScreenshotOptions,
-) -> Result<(String, String), String> {
+) -> Result<ScreenshotResult, String> {
+    let target_rect = if options.annotate {
+        match options.selector.as_deref() {
+            Some(selector) => get_rect_for_selector(client, session_id, ref_map, selector).await?,
+            None => None,
+        }
+    } else {
+        None
+    };
+
+    let raw_annotations = if options.annotate {
+        collect_annotations(client, session_id, ref_map).await?
+    } else {
+        Vec::new()
+    };
+
+    let overlay_items = filter_annotations(raw_annotations, target_rect.as_ref());
+    let overlay_injected = if options.annotate && !overlay_items.is_empty() {
+        inject_annotation_overlay(client, session_id, &overlay_items).await?;
+        true
+    } else {
+        false
+    };
+
+    let base64 = capture_screenshot_base64(client, session_id, ref_map, options).await;
+
+    if overlay_injected {
+        let _ = remove_annotation_overlay(client, session_id).await;
+    }
+
+    let base64 = base64?;
+    let annotations = if options.annotate {
+        let scroll = if options.full_page {
+            Some(get_scroll_offsets(client, session_id).await?)
+        } else {
+            None
+        };
+        project_annotations(&overlay_items, target_rect.as_ref(), scroll)
+    } else {
+        Vec::new()
+    };
+
+    let ext = if options.format == "jpeg" {
+        "jpg"
+    } else {
+        "png"
+    };
+    let path = save_screenshot(&base64, options.path.as_deref(), ext)?;
+
+    Ok(ScreenshotResult {
+        path,
+        base64,
+        annotations,
+    })
+}
+
+async fn capture_screenshot_base64(
+    client: &CdpClient,
+    session_id: &str,
+    ref_map: &RefMap,
+    options: &ScreenshotOptions,
+) -> Result<String, String> {
     let mut params = CaptureScreenshotParams {
         format: Some(options.format.clone()),
         quality: if options.format == "jpeg" {
@@ -64,40 +193,12 @@ pub async fn take_screenshot(
             });
         }
     } else if let Some(ref selector) = options.selector {
-        // Element screenshot via bounding box
-        let object_id =
-            super::element::resolve_element_object_id(client, session_id, ref_map, selector)
-                .await?;
-
-        let result: EvaluateResult = client
-            .send_command_typed(
-                "Runtime.callFunctionOn",
-                &CallFunctionOnParams {
-                    function_declaration: r#"function() {
-                        const rect = this.getBoundingClientRect();
-                        return { x: rect.x, y: rect.y, width: rect.width, height: rect.height };
-                    }"#
-                    .to_string(),
-                    object_id: Some(object_id),
-                    arguments: None,
-                    return_by_value: Some(true),
-                    await_promise: Some(false),
-                },
-                Some(session_id),
-            )
-            .await?;
-
-        if let Some(rect) = result.result.value {
-            let x = rect.get("x").and_then(|v| v.as_f64()).unwrap_or(0.0);
-            let y = rect.get("y").and_then(|v| v.as_f64()).unwrap_or(0.0);
-            let w = rect.get("width").and_then(|v| v.as_f64()).unwrap_or(100.0);
-            let h = rect.get("height").and_then(|v| v.as_f64()).unwrap_or(100.0);
-
+        if let Some(rect) = get_rect_for_selector(client, session_id, ref_map, selector).await? {
             params.clip = Some(Viewport {
-                x,
-                y,
-                width: w,
-                height: h,
+                x: rect.x,
+                y: rect.y,
+                width: rect.width,
+                height: rect.height,
                 scale: 1.0,
             });
         }
@@ -107,14 +208,280 @@ pub async fn take_screenshot(
         .send_command_typed("Page.captureScreenshot", &params, Some(session_id))
         .await?;
 
-    let ext = if options.format == "jpeg" {
-        "jpg"
-    } else {
-        "png"
-    };
+    Ok(result.data)
+}
 
-    let save_path = match &options.path {
-        Some(p) => p.clone(),
+async fn collect_annotations(
+    client: &CdpClient,
+    session_id: &str,
+    ref_map: &RefMap,
+) -> Result<Vec<RawAnnotation>, String> {
+    let mut annotations = Vec::new();
+
+    for (ref_id, entry) in ref_map.entries_sorted() {
+        let object_id =
+            match super::element::resolve_element_object_id(client, session_id, ref_map, &ref_id)
+                .await
+            {
+                Ok(id) => id,
+                Err(_) => continue,
+            };
+
+        let Some(rect) = get_rect_for_object(client, session_id, &object_id).await? else {
+            continue;
+        };
+
+        if rect.width <= 0.0 || rect.height <= 0.0 {
+            continue;
+        }
+
+        let number = ref_id
+            .strip_prefix('e')
+            .and_then(|n| n.parse::<u64>().ok())
+            .unwrap_or(0);
+
+        annotations.push(RawAnnotation {
+            ref_id,
+            number,
+            role: entry.role,
+            name: (!entry.name.is_empty()).then_some(entry.name),
+            rect,
+        });
+    }
+
+    Ok(annotations)
+}
+
+async fn get_rect_for_selector(
+    client: &CdpClient,
+    session_id: &str,
+    ref_map: &RefMap,
+    selector: &str,
+) -> Result<Option<Rect>, String> {
+    let object_id =
+        super::element::resolve_element_object_id(client, session_id, ref_map, selector).await?;
+    get_rect_for_object(client, session_id, &object_id).await
+}
+
+async fn get_rect_for_object(
+    client: &CdpClient,
+    session_id: &str,
+    object_id: &str,
+) -> Result<Option<Rect>, String> {
+    let result: EvaluateResult = client
+        .send_command_typed(
+            "Runtime.callFunctionOn",
+            &CallFunctionOnParams {
+                function_declaration: r#"function() {
+                    const rect = this.getBoundingClientRect();
+                    return { x: rect.x, y: rect.y, width: rect.width, height: rect.height };
+                }"#
+                .to_string(),
+                object_id: Some(object_id.to_string()),
+                arguments: None,
+                return_by_value: Some(true),
+                await_promise: Some(false),
+            },
+            Some(session_id),
+        )
+        .await?;
+
+    Ok(result.result.value.as_ref().and_then(parse_rect))
+}
+
+fn parse_rect(value: &Value) -> Option<Rect> {
+    Some(Rect {
+        x: value.get("x")?.as_f64()?,
+        y: value.get("y")?.as_f64()?,
+        width: value.get("width")?.as_f64()?,
+        height: value.get("height")?.as_f64()?,
+    })
+}
+
+fn filter_annotations(
+    annotations: Vec<RawAnnotation>,
+    target_rect: Option<&Rect>,
+) -> Vec<RawAnnotation> {
+    let mut items = annotations
+        .into_iter()
+        .filter(|annotation| match target_rect {
+            Some(target) => overlaps(&annotation.rect, target),
+            None => true,
+        })
+        .collect::<Vec<_>>();
+
+    items.sort_by_key(|annotation| annotation.number);
+    items
+}
+
+fn overlaps(left: &Rect, right: &Rect) -> bool {
+    let left_x2 = left.x + left.width;
+    let left_y2 = left.y + left.height;
+    let right_x2 = right.x + right.width;
+    let right_y2 = right.y + right.height;
+
+    left.x < right_x2 && left_x2 > right.x && left.y < right_y2 && left_y2 > right.y
+}
+
+async fn inject_annotation_overlay(
+    client: &CdpClient,
+    session_id: &str,
+    annotations: &[RawAnnotation],
+) -> Result<(), String> {
+    let overlay_data = annotations
+        .iter()
+        .map(|annotation| {
+            serde_json::json!({
+                "number": annotation.number,
+                "x": round(annotation.rect.x),
+                "y": round(annotation.rect.y),
+                "width": round(annotation.rect.width),
+                "height": round(annotation.rect.height),
+            })
+        })
+        .collect::<Vec<_>>();
+
+    let expression = format!(
+        r#"(() => {{
+            var items = {items};
+            var id = {overlay_id};
+            var existing = document.getElementById(id);
+            if (existing) existing.remove();
+            var sx = window.scrollX || 0;
+            var sy = window.scrollY || 0;
+            var c = document.createElement('div');
+            c.id = id;
+            c.style.cssText = 'position:absolute;top:0;left:0;width:0;height:0;pointer-events:none;z-index:2147483647;';
+            for (var i = 0; i < items.length; i++) {{
+                var it = items[i];
+                var dx = it.x + sx;
+                var dy = it.y + sy;
+                var b = document.createElement('div');
+                b.style.cssText = 'position:absolute;left:' + dx + 'px;top:' + dy + 'px;width:' + it.width + 'px;height:' + it.height + 'px;border:2px solid rgba(255,0,0,0.8);box-sizing:border-box;pointer-events:none;';
+                var l = document.createElement('div');
+                l.textContent = String(it.number);
+                var labelTop = dy < 14 ? '2px' : '-14px';
+                l.style.cssText = 'position:absolute;top:' + labelTop + ';left:-2px;background:rgba(255,0,0,0.9);color:#fff;font:bold 11px/14px monospace;padding:0 4px;border-radius:2px;white-space:nowrap;';
+                b.appendChild(l);
+                c.appendChild(b);
+            }}
+            document.documentElement.appendChild(c);
+            return true;
+        }})()"#,
+        items = serde_json::to_string(&overlay_data).unwrap_or_else(|_| "[]".to_string()),
+        overlay_id =
+            serde_json::to_string(ANNOTATION_OVERLAY_ID).unwrap_or_else(|_| "\"\"".to_string()),
+    );
+
+    let _: EvaluateResult = client
+        .send_command_typed(
+            "Runtime.evaluate",
+            &EvaluateParams {
+                expression,
+                return_by_value: Some(true),
+                await_promise: Some(false),
+            },
+            Some(session_id),
+        )
+        .await?;
+
+    Ok(())
+}
+
+async fn remove_annotation_overlay(client: &CdpClient, session_id: &str) -> Result<(), String> {
+    let expression = format!(
+        r#"(() => {{
+            var el = document.getElementById({overlay_id});
+            if (el) el.remove();
+            return true;
+        }})()"#,
+        overlay_id =
+            serde_json::to_string(ANNOTATION_OVERLAY_ID).unwrap_or_else(|_| "\"\"".to_string()),
+    );
+
+    let _: EvaluateResult = client
+        .send_command_typed(
+            "Runtime.evaluate",
+            &EvaluateParams {
+                expression,
+                return_by_value: Some(true),
+                await_promise: Some(false),
+            },
+            Some(session_id),
+        )
+        .await?;
+
+    Ok(())
+}
+
+async fn get_scroll_offsets(client: &CdpClient, session_id: &str) -> Result<(f64, f64), String> {
+    let result: EvaluateResult = client
+        .send_command_typed(
+            "Runtime.evaluate",
+            &EvaluateParams {
+                expression: "({x: window.scrollX || 0, y: window.scrollY || 0})".to_string(),
+                return_by_value: Some(true),
+                await_promise: Some(false),
+            },
+            Some(session_id),
+        )
+        .await?;
+
+    let value = result.result.value.unwrap_or(Value::Null);
+    let x = value.get("x").and_then(|v| v.as_f64()).unwrap_or(0.0);
+    let y = value.get("y").and_then(|v| v.as_f64()).unwrap_or(0.0);
+    Ok((x, y))
+}
+
+fn project_annotations(
+    annotations: &[RawAnnotation],
+    target_rect: Option<&Rect>,
+    scroll: Option<(f64, f64)>,
+) -> Vec<ScreenshotAnnotation> {
+    annotations
+        .iter()
+        .map(|annotation| {
+            let rect = if let Some(target) = target_rect {
+                Rect {
+                    x: annotation.rect.x - target.x,
+                    y: annotation.rect.y - target.y,
+                    width: annotation.rect.width,
+                    height: annotation.rect.height,
+                }
+            } else if let Some((scroll_x, scroll_y)) = scroll {
+                Rect {
+                    x: annotation.rect.x + scroll_x,
+                    y: annotation.rect.y + scroll_y,
+                    width: annotation.rect.width,
+                    height: annotation.rect.height,
+                }
+            } else {
+                annotation.rect.clone()
+            };
+
+            ScreenshotAnnotation {
+                ref_id: annotation.ref_id.clone(),
+                number: annotation.number,
+                role: annotation.role.clone(),
+                name: annotation.name.clone(),
+                box_: AnnotationBox {
+                    x: round(rect.x),
+                    y: round(rect.y),
+                    width: round(rect.width),
+                    height: round(rect.height),
+                },
+            }
+        })
+        .collect()
+}
+
+fn save_screenshot(
+    base64_data: &str,
+    explicit_path: Option<&str>,
+    ext: &str,
+) -> Result<String, String> {
+    let save_path = match explicit_path {
+        Some(path) => path.to_string(),
         None => {
             let dir = get_screenshot_dir();
             let _ = std::fs::create_dir_all(&dir);
@@ -127,13 +494,17 @@ pub async fn take_screenshot(
         }
     };
 
-    let bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, &result.data)
+    let bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, base64_data)
         .map_err(|e| format!("Failed to decode screenshot: {}", e))?;
 
     std::fs::write(&save_path, &bytes)
         .map_err(|e| format!("Failed to save screenshot to {}: {}", save_path, e))?;
 
-    Ok((save_path, result.data))
+    Ok(save_path)
+}
+
+fn round(value: f64) -> i64 {
+    value.round() as i64
 }
 
 fn get_screenshot_dir() -> PathBuf {
@@ -145,3 +516,96 @@ fn get_screenshot_dir() -> PathBuf {
             .join("screenshots")
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn filters_annotations_to_target_overlap() {
+        let annotations = vec![
+            RawAnnotation {
+                ref_id: "e1".to_string(),
+                number: 1,
+                role: "button".to_string(),
+                name: Some("Inside".to_string()),
+                rect: Rect {
+                    x: 10.0,
+                    y: 10.0,
+                    width: 50.0,
+                    height: 20.0,
+                },
+            },
+            RawAnnotation {
+                ref_id: "e2".to_string(),
+                number: 2,
+                role: "button".to_string(),
+                name: Some("Outside".to_string()),
+                rect: Rect {
+                    x: 200.0,
+                    y: 200.0,
+                    width: 40.0,
+                    height: 20.0,
+                },
+            },
+        ];
+
+        let target = Rect {
+            x: 0.0,
+            y: 0.0,
+            width: 100.0,
+            height: 100.0,
+        };
+
+        let filtered = filter_annotations(annotations, Some(&target));
+        assert_eq!(filtered.len(), 1);
+        assert_eq!(filtered[0].ref_id, "e1");
+    }
+
+    #[test]
+    fn projects_selector_annotations_relative_to_target() {
+        let annotations = vec![RawAnnotation {
+            ref_id: "e1".to_string(),
+            number: 1,
+            role: "button".to_string(),
+            name: Some("Inside".to_string()),
+            rect: Rect {
+                x: 25.0,
+                y: 35.0,
+                width: 40.0,
+                height: 20.0,
+            },
+        }];
+
+        let target = Rect {
+            x: 10.0,
+            y: 15.0,
+            width: 100.0,
+            height: 100.0,
+        };
+
+        let projected = project_annotations(&annotations, Some(&target), None);
+        assert_eq!(projected[0].box_.x, 15);
+        assert_eq!(projected[0].box_.y, 20);
+    }
+
+    #[test]
+    fn projects_full_page_annotations_to_document_space() {
+        let annotations = vec![RawAnnotation {
+            ref_id: "e1".to_string(),
+            number: 1,
+            role: "button".to_string(),
+            name: Some("Bottom".to_string()),
+            rect: Rect {
+                x: 5.0,
+                y: 12.0,
+                width: 40.0,
+                height: 20.0,
+            },
+        }];
+
+        let projected = project_annotations(&annotations, None, Some((10.0, 1000.0)));
+        assert_eq!(projected[0].box_.x, 15);
+        assert_eq!(projected[0].box_.y, 1012);
+    }
+}