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