Skip to main content

batuta/agent/tool/
browser.rs

1//! BrowserTool — wraps `jugar_probar::Browser` for headless Chromium.
2//!
3//! Provides agent access to browser automation: navigate, screenshot,
4//! evaluate JS/WASM, click elements. Sovereign privacy tier restricts
5//! navigation to localhost/file:// URLs only.
6//!
7//! Feature-gated: requires `agents-browser` feature for
8//! `jugar_probar` access.
9
10use 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
20/// Tool that wraps jugar-probar for headless Chromium automation.
21pub struct BrowserTool {
22    config: BrowserConfig,
23    privacy_tier: PrivacyTier,
24    page: Mutex<Option<Page>>,
25}
26
27impl BrowserTool {
28    /// Create a new BrowserTool with the given privacy tier.
29    pub fn new(privacy_tier: PrivacyTier) -> Self {
30        Self { config: BrowserConfig::default(), privacy_tier, page: Mutex::new(None) }
31    }
32
33    /// Check if a URL is allowed under the current privacy tier.
34    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}