use super::manager::BrowserManager;
use crate::brain::tools::error::Result;
use crate::brain::tools::r#trait::{Tool, ToolCapability, ToolExecutionContext, ToolResult};
use async_trait::async_trait;
use base64::Engine;
use serde_json::Value;
use std::sync::Arc;
pub struct BrowserScreenshotTool {
manager: Arc<BrowserManager>,
}
impl BrowserScreenshotTool {
pub fn new(manager: Arc<BrowserManager>) -> Self {
Self { manager }
}
}
#[async_trait]
impl Tool for BrowserScreenshotTool {
fn name(&self) -> &str {
"browser_screenshot"
}
fn description(&self) -> &str {
"Capture a screenshot of the current page. Returns base64-encoded PNG. \
Optionally screenshot a specific element by CSS selector."
}
fn input_schema(&self) -> Value {
serde_json::json!({
"type": "object",
"properties": {
"selector": {
"type": "string",
"description": "CSS selector of element to screenshot (default: full page)"
}
}
})
}
fn capabilities(&self) -> Vec<ToolCapability> {
vec![ToolCapability::Network]
}
fn requires_approval(&self) -> bool {
true
}
async fn execute(&self, input: Value, context: &ToolExecutionContext) -> Result<ToolResult> {
let selector = input["selector"].as_str();
let page = match self
.manager
.get_or_create_session_page(context.session_id)
.await
{
Ok(p) => p,
Err(e) => return Ok(ToolResult::error(format!("Browser error: {e}"))),
};
if let Err(e) = page.url().await {
return Ok(ToolResult::error(format!(
"Screenshot failed: page is not responding (CDP connection may be dead). \
Underlying error: {e}. Try `browser_navigate` to reset the page or \
`browser_close` if the session is stuck."
)));
}
let bytes = if let Some(sel) = selector {
let element = match page.find_element(sel).await {
Ok(el) => el,
Err(e) => return Ok(ToolResult::error(format!("Element '{sel}' not found: {e}"))),
};
match element
.screenshot(
chromiumoxide::cdp::browser_protocol::page::CaptureScreenshotFormat::Png,
)
.await
{
Ok(b) => b,
Err(e) => return Ok(ToolResult::error(format!("Element screenshot failed: {e}"))),
}
} else {
match page
.screenshot(
chromiumoxide::cdp::browser_protocol::page::CaptureScreenshotParams::builder()
.format(
chromiumoxide::cdp::browser_protocol::page::CaptureScreenshotFormat::Png,
)
.build(),
)
.await
{
Ok(b) => b,
Err(e) => return Ok(ToolResult::error(format!("Screenshot failed: {e}"))),
}
};
if selector.is_none() {
let hash = BrowserManager::hash_screenshot_bytes(&bytes);
if Some(hash) == self.manager.last_screenshot_hash(context.session_id).await {
return Ok(ToolResult::error(
"Page is identical to your last screenshot. The previous action \
(or no action) produced no visible change. Do not screenshot again — \
take a different action: `browser_click` something, `browser_type` \
into a field, `browser_navigate` to a new URL, or `browser_find` to \
locate an element you can interact with."
.to_string(),
));
}
self.manager
.set_last_screenshot_hash(context.session_id, hash)
.await;
}
let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
Ok(ToolResult::success(format!("data:image/png;base64,{b64}")))
}
}