Add screenshot support

Cargo.toml

@@ -15,6 +15,7 @@ toml = "0.9"
 anyhow = "1.0"
 uuid = { version = "1.0", features = ["v4"] }
 url = "2.5"
+base64 = "0.22"
 
 [target.'cfg(windows)'.dependencies]
 winreg = "0.52"

extension/manifest.json

@@ -3,7 +3,7 @@
   "name": "Browser CLI",
   "version": "0.2.1",
   "description": "Bridges the browser to browser-cli: manages sessions, routes commands to content scripts, and streams structured page snapshots over Native Messaging.",
-  "permissions": ["nativeMessaging", "tabs", "scripting", "storage", "webNavigation"],
+  "permissions": ["nativeMessaging", "tabs", "scripting", "storage", "webNavigation", "activeTab"],
   "background": {
     "scripts": ["dist/background/service-worker.js"],
     "service_worker": "dist/background/service-worker.js"

extension/src/background/service-worker.ts

@@ -85,6 +85,8 @@ async function handleRequest(req: Request): Promise<Response> {
       case 'type':
       case 'wait':
         return await forwardToContent(req);
+      case 'screenshot':
+        return await handleScreenshot(req);
       default:
         return { id: req.id, ok: false, error: `Unknown action: ${req.action}` };
     }
@@ -295,6 +297,51 @@ async function forwardToContent(req: Request): Promise<Response> {
   };
 }
 
+async function handleScreenshot(req: Request): Promise<Response> {
+  const session = sessionFromRequest(req);
+  if (!session.ok) {
+    return { id: req.id, ok: false, error: session.error };
+  }
+
+  await ensureTabLoaded(session.value.tab_id);
+
+  const quality = typeof req.params.quality === 'number' ? req.params.quality : undefined;
+  const format: 'png' | 'jpeg' = quality !== undefined ? 'jpeg' : 'png';
+  const options: Parameters<typeof chrome.tabs.captureVisibleTab>[1] = { format };
+  if (quality !== undefined) {
+    options.quality = quality;
+  }
+
+  // Get the window ID for the session's tab
+  const tab = await chrome.tabs.get(session.value.tab_id);
+  if (!tab.windowId) {
+    return { id: req.id, ok: false, error: 'Could not determine window for tab' };
+  }
+
+  // Record the currently active tab so we can restore focus after capture
+  const [previousTab] = await chrome.tabs.query({ active: true, windowId: tab.windowId });
+
+  // Ensure the tab is active in its window before capturing
+  await chrome.tabs.update(session.value.tab_id, { active: true });
+
+  const dataUrl = await chrome.tabs.captureVisibleTab(tab.windowId, options);
+  const base64Data = dataUrl.split(',')[1];
+
+  // Restore previously active tab if it was different
+  if (previousTab && previousTab.id !== undefined && previousTab.id !== tab.id) {
+    await chrome.tabs.update(previousTab.id, { active: true });
+  }
+
+  return {
+    id: req.id,
+    ok: true,
+    data: {
+      image: base64Data,
+      format: format,
+    },
+  };
+}
+
 function sessionFromRequest(
   req: Request,
 ): { ok: true; value: Session } | { ok: false; error: string } {

src/cli/commands.rs

@@ -920,6 +920,75 @@ pub fn plugin_list(json_mode: bool) -> Result<()> {
     Ok(())
 }
 
+pub async fn screenshot(
+    session_id: &str,
+    output: Option<&str>,
+    full_page: bool,
+    quality: Option<u32>,
+    json_mode: bool,
+) -> Result<()> {
+    use base64::Engine as _;
+    use std::time::{SystemTime, UNIX_EPOCH};
+
+    if full_page {
+        eprintln!("Warning: --full-page is not yet supported; capturing viewport only.");
+    }
+
+    let mut params = json!({ "session_id": session_id, "full_page": false });
+    if let Some(q) = quality {
+        params["quality"] = json!(q);
+    }
+
+    let data = send_ok(Request::new(actions::SCREENSHOT, params)).await?;
+
+    let image_b64 = data
+        .get("image")
+        .and_then(|v| v.as_str())
+        .ok_or_else(|| anyhow::anyhow!("missing image data in response"))?;
+    let format = data
+        .get("format")
+        .and_then(|v| v.as_str())
+        .unwrap_or("png");
+
+    let image_bytes = base64::engine::general_purpose::STANDARD
+        .decode(image_b64)
+        .map_err(|e| anyhow::anyhow!("failed to decode base64 image: {e}"))?;
+
+    let extension = if format == "jpeg" { "jpg" } else { "png" };
+    let output_path = match output {
+        Some(p) => PathBuf::from(p),
+        None => {
+            let ts = SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .unwrap_or_default()
+                .as_millis();
+            PathBuf::from(format!("screenshot-{ts}.{extension}"))
+        }
+    };
+
+    fs::write(&output_path, &image_bytes)?;
+
+    let size_bytes = image_bytes.len();
+
+    if json_mode {
+        print_json(&json!({
+            "session_id": session_id,
+            "path": output_path.display().to_string(),
+            "format": format,
+            "size_bytes": size_bytes,
+        }))?;
+    } else {
+        println!(
+            "Screenshot saved: {} ({} bytes, {})",
+            output_path.display(),
+            size_bytes,
+            format
+        );
+    }
+
+    Ok(())
+}
+
 async fn fetch_snapshot(
     session_id: &str,
     action: &str,

src/main.rs

@@ -276,6 +276,23 @@ enum Command {
         #[arg(long, short = 'v')]
         verbose: bool,
     },
+    /// Capture screenshot of the current page
+    Screenshot {
+        /// Session ID
+        session_id: String,
+        /// Output file path (default: screenshot-<timestamp>.png)
+        #[arg(short, long)]
+        output: Option<String>,
+        /// Capture full page instead of just the viewport
+        #[arg(long)]
+        full_page: bool,
+        /// Image quality for JPEG (0-100, default: PNG format)
+        #[arg(long, value_parser = clap::value_parser!(u32).range(0..=100))]
+        quality: Option<u32>,
+        /// Output as JSON
+        #[arg(long)]
+        json: bool,
+    },
     /// Manage and run plugins
     Plugin {
         #[command(subcommand)]
@@ -465,6 +482,13 @@ async fn main() -> anyhow::Result<()> {
             json,
             verbose,
         } => cli::commands::view(session_id, target, page, fresh, json, verbose).await?,
+        Command::Screenshot {
+            ref session_id,
+            ref output,
+            full_page,
+            quality,
+            json,
+        } => cli::commands::screenshot(session_id, output.as_deref(), full_page, quality, json).await?,
         Command::Plugin { ref cmd } => match cmd {
             PluginCommand::Run {
                 name,

src/protocol/messages.rs

@@ -152,6 +152,7 @@ pub mod actions {
     pub const TYPE: &str = "type";
     pub const WAIT: &str = "wait";
     pub const GET_TEXT: &str = "get_text";
+    pub const SCREENSHOT: &str = "screenshot";
 }
 
 pub const PAGE_CHUNK_TYPE: &str = "page_chunk";