fix(screenshot): support refs and improve error messages (#141)

* fix(screenshot): support refs and improve error messages

* fix(cli): support selector argument in screenshot command

* Fix CSS class selectors being treated as file paths

* fix(test): update screenshot test assertions

---------

Co-authored-by: Chris Tate <chris@ctate.dev>

cli/src/commands.rs

@@ -335,11 +335,33 @@ pub fn parse_command(args: &[String], flags: &Flags) -> Result<Value, ParseError
 
         // === Screenshot/PDF ===
         "screenshot" => {
-            let mut cmd = json!({ "id": id, "action": "screenshot", "fullPage": flags.full });
+            // screenshot [selector] [path]
-            if let Some(path) = rest.get(0) {
+            // selector: @ref or CSS selector
-                cmd["path"] = json!(path);
+            // path: file path (contains / or . or ends with known extension)
+            let (selector, path) = match (rest.get(0), rest.get(1)) {
+                (Some(first), Some(second)) => {
+                    // Two args: first is selector, second is path
+                    (Some(*first), Some(*second))
                 }
-            Ok(cmd)
+                (Some(first), None) => {
+                    // One arg: determine if it's a selector or a path
+                    let is_relative_path = first.starts_with("./") || first.starts_with("../");
+                    let is_selector = !is_relative_path
+                        && (first.starts_with('.') || first.starts_with('#') || first.starts_with('@'));
+                    let has_path_extension = first.ends_with(".png")
+                        || first.ends_with(".jpg")
+                        || first.ends_with(".jpeg")
+                        || first.ends_with(".webp");
+                    let is_path = is_relative_path || first.contains('/') || has_path_extension;
+                    if is_selector || !is_path {
+                        (Some(*first), None)
+                    } else {
+                        (None, Some(*first))
+                    }
+                }
+                _ => (None, None),
+            };
+            Ok(json!({ "id": id, "action": "screenshot", "path": path, "selector": selector, "fullPage": flags.full }))
         }
         "pdf" => {
             let path = rest.get(0).ok_or_else(|| ParseError::MissingArguments {
@@ -1504,7 +1526,8 @@ mod tests {
     fn test_screenshot() {
         let cmd = parse_command(&args("screenshot"), &default_flags()).unwrap();
         assert_eq!(cmd["action"], "screenshot");
-        assert!(cmd.get("path").is_none());
+        assert_eq!(cmd["path"], serde_json::Value::Null);
+        assert_eq!(cmd["selector"], serde_json::Value::Null);
     }
 
     #[test]
@@ -1523,6 +1546,46 @@ mod tests {
         assert_eq!(cmd["fullPage"], true);
     }
 
+    #[test]
+    fn test_screenshot_with_ref() {
+        let cmd = parse_command(&args("screenshot @e1"), &default_flags()).unwrap();
+        assert_eq!(cmd["action"], "screenshot");
+        assert_eq!(cmd["selector"], "@e1");
+        assert_eq!(cmd["path"], serde_json::Value::Null);
+    }
+
+    #[test]
+    fn test_screenshot_with_css_class() {
+        let cmd = parse_command(&args("screenshot .my-button"), &default_flags()).unwrap();
+        assert_eq!(cmd["action"], "screenshot");
+        assert_eq!(cmd["selector"], ".my-button");
+        assert_eq!(cmd["path"], serde_json::Value::Null);
+    }
+
+    #[test]
+    fn test_screenshot_with_css_id() {
+        let cmd = parse_command(&args("screenshot #header"), &default_flags()).unwrap();
+        assert_eq!(cmd["action"], "screenshot");
+        assert_eq!(cmd["selector"], "#header");
+        assert_eq!(cmd["path"], serde_json::Value::Null);
+    }
+
+    #[test]
+    fn test_screenshot_with_path() {
+        let cmd = parse_command(&args("screenshot ./output.png"), &default_flags()).unwrap();
+        assert_eq!(cmd["action"], "screenshot");
+        assert_eq!(cmd["selector"], serde_json::Value::Null);
+        assert_eq!(cmd["path"], "./output.png");
+    }
+
+    #[test]
+    fn test_screenshot_with_selector_and_path() {
+        let cmd = parse_command(&args("screenshot .btn ./button.png"), &default_flags()).unwrap();
+        assert_eq!(cmd["action"], "screenshot");
+        assert_eq!(cmd["selector"], ".btn");
+        assert_eq!(cmd["path"], "./button.png");
+    }
+
     // === Snapshot ===
 
     #[test]

src/actions.ts

@@ -549,9 +549,10 @@ async function handleScreenshot(
 
   let target: Page | ReturnType<Page['locator']> = page;
   if (command.selector) {
-    target = page.locator(command.selector);
+    target = browser.getLocator(command.selector);
   }
 
+  try {
     if (command.path) {
       await target.screenshot({ ...options, path: command.path });
       return successResponse(command.id, { path: command.path });
@@ -559,6 +560,12 @@ async function handleScreenshot(
       const buffer = await target.screenshot(options);
       return successResponse(command.id, { base64: buffer.toString('base64') });
     }
+  } catch (error) {
+    if (command.selector) {
+      throw toAIFriendlyError(error, command.selector);
+    }
+    throw error;
+  }
 }
 
 async function handleSnapshot(