batuta/agent/tool/
browser.rs1use async_trait::async_trait;
11use tokio::sync::Mutex;
12
13use super::{Tool, ToolResult};
14use crate::agent::capability::Capability;
15use crate::agent::driver::ToolDefinition;
16use crate::serve::backends::PrivacyTier;
17
18use jugar_probar::{Browser, BrowserConfig, Page};
19
20pub struct BrowserTool {
22 config: BrowserConfig,
23 privacy_tier: PrivacyTier,
24 page: Mutex<Option<Page>>,
25}
26
27impl BrowserTool {
28 pub fn new(privacy_tier: PrivacyTier) -> Self {
30 Self { config: BrowserConfig::default(), privacy_tier, page: Mutex::new(None) }
31 }
32
33 fn is_url_allowed(&self, url: &str) -> bool {
35 match self.privacy_tier {
36 PrivacyTier::Sovereign => {
37 url.starts_with("http://localhost")
38 || url.starts_with("http://127.0.0.1")
39 || url.starts_with("https://localhost")
40 || url.starts_with("https://127.0.0.1")
41 || url.starts_with("file://")
42 }
43 PrivacyTier::Private | PrivacyTier::Standard => true,
44 }
45 }
46}
47
48#[async_trait]
49impl Tool for BrowserTool {
50 fn name(&self) -> &'static str {
51 "browser"
52 }
53
54 fn definition(&self) -> ToolDefinition {
55 ToolDefinition {
56 name: "browser".into(),
57 description: "Headless browser automation (navigate, \
58 screenshot, evaluate JS/WASM)"
59 .into(),
60 input_schema: serde_json::json!({
61 "type": "object",
62 "properties": {
63 "action": {
64 "type": "string",
65 "enum": [
66 "navigate", "screenshot", "evaluate",
67 "eval_wasm", "click", "wait_wasm",
68 "console"
69 ],
70 "description": "Browser action to perform"
71 },
72 "url": {
73 "type": "string",
74 "description": "URL to navigate to (navigate)"
75 },
76 "selector": {
77 "type": "string",
78 "description": "CSS selector (screenshot, click)"
79 },
80 "expression": {
81 "type": "string",
82 "description": "JS/WASM expression (evaluate, eval_wasm)"
83 },
84 "clear": {
85 "type": "boolean",
86 "description": "Clear console messages (console)"
87 }
88 },
89 "required": ["action"]
90 }),
91 }
92 }
93
94 async fn execute(&self, input: serde_json::Value) -> ToolResult {
95 let action = match input.get("action").and_then(|a| a.as_str()) {
96 Some(a) => a,
97 None => {
98 return ToolResult::error("missing required field: action");
99 }
100 };
101
102 match action {
103 "navigate" => self.handle_navigate(&input).await,
104 "screenshot" => self.handle_screenshot().await,
105 "evaluate" => self.handle_evaluate(&input).await,
106 "eval_wasm" => self.handle_eval_wasm(&input).await,
107 "click" => self.handle_click(&input).await,
108 "wait_wasm" => self.handle_wait_wasm().await,
109 "console" => self.handle_console().await,
110 _ => ToolResult::error(format!("unknown action: {action}")),
111 }
112 }
113
114 fn required_capability(&self) -> Capability {
115 Capability::Browser
116 }
117
118 fn timeout(&self) -> std::time::Duration {
119 std::time::Duration::from_secs(30)
120 }
121}
122
123impl BrowserTool {
124 async fn ensure_page(&self) -> Result<(), ToolResult> {
125 let mut guard = self.page.lock().await;
126 if guard.is_none() {
127 let browser = Browser::launch(self.config.clone())
128 .await
129 .map_err(|e| ToolResult::error(format!("browser launch: {e}")))?;
130 let new_page = browser
131 .new_page()
132 .await
133 .map_err(|e| ToolResult::error(format!("new page: {e}")))?;
134 *guard = Some(new_page);
135 }
136 Ok(())
137 }
138
139 async fn handle_navigate(&self, input: &serde_json::Value) -> ToolResult {
140 let url = match input.get("url").and_then(|u| u.as_str()) {
141 Some(u) => u,
142 None => return ToolResult::error("navigate: missing url"),
143 };
144
145 if !self.is_url_allowed(url) {
146 return ToolResult::error(format!(
147 "navigate blocked: URL '{url}' not allowed \
148 under {:?} privacy tier",
149 self.privacy_tier
150 ));
151 }
152
153 if let Err(e) = self.ensure_page().await {
154 return e;
155 }
156
157 let mut guard = self.page.lock().await;
158 let Some(page) = guard.as_mut() else {
159 return ToolResult::error("browser page not initialized");
160 };
161 match page.goto(url).await {
162 Ok(()) => {
163 let current = page.current_url().to_string();
164 ToolResult::success(format!("Navigated to: {current}"))
165 }
166 Err(e) => ToolResult::error(format!("navigate failed: {e}")),
167 }
168 }
169
170 async fn handle_screenshot(&self) -> ToolResult {
171 if let Err(e) = self.ensure_page().await {
172 return e;
173 }
174 let guard = self.page.lock().await;
175 let Some(page) = guard.as_ref() else {
176 return ToolResult::error("browser page not initialized");
177 };
178 match page.screenshot().await {
179 Ok(bytes) => {
180 use base64::Engine;
181 let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
182 ToolResult::success(format!(
183 "Screenshot ({} bytes): \
184 data:image/png;base64,{b64}",
185 bytes.len()
186 ))
187 }
188 Err(e) => ToolResult::error(format!("screenshot failed: {e}")),
189 }
190 }
191
192 async fn handle_evaluate(&self, input: &serde_json::Value) -> ToolResult {
193 let expr = match input.get("expression").and_then(|e| e.as_str()) {
194 Some(e) => e,
195 None => {
196 return ToolResult::error("evaluate: missing expression");
197 }
198 };
199
200 if let Err(e) = self.ensure_page().await {
201 return e;
202 }
203 let guard = self.page.lock().await;
204 let Some(page) = guard.as_ref() else {
205 return ToolResult::error("browser page not initialized");
206 };
207 match page.evaluate(expr).await {
208 Ok(val) => ToolResult::success(format!("{val:?}")),
209 Err(e) => ToolResult::error(format!("evaluate failed: {e}")),
210 }
211 }
212
213 async fn handle_eval_wasm(&self, input: &serde_json::Value) -> ToolResult {
214 let expr = match input.get("expression").and_then(|e| e.as_str()) {
215 Some(e) => e,
216 None => {
217 return ToolResult::error("eval_wasm: missing expression");
218 }
219 };
220
221 if let Err(e) = self.ensure_page().await {
222 return e;
223 }
224 let guard = self.page.lock().await;
225 let Some(page) = guard.as_ref() else {
226 return ToolResult::error("browser page not initialized");
227 };
228 match page.eval_wasm::<serde_json::Value>(expr).await {
229 Ok(val) => ToolResult::success(
230 serde_json::to_string_pretty(&val)
231 .unwrap_or_else(|e| format!("{val:?} (serialize error: {e})")),
232 ),
233 Err(e) => ToolResult::error(format!("eval_wasm failed: {e}")),
234 }
235 }
236
237 async fn handle_click(&self, input: &serde_json::Value) -> ToolResult {
238 let selector = match input.get("selector").and_then(|s| s.as_str()) {
239 Some(s) => s,
240 None => {
241 return ToolResult::error("click: missing selector");
242 }
243 };
244
245 if let Err(e) = self.ensure_page().await {
246 return e;
247 }
248 let guard = self.page.lock().await;
249 let Some(page) = guard.as_ref() else {
250 return ToolResult::error("browser page not initialized");
251 };
252 match page.click(selector).await {
253 Ok(()) => ToolResult::success(format!("Clicked: {selector}")),
254 Err(e) => ToolResult::error(format!("click failed: {e}")),
255 }
256 }
257
258 async fn handle_wait_wasm(&self) -> ToolResult {
259 if let Err(e) = self.ensure_page().await {
260 return e;
261 }
262 let mut guard = self.page.lock().await;
263 let Some(page) = guard.as_mut() else {
264 return ToolResult::error("browser page not initialized");
265 };
266 match page.wait_for_wasm_ready().await {
267 Ok(()) => ToolResult::success("WASM runtime ready"),
268 Err(e) => ToolResult::error(format!("wait_wasm failed: {e}")),
269 }
270 }
271
272 async fn handle_console(&self) -> ToolResult {
273 if let Err(e) = self.ensure_page().await {
274 return e;
275 }
276 let guard = self.page.lock().await;
277 let Some(page) = guard.as_ref() else {
278 return ToolResult::error("browser page not initialized");
279 };
280 let msgs = page.console_messages().await;
281 let formatted: Vec<String> =
282 msgs.iter().map(|m| format!("[{}] {}", m.level, m.text)).collect();
283 ToolResult::success(formatted.join("\n"))
284 }
285}
286
287#[cfg(test)]
288mod tests {
289 use super::*;
290
291 #[test]
292 fn test_sovereign_url_restriction() {
293 let tool = BrowserTool::new(PrivacyTier::Sovereign);
294 assert!(tool.is_url_allowed("http://localhost:8080"));
295 assert!(tool.is_url_allowed("http://127.0.0.1:3000"));
296 assert!(tool.is_url_allowed("https://localhost"));
297 assert!(tool.is_url_allowed("file:///tmp/test.html"));
298 assert!(!tool.is_url_allowed("https://example.com"));
299 assert!(!tool.is_url_allowed("http://evil.com"));
300 }
301
302 #[test]
303 fn test_standard_allows_all() {
304 let tool = BrowserTool::new(PrivacyTier::Standard);
305 assert!(tool.is_url_allowed("https://example.com"));
306 assert!(tool.is_url_allowed("http://localhost:8080"));
307 }
308
309 #[test]
310 fn test_private_allows_all() {
311 let tool = BrowserTool::new(PrivacyTier::Private);
312 assert!(tool.is_url_allowed("https://example.com"));
313 }
314
315 #[test]
316 fn test_tool_metadata() {
317 let tool = BrowserTool::new(PrivacyTier::Sovereign);
318 assert_eq!(tool.name(), "browser");
319 assert_eq!(tool.required_capability(), Capability::Browser);
320 assert_eq!(tool.timeout(), std::time::Duration::from_secs(30));
321 }
322
323 #[test]
324 fn test_definition_schema() {
325 let tool = BrowserTool::new(PrivacyTier::Sovereign);
326 let def = tool.definition();
327 assert_eq!(def.name, "browser");
328 let props = def.input_schema.get("properties");
329 assert!(props.is_some());
330 let action = props.expect("props exists").get("action");
331 assert!(action.is_some());
332 }
333
334 #[tokio::test]
335 async fn test_missing_action() {
336 let tool = BrowserTool::new(PrivacyTier::Sovereign);
337 let result = tool.execute(serde_json::json!({})).await;
338 assert!(result.is_error);
339 assert!(result.content.contains("action"));
340 }
341
342 #[tokio::test]
343 async fn test_unknown_action() {
344 let tool = BrowserTool::new(PrivacyTier::Sovereign);
345 let result = tool.execute(serde_json::json!({"action": "fly"})).await;
346 assert!(result.is_error);
347 assert!(result.content.contains("unknown action"));
348 }
349
350 #[tokio::test]
351 async fn test_navigate_missing_url() {
352 let tool = BrowserTool::new(PrivacyTier::Sovereign);
353 let result = tool.execute(serde_json::json!({"action": "navigate"})).await;
354 assert!(result.is_error);
355 assert!(result.content.contains("missing url"));
356 }
357
358 #[tokio::test]
359 async fn test_navigate_blocked_by_privacy() {
360 let tool = BrowserTool::new(PrivacyTier::Sovereign);
361 let result = tool
362 .execute(serde_json::json!({
363 "action": "navigate",
364 "url": "https://example.com"
365 }))
366 .await;
367 assert!(result.is_error);
368 assert!(result.content.contains("blocked"));
369 }
370
371 #[tokio::test]
372 async fn test_evaluate_missing_expression() {
373 let tool = BrowserTool::new(PrivacyTier::Sovereign);
374 let result = tool.execute(serde_json::json!({"action": "evaluate"})).await;
375 assert!(result.is_error);
376 assert!(result.content.contains("missing expression"));
377 }
378
379 #[tokio::test]
380 async fn test_eval_wasm_missing_expression() {
381 let tool = BrowserTool::new(PrivacyTier::Sovereign);
382 let result = tool.execute(serde_json::json!({"action": "eval_wasm"})).await;
383 assert!(result.is_error);
384 assert!(result.content.contains("missing expression"));
385 }
386
387 #[tokio::test]
388 async fn test_click_missing_selector() {
389 let tool = BrowserTool::new(PrivacyTier::Sovereign);
390 let result = tool.execute(serde_json::json!({"action": "click"})).await;
391 assert!(result.is_error);
392 assert!(result.content.contains("missing selector"));
393 }
394}