browser_use/browser/
session.rs

1use crate::browser::config::{ConnectionOptions, LaunchOptions};
2use crate::dom::DomTree;
3use crate::error::{BrowserError, Result};
4use crate::tools::{ToolContext, ToolRegistry};
5use headless_chrome::{Browser, Tab};
6use std::sync::Arc;
7use std::time::Duration;
8
9/// Browser session that manages a Chrome/Chromium instance
10pub struct BrowserSession {
11    /// The underlying headless_chrome Browser instance
12    browser: Browser,
13
14    /// The active tab for browser operations
15    active_tab: Arc<Tab>,
16
17    /// Tool registry for executing browser automation tools
18    tool_registry: ToolRegistry,
19}
20
21impl BrowserSession {
22    /// Launch a new browser instance with the given options
23    pub fn launch(options: LaunchOptions) -> Result<Self> {
24        let mut launch_opts = headless_chrome::LaunchOptions::default();
25
26        // Set the browser's idle timeout to 1 hour (default is 30 seconds) to prevent the session from closing too soon
27        launch_opts.idle_browser_timeout = Duration::from_secs(60 * 60);
28
29        // Configure headless mode
30        launch_opts.headless = options.headless;
31
32        // Set window size
33        launch_opts.window_size = Some((options.window_width, options.window_height));
34
35        // Set Chrome binary path if provided
36        if let Some(path) = options.chrome_path {
37            launch_opts.path = Some(path);
38        }
39
40        // Set user data directory if provided
41        if let Some(dir) = options.user_data_dir {
42            launch_opts.user_data_dir = Some(dir);
43        }
44
45        // Set sandbox mode
46        launch_opts.sandbox = options.sandbox;
47
48        // Launch browser
49        let browser =
50            Browser::new(launch_opts).map_err(|e| BrowserError::LaunchFailed(e.to_string()))?;
51
52        // Get or create the first tab
53        let active_tab = browser
54            .new_tab()
55            .map_err(|e| BrowserError::LaunchFailed(format!("Failed to create tab: {}", e)))?;
56
57        Ok(Self {
58            browser,
59            active_tab,
60            tool_registry: ToolRegistry::with_defaults(),
61        })
62    }
63
64    /// Connect to an existing browser instance via WebSocket
65    pub fn connect(options: ConnectionOptions) -> Result<Self> {
66        let browser = Browser::connect(options.ws_url)
67            .map_err(|e| BrowserError::ConnectionFailed(e.to_string()))?;
68
69        // Get the first available tab
70        let active_tab = browser
71            .get_tabs()
72            .lock()
73            .map_err(|e| BrowserError::ConnectionFailed(format!("Failed to get tabs: {}", e)))?
74            .first()
75            .ok_or_else(|| BrowserError::ConnectionFailed("No tabs available".to_string()))?
76            .clone();
77
78        Ok(Self {
79            browser,
80            active_tab,
81            tool_registry: ToolRegistry::with_defaults(),
82        })
83    }
84
85    /// Launch a browser with default options
86    pub fn new() -> Result<Self> {
87        Self::launch(LaunchOptions::default())
88    }
89
90    /// Get the active tab
91    pub fn tab(&self) -> &Arc<Tab> {
92        &self.active_tab
93    }
94
95    /// Create a new tab and set it as active
96    pub fn new_tab(&mut self) -> Result<Arc<Tab>> {
97        let tab = self.browser.new_tab().map_err(|e| {
98            BrowserError::TabOperationFailed(format!("Failed to create tab: {}", e))
99        })?;
100
101        self.active_tab = tab.clone();
102        Ok(tab)
103    }
104
105    /// Switch to a specific tab by index
106    pub fn switch_tab(&mut self, index: usize) -> Result<()> {
107        let tabs =
108            self.browser.get_tabs().lock().map_err(|e| {
109                BrowserError::TabOperationFailed(format!("Failed to get tabs: {}", e))
110            })?;
111
112        let tab = tabs
113            .get(index)
114            .ok_or_else(|| {
115                BrowserError::TabOperationFailed(format!("Tab index {} out of range", index))
116            })?
117            .clone();
118
119        self.active_tab = tab;
120        Ok(())
121    }
122
123    /// Get all tabs
124    pub fn get_tabs(&self) -> Result<Vec<Arc<Tab>>> {
125        let tabs = self
126            .browser
127            .get_tabs()
128            .lock()
129            .map_err(|e| BrowserError::TabOperationFailed(format!("Failed to get tabs: {}", e)))?
130            .clone();
131
132        Ok(tabs)
133    }
134
135    /// Close the active tab
136    pub fn close_active_tab(&mut self) -> Result<()> {
137        self.active_tab
138            .close(true)
139            .map_err(|e| BrowserError::TabOperationFailed(format!("Failed to close tab: {}", e)))?;
140
141        // Switch to another tab if available
142        let tabs = self.get_tabs()?;
143        if !tabs.is_empty() {
144            self.active_tab = tabs[0].clone();
145        }
146
147        Ok(())
148    }
149
150    /// Get the underlying Browser instance
151    pub fn browser(&self) -> &Browser {
152        &self.browser
153    }
154
155    /// Navigate to a URL using the active tab
156    pub fn navigate(&self, url: &str) -> Result<()> {
157        self.active_tab.navigate_to(url).map_err(|e| {
158            BrowserError::NavigationFailed(format!("Failed to navigate to {}: {}", url, e))
159        })?;
160
161        Ok(())
162    }
163
164    /// Wait for navigation to complete
165    pub fn wait_for_navigation(&self) -> Result<()> {
166        self.active_tab
167            .wait_until_navigated()
168            .map_err(|e| BrowserError::NavigationFailed(format!("Navigation timeout: {}", e)))?;
169
170        Ok(())
171    }
172
173    /// Extract the DOM tree from the active tab
174    pub fn extract_dom(&self) -> Result<DomTree> {
175        DomTree::from_tab(&self.active_tab)
176    }
177
178    /// Extract and simplify the DOM tree from the active tab
179    pub fn extract_simplified_dom(&self) -> Result<DomTree> {
180        let mut tree = self.extract_dom()?;
181        tree.simplify();
182        Ok(tree)
183    }
184
185    /// Get element selector by index from the last extracted DOM
186    /// Note: You need to extract the DOM first using extract_dom()
187    pub fn find_element<'a>(&'a self, css_selector: &str) -> Result<headless_chrome::Element<'a>> {
188        self.active_tab.find_element(css_selector).map_err(|e| {
189            BrowserError::ElementNotFound(format!("Element '{}' not found: {}", css_selector, e))
190        })
191    }
192
193    /// Get the tool registry
194    pub fn tool_registry(&self) -> &ToolRegistry {
195        &self.tool_registry
196    }
197
198    /// Get mutable tool registry
199    pub fn tool_registry_mut(&mut self) -> &mut ToolRegistry {
200        &mut self.tool_registry
201    }
202
203    /// Execute a tool by name
204    pub fn execute_tool(
205        &self,
206        name: &str,
207        params: serde_json::Value,
208    ) -> Result<crate::tools::ToolResult> {
209        let mut context = ToolContext::new(self);
210        self.tool_registry.execute(name, params, &mut context)
211    }
212
213    /// Navigate back in browser history
214    pub fn go_back(&self) -> Result<()> {
215        let go_back_js = r#"
216            (function() {
217                window.history.back();
218                return true;
219            })()
220        "#;
221
222        self.active_tab
223            .evaluate(go_back_js, false)
224            .map_err(|e| BrowserError::NavigationFailed(format!("Failed to go back: {}", e)))?;
225
226        // Wait a moment for navigation
227        std::thread::sleep(std::time::Duration::from_millis(300));
228
229        Ok(())
230    }
231
232    /// Navigate forward in browser history
233    pub fn go_forward(&self) -> Result<()> {
234        let go_forward_js = r#"
235            (function() {
236                window.history.forward();
237                return true;
238            })()
239        "#;
240
241        self.active_tab
242            .evaluate(go_forward_js, false)
243            .map_err(|e| BrowserError::NavigationFailed(format!("Failed to go forward: {}", e)))?;
244
245        // Wait a moment for navigation
246        std::thread::sleep(std::time::Duration::from_millis(300));
247
248        Ok(())
249    }
250
251    /// Close the browser
252    pub fn close(&self) -> Result<()> {
253        // Note: The Browser struct doesn't have a public close method in headless_chrome
254        // The browser will be closed when the Browser instance is dropped
255        // We can close all tabs to effectively shut down
256        let tabs = self.get_tabs()?;
257        for tab in tabs {
258            let _ = tab.close(false); // Ignore errors on individual tab closes
259        }
260        Ok(())
261    }
262}
263
264impl Default for BrowserSession {
265    fn default() -> Self {
266        Self::new().expect("Failed to create default browser session")
267    }
268}
269
270#[cfg(test)]
271mod tests {
272    use super::*;
273
274    #[test]
275    fn test_launch_options_builder() {
276        let opts = LaunchOptions::new().headless(true).window_size(800, 600);
277
278        assert!(opts.headless);
279        assert_eq!(opts.window_width, 800);
280        assert_eq!(opts.window_height, 600);
281    }
282
283    #[test]
284    fn test_connection_options() {
285        let opts = ConnectionOptions::new("ws://localhost:9222").timeout(5000);
286
287        assert_eq!(opts.ws_url, "ws://localhost:9222");
288        assert_eq!(opts.timeout, 5000);
289    }
290
291    // Integration tests (require Chrome to be installed)
292    #[test]
293    #[ignore] // Ignore by default, run with: cargo test -- --ignored
294    fn test_launch_browser() {
295        let result = BrowserSession::launch(LaunchOptions::new().headless(true));
296        assert!(result.is_ok());
297    }
298
299    #[test]
300    #[ignore]
301    fn test_navigate() {
302        let session = BrowserSession::launch(LaunchOptions::new().headless(true))
303            .expect("Failed to launch browser");
304
305        let result = session.navigate("about:blank");
306        assert!(result.is_ok());
307    }
308
309    #[test]
310    #[ignore]
311    fn test_new_tab() {
312        let mut session = BrowserSession::launch(LaunchOptions::new().headless(true))
313            .expect("Failed to launch browser");
314
315        let result = session.new_tab();
316        assert!(result.is_ok());
317
318        let tabs = session.get_tabs().expect("Failed to get tabs");
319        assert!(tabs.len() >= 2);
320    }
321}