adk_browser/tools/
screenshot.rs

1//! Screenshot tool for capturing page images.
2
3use crate::session::BrowserSession;
4use adk_core::{Result, Tool, ToolContext};
5use async_trait::async_trait;
6use serde_json::{json, Value};
7use std::sync::Arc;
8
9/// Tool for taking screenshots of the page.
10pub struct ScreenshotTool {
11    browser: Arc<BrowserSession>,
12}
13
14impl ScreenshotTool {
15    /// Create a new screenshot tool with a shared browser session.
16    pub fn new(browser: Arc<BrowserSession>) -> Self {
17        Self { browser }
18    }
19}
20
21#[async_trait]
22impl Tool for ScreenshotTool {
23    fn name(&self) -> &str {
24        "browser_screenshot"
25    }
26
27    fn description(&self) -> &str {
28        "Take a screenshot of the current page or a specific element. Returns base64-encoded PNG image."
29    }
30
31    fn parameters_schema(&self) -> Option<Value> {
32        Some(json!({
33            "type": "object",
34            "properties": {
35                "selector": {
36                    "type": "string",
37                    "description": "Optional CSS selector to screenshot a specific element. If not provided, captures the full page."
38                },
39                "save_to_artifacts": {
40                    "type": "boolean",
41                    "description": "Whether to save the screenshot to artifacts (default: false)"
42                },
43                "artifact_name": {
44                    "type": "string",
45                    "description": "Name for the artifact if saving (default: 'screenshot.png')"
46                }
47            }
48        }))
49    }
50
51    fn response_schema(&self) -> Option<Value> {
52        Some(json!({
53            "type": "object",
54            "properties": {
55                "success": { "type": "boolean" },
56                "base64_image": { "type": "string" },
57                "saved_to_artifacts": { "type": "boolean" },
58                "artifact_name": { "type": "string" }
59            }
60        }))
61    }
62
63    async fn execute(&self, ctx: Arc<dyn ToolContext>, args: Value) -> Result<Value> {
64        let selector = args.get("selector").and_then(|v| v.as_str());
65        let save_to_artifacts =
66            args.get("save_to_artifacts").and_then(|v| v.as_bool()).unwrap_or(false);
67        let artifact_name =
68            args.get("artifact_name").and_then(|v| v.as_str()).unwrap_or("screenshot.png");
69
70        // Take screenshot
71        let base64_image = if let Some(sel) = selector {
72            self.browser.screenshot_element(sel).await?
73        } else {
74            self.browser.screenshot().await?
75        };
76
77        // Optionally save to artifacts
78        let mut saved = false;
79        if save_to_artifacts {
80            if let Some(artifacts) = ctx.artifacts() {
81                use base64::Engine;
82                let image_data =
83                    base64::engine::general_purpose::STANDARD.decode(&base64_image).map_err(
84                        |e| adk_core::AdkError::Tool(format!("Failed to decode base64: {}", e)),
85                    )?;
86
87                let part = adk_core::Part::InlineData {
88                    mime_type: "image/png".to_string(),
89                    data: image_data,
90                };
91
92                artifacts.save(artifact_name, &part).await?;
93                saved = true;
94            }
95        }
96
97        Ok(json!({
98            "success": true,
99            "base64_image": base64_image,
100            "saved_to_artifacts": saved,
101            "artifact_name": if saved { Some(artifact_name) } else { None }
102        }))
103    }
104}