use super::{PreparedImage, ToolExecError, context::ToolContext, human_size, resolve_path};
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
use choreo_keystore::ServiceCredential;
use headless_chrome::Tab;
use headless_chrome::protocol::cdp::Page;
use headless_chrome::protocol::cdp::Page::CaptureScreenshotFormatOption;
use headless_chrome::{Browser, LaunchOptions};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::path::Path;
use std::time::Duration;
use tracing::{debug, info, warn};
use url::Url;
#[derive(Debug, Clone, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum WebpageAction {
Content,
Text,
Screenshot,
Pdf,
}
impl WebpageAction {
fn as_str(&self) -> &'static str {
match self {
WebpageAction::Content => "content",
WebpageAction::Text => "text",
WebpageAction::Screenshot => "screenshot",
WebpageAction::Pdf => "pdf",
}
}
}
#[derive(Debug, Default, Deserialize, JsonSchema)]
pub struct RetrieveWebpageArgs {
url: String,
action: Option<WebpageAction>,
wait_ms: Option<u64>,
timeout_ms: Option<u64>,
width: Option<u32>,
height: Option<u32>,
full_page: Option<bool>,
selector: Option<String>,
output_path: Option<String>,
}
const DEFAULT_TIMEOUT_MS: u64 = 30_000;
const MAX_TIMEOUT_MS: u64 = 120_000;
const MAX_WAIT_MS: u64 = 30_000;
const CANDIDATE_NAMES: &[&str] = &[
"chromium",
"chromium-browser",
"google-chrome",
"google-chrome-stable",
"google-chrome-beta",
"google-chrome-unstable",
"chrome",
"microsoft-edge",
"brave-browser",
];
fn resolve_browser_binary() -> Option<std::path::PathBuf> {
for var in ["CHROMIUM_BIN", "CHROME_BIN"] {
if let Some(p) = std::env::var_os(var) {
let candidate = std::path::PathBuf::from(p);
if is_executable(&candidate) {
return Some(candidate);
}
}
}
if let Some(path_var) = std::env::var_os("PATH") {
for dir in std::env::split_paths(&path_var) {
for name in CANDIDATE_NAMES {
let candidate = dir.join(name);
if is_executable(&candidate) {
return Some(candidate);
}
}
}
}
#[cfg(target_os = "macos")]
{
let mac_paths = [
"/Applications/Chromium.app/Contents/MacOS/Chromium",
"/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
"/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge",
];
for p in mac_paths {
let candidate = std::path::PathBuf::from(p);
if is_executable(&candidate) {
return Some(candidate);
}
}
}
#[cfg(target_os = "windows")]
{
for base in [
"C:/Program Files/Chromium/Application/chrome.exe",
"C:/Program Files/Google/Chrome/Application/chrome.exe",
"C:/Program Files (x86)/Google/Chrome/Application/chrome.exe",
"C:/Program Files (x86)/Microsoft/Edge/Application/msedge.exe",
] {
let candidate = std::path::PathBuf::from(base);
if is_executable(&candidate) {
return Some(candidate);
}
}
}
None
}
fn is_executable(path: &std::path::Path) -> bool {
if !path.is_file() {
return false;
}
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
path.metadata()
.map(|m| m.permissions().mode() & 0o111 != 0)
.unwrap_or(false)
}
#[cfg(not(unix))]
{
true
}
}
fn validate_url(url: &str) -> Result<(), ToolExecError> {
let parsed = Url::parse(url).map_err(|e| ToolExecError(format!("invalid URL '{url}': {e}")))?;
match parsed.scheme() {
"http" | "https" | "file" => Ok(()),
other => Err(ToolExecError(format!(
"unsupported URL scheme '{other}'; only http/https/file are allowed"
))),
}
}
fn text_expression(selector: Option<&str>) -> String {
match selector {
Some(sel) => {
let sel = serde_json::to_string(sel).unwrap_or_else(|_| "\"body\"".to_string());
format!(
"(() => {{ const e = document.querySelector({sel}); return e ? e.innerText : ''; }})()"
)
}
None => "(() => { const e = document.body; return e ? e.innerText : ''; })()".to_string(),
}
}
fn html_expression(selector: &str) -> String {
let sel = serde_json::to_string(selector).unwrap_or_else(|_| "\"html\"".to_string());
format!("(() => {{ const e = document.querySelector({sel}); return e ? e.outerHTML : ''; }})()")
}
const PAGE_SIZE_JS: &str = "(() => { const d = document.documentElement; const w = Math.max(d.scrollWidth, d.clientWidth); \
const h = Math.max(d.scrollHeight, d.clientHeight); return w + 'x' + h; })()";
fn page_content_size(tab: &Tab) -> Result<(f64, f64), ToolExecError> {
let obj = tab
.evaluate(PAGE_SIZE_JS, false)
.map_err(|e| ToolExecError(format!("failed to measure page size: {e:#}")))?;
let raw = remote_text(&obj);
let (w, h) = raw
.split_once('x')
.and_then(|(w, h)| Some((w.parse::<f64>().ok()?, h.parse::<f64>().ok()?)))
.unwrap_or((0.0, 0.0));
debug!(raw_width = %raw, "measured page content size");
Ok((w, h))
}
const ELEMENT_BOX_JS_TEMPLATE: &str = "(() => { const e = document.querySelector({sel}); if (!e) return ''; \
const r = e.getBoundingClientRect(); \
return (r.left + window.scrollX) + ',' + (r.top + window.scrollY) + ',' + r.width + ',' + r.height; })()";
fn element_document_box(
tab: &Tab,
selector: &str,
) -> Result<Option<(f64, f64, f64, f64)>, ToolExecError> {
let sel = serde_json::to_string(selector)
.map_err(|e| ToolExecError(format!("failed to encode selector: {e}")))?;
let expr = ELEMENT_BOX_JS_TEMPLATE.replace("{sel}", &sel);
let obj = tab
.evaluate(&expr, false)
.map_err(|e| ToolExecError(format!("failed to measure element box: {e:#}")))?;
let raw = remote_text(&obj);
let parsed: Option<(f64, f64, f64, f64)> = raw
.split(',')
.map(str::parse::<f64>)
.collect::<Result<Vec<_>, _>>()
.ok()
.and_then(|v| {
<[f64; 4]>::try_from(v)
.ok()
.map(|a| (a[0], a[1], a[2], a[3]))
});
debug!(raw = %raw, "measured element document box");
Ok(parsed)
}
fn capture_screenshot(
tab: &Tab,
selector: Option<&str>,
full_page: bool,
) -> Result<Vec<u8>, ToolExecError> {
if let Some(sel) = selector {
let Some((x, y, w, h)) = element_document_box(tab, sel)? else {
return Err(ToolExecError(format!(
"selector '{sel}' matched no element"
)));
};
if w <= 0.0 || h <= 0.0 {
return Err(ToolExecError(format!(
"selector '{sel}' matched an element with a zero-size box"
)));
}
let result = tab
.call_method(Page::CaptureScreenshot {
format: Some(CaptureScreenshotFormatOption::Png),
quality: None,
clip: Some(Page::Viewport {
x,
y,
width: w,
height: h,
scale: 1.0,
}),
from_surface: Some(true),
capture_beyond_viewport: Some(true),
optimize_for_speed: None,
})
.map_err(|e| ToolExecError(format!("failed to screenshot element: {e:#}")))?;
return BASE64
.decode(result.data)
.map_err(|e| ToolExecError(format!("element screenshot decode failed: {e}")));
}
if full_page {
let (w, h) = page_content_size(tab)?;
if w > 0.0 && h > 0.0 {
let result = tab
.call_method(Page::CaptureScreenshot {
format: Some(CaptureScreenshotFormatOption::Png),
quality: None,
clip: Some(Page::Viewport {
x: 0.0,
y: 0.0,
width: w,
height: h,
scale: 1.0,
}),
from_surface: Some(true),
capture_beyond_viewport: Some(true),
optimize_for_speed: None,
})
.map_err(|e| {
ToolExecError(format!("failed to capture full-page screenshot: {e:#}"))
})?;
return BASE64
.decode(result.data)
.map_err(|e| ToolExecError(format!("full-page screenshot decode failed: {e}")));
}
debug!(width = %w, height = %h, "content size was unusable; falling back to viewport shot");
}
tab.capture_screenshot(CaptureScreenshotFormatOption::Png, None, None, true)
.map_err(|e| ToolExecError(format!("failed to capture screenshot: {e:#}")))
}
fn png_dimensions(bytes: &[u8]) -> Option<(u32, u32)> {
const SIGNATURE: [u8; 8] = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
if bytes.len() >= 24 && bytes.get(..8) == Some(&SIGNATURE) && bytes.get(12..16) == Some(b"IHDR")
{
let width = u32::from_be_bytes(bytes.get(16..20)?.try_into().ok()?);
let height = u32::from_be_bytes(bytes.get(20..24)?.try_into().ok()?);
Some((width, height))
} else {
None
}
}
#[derive(Debug)]
pub struct RetrieveWebpageReturn {
pub text: String,
pub image: Option<PreparedImage>,
}
impl Serialize for RetrieveWebpageReturn {
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
serializer.serialize_str(&self.text)
}
}
impl JsonSchema for RetrieveWebpageReturn {
fn schema_name() -> std::borrow::Cow<'static, str> {
std::borrow::Cow::Borrowed("RetrieveWebpageReturn")
}
fn json_schema(_gen: &mut schemars::SchemaGenerator) -> schemars::Schema {
schemars::json_schema!({ "type": "string" })
}
}
pub struct RetrieveWebpage {}
impl RetrieveWebpage {
pub fn new() -> Self {
RetrieveWebpage {}
}
}
impl Default for RetrieveWebpage {
fn default() -> Self {
Self::new()
}
}
impl super::Tool for RetrieveWebpage {
type Args = RetrieveWebpageArgs;
type Return = RetrieveWebpageReturn;
type Error = ToolExecError;
fn name(&self) -> &'static str {
"retrieve_webpage"
}
fn description(&self) -> &'static str {
"Render a URL in a local headless Chromium/Chrome and return page content (HTML), plain text, a screenshot (PNG), or a PDF. Runs locally and offline; requires a chromium/chrome binary already installed (prefers chromium; override with CHROMIUM_BIN). Screenshots are returned inline or saved to output_path; PDFs require output_path. One-shot per call — no persistent session."
}
fn describe_invocation(&self, args: &Self::Args) -> String {
let action = args
.action
.as_ref()
.map(|a| a.as_str())
.unwrap_or(WebpageAction::Content.as_str());
let mut parts = vec![format!(
"Retrieving web page ({action}). URL: {}.",
args.url
)];
if let Some(sel) = args.selector.as_deref() {
parts.push(format!(" Selector: {sel}."));
}
if let Some(out) = args.output_path.as_deref() {
parts.push(format!(" Output: {out}."));
}
parts.concat()
}
fn return_string(ret: &Self::Return) -> String {
ret.text.clone()
}
fn execute(
&self,
args: Self::Args,
_x_credentials: Option<&ServiceCredential>,
working_dir: Option<&Path>,
_ctx: Option<&ToolContext>,
) -> Result<Self::Return, Self::Error> {
let action = args.action.clone().unwrap_or(WebpageAction::Content);
let url = args.url.trim();
let action_str = action.as_str();
validate_url(url)?;
info!(url, action = action_str, "retrieve_webpage: rendering page");
let binary = resolve_browser_binary().ok_or_else(|| {
warn!("retrieve_webpage: no chromium/chrome binary found on PATH");
ToolExecError(
"no chromium or chrome binary found on PATH (or in standard locations). \
Install Chromium/Chrome, or set CHROMIUM_BIN / CHROME_BIN to its path"
.to_string(),
)
})?;
debug!(path = %binary.display(), "resolved browser binary");
let timeout_ms = args
.timeout_ms
.unwrap_or(DEFAULT_TIMEOUT_MS)
.min(MAX_TIMEOUT_MS);
let wait_ms = args.wait_ms.map(|ms| ms.min(MAX_WAIT_MS));
let mut builder = LaunchOptions::default_builder();
builder.headless(true);
builder.path(Some(binary));
builder.window_size(Some((
args.width.unwrap_or(1280),
args.height.unwrap_or(800),
)));
builder.idle_browser_timeout(Duration::from_millis(
timeout_ms.saturating_mul(2).max(60_000),
));
let options = builder.build().map_err(|e| ToolExecError(e.to_string()))?;
let browser = Browser::new(options)
.map_err(|e| ToolExecError(format!("failed to launch headless browser: {e:#}")))?;
debug!("launched headless browser");
let outcome = (|| -> Result<RetrieveWebpageReturn, ToolExecError> {
let tab = browser
.new_tab()
.map_err(|e| ToolExecError(format!("failed to open a tab: {e:#}")))?;
tab.set_default_timeout(Duration::from_millis(timeout_ms));
tab.navigate_to(url)
.map_err(|e| ToolExecError(format!("navigation to {url} failed: {e:#}")))?;
tab.wait_until_navigated()
.map_err(|e| ToolExecError(format!("page never finished loading: {e:#}")))?;
if let Some(ms) = wait_ms
&& ms > 0
{
std::thread::sleep(Duration::from_millis(ms));
}
debug!(timeout_ms, wait_ms, "page navigated; capturing");
Self::capture(&tab, &args, action, url, working_dir)
})();
match &outcome {
Ok(..) => info!(url, action = action_str, "retrieve_webpage: ok"),
Err(e) => warn!(url, action = action_str, error = %e, "retrieve_webpage: failed"),
}
outcome
}
fn extract_image(&self, ret: &Self::Return) -> Option<PreparedImage> {
ret.image.clone()
}
}
impl RetrieveWebpage {
fn capture(
tab: &Tab,
args: &RetrieveWebpageArgs,
action: WebpageAction,
url: &str,
working_dir: Option<&Path>,
) -> Result<RetrieveWebpageReturn, ToolExecError> {
match action {
WebpageAction::Content => match args.selector.as_deref() {
Some(sel) => {
let obj = tab
.evaluate(&html_expression(sel), false)
.map_err(|e| ToolExecError(format!("failed to extract HTML: {e:#}")))?;
Ok(RetrieveWebpageReturn {
text: remote_text(&obj),
image: None,
})
}
None => tab
.get_content()
.map(|text| RetrieveWebpageReturn { text, image: None })
.map_err(|e| ToolExecError(format!("failed to get page HTML: {e:#}"))),
},
WebpageAction::Text => {
let expr = text_expression(args.selector.as_deref());
let obj = tab
.evaluate(&expr, false)
.map_err(|e| ToolExecError(format!("failed to extract text: {e:#}")))?;
Ok(RetrieveWebpageReturn {
text: remote_text(&obj),
image: None,
})
}
WebpageAction::Screenshot => {
let bytes = capture_screenshot(
tab,
args.selector.as_deref(),
args.full_page.unwrap_or(true),
)?;
let (width, height) = png_dimensions(&bytes).unwrap_or((0, 0));
let size = bytes.len();
let alt = Some(format!("Screenshot of {url}"));
let message = match args.output_path.as_deref() {
Some(out) => {
let path = resolve_path(out, working_dir);
write_bytes_with_dirs(&path, &bytes)?;
format!(
"captured screenshot ({width}x{height}, PNG, {size}); saved to {path}",
size = human_size(size as u64),
path = path.display(),
)
}
None => {
format!(
"captured screenshot ({width}x{height}, PNG, {size})",
size = human_size(size as u64),
)
}
};
Ok(RetrieveWebpageReturn {
text: message,
image: Some(PreparedImage {
mime_type: "image/png".to_string(),
data: bytes,
width,
height,
alt,
}),
})
}
WebpageAction::Pdf => {
let out = args.output_path.as_deref().ok_or_else(|| {
ToolExecError(
"pdf action requires output_path so the binary can be saved".to_string(),
)
})?;
let bytes = tab
.print_to_pdf(None)
.map_err(|e| ToolExecError(format!("failed to render PDF: {e:#}")))?;
let path = resolve_path(out, working_dir);
write_bytes_with_dirs(&path, &bytes)?;
Ok(RetrieveWebpageReturn {
text: format!(
"saved PDF ({size}) to {path}",
size = human_size(bytes.len() as u64),
path = path.display(),
),
image: None,
})
}
}
}
}
fn extract_text(value: &serde_json::Value) -> String {
value.as_str().map(str::to_owned).unwrap_or_default()
}
fn remote_text(object: &headless_chrome::protocol::cdp::Runtime::RemoteObject) -> String {
object.value.as_ref().map(extract_text).unwrap_or_default()
}
fn write_bytes_with_dirs(path: &Path, bytes: &[u8]) -> Result<(), ToolExecError> {
if let Some(parent) = path.parent()
&& !parent.as_os_str().is_empty()
{
std::fs::create_dir_all(parent)
.map_err(|e| ToolExecError(format!("failed to create output dir: {e}")))?;
}
std::fs::write(path, bytes)
.map_err(|e| ToolExecError(format!("failed to write '{}': {e}", path.display())))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tools::Tool;
#[test]
fn validate_url_accepts_http_https_and_file() {
for u in [
"https://example.com",
"http://example.com/path?q=1",
"file:///etc/passwd",
"file:///tmp/foo.html",
"file://localhost/etc/passwd",
] {
assert!(validate_url(u).is_ok(), "{u} should be accepted");
}
}
#[test]
fn validate_url_rejects_other_schemes() {
for u in [
"javascript:alert(1)",
"ftp://x",
"data:text/html,hi",
"not a url",
] {
let err = validate_url(u).unwrap_err();
assert!(!err.to_string().is_empty());
}
}
#[test]
fn text_expression_embeds_selector_safely() {
let expr = text_expression(Some("div[data-x=\"y\"]"));
assert!(expr.contains("\\\"y\\\""));
assert!(expr.contains("innerText"));
}
#[test]
fn text_expression_defaults_to_body() {
let expr = text_expression(None);
assert!(expr.contains("document.body"));
assert!(expr.contains("innerText"));
}
#[test]
fn html_expression_uses_outer_html() {
let expr = html_expression("#main");
assert!(expr.contains("#main"));
assert!(expr.contains("outerHTML"));
}
#[test]
fn describe_invocation_defaults_to_content() {
let tool = RetrieveWebpage::new();
let args = RetrieveWebpageArgs {
url: "https://example.com".to_string(),
..RetrieveWebpageArgs::default()
};
let desc = tool.describe_invocation(&args);
assert!(desc.contains("content"));
assert!(desc.contains("https://example.com"));
}
#[test]
fn describe_invocation_includes_selector_and_output() {
let tool = RetrieveWebpage::new();
let args = RetrieveWebpageArgs {
url: "https://example.com".to_string(),
action: Some(WebpageAction::Screenshot),
selector: Some("#main".to_string()),
output_path: Some("shot.png".to_string()),
..RetrieveWebpageArgs::default()
};
let desc = tool.describe_invocation(&args);
assert!(desc.contains("screenshot"));
assert!(desc.contains("#main"));
assert!(desc.contains("shot.png"));
}
#[test]
fn extract_text_handles_string_and_absent_value() {
assert_eq!(extract_text(&serde_json::json!("hello")), "hello");
assert_eq!(extract_text(&serde_json::Value::Null), "");
assert_eq!(extract_text(&serde_json::json!(42)), "");
}
#[test]
fn png_dimensions_reads_ihdr() {
let mut png = vec![0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A];
png.extend_from_slice(&[0, 0, 0, 13]); png.extend_from_slice(b"IHDR");
png.extend_from_slice(&12u32.to_be_bytes());
png.extend_from_slice(&3456u32.to_be_bytes());
assert_eq!(png_dimensions(&png), Some((12, 3456)));
}
#[test]
fn png_dimensions_rejects_garbage() {
assert_eq!(png_dimensions(b"not a png"), None);
assert_eq!(png_dimensions(&[0u8; 32]), None);
}
}