Skip to main content

adk_browser/
toolset.rs

1//! Browser toolset that provides all browser tools as a collection.
2
3use crate::pool::BrowserSessionPool;
4use crate::session::BrowserSession;
5use crate::tools::*;
6use adk_core::{ReadonlyContext, Result, Tool, Toolset};
7use async_trait::async_trait;
8use std::sync::Arc;
9
10/// Internal abstraction for session acquisition.
11/// Not exposed in the public API.
12enum SessionResolver {
13    /// Fixed session — always returns the same session.
14    Fixed(Arc<BrowserSession>),
15    /// Pool-backed — resolves session from pool using user_id from context.
16    Pool(Arc<BrowserSessionPool>),
17}
18
19impl SessionResolver {
20    async fn resolve(&self, ctx: &Arc<dyn ReadonlyContext>) -> Result<Arc<BrowserSession>> {
21        match self {
22            SessionResolver::Fixed(session) => Ok(session.clone()),
23            SessionResolver::Pool(pool) => pool.get_or_create(ctx.user_id()).await,
24        }
25    }
26}
27
28/// Pre-configured tool profiles for common use cases.
29///
30/// Instead of using all 46 tools (which overwhelms LLM context windows),
31/// select a profile that matches your agent's task.
32///
33/// # Example
34///
35/// ```rust,ignore
36/// use adk_browser::{BrowserToolset, BrowserProfile, BrowserSession, BrowserConfig};
37/// use std::sync::Arc;
38///
39/// let browser = Arc::new(BrowserSession::new(BrowserConfig::default()));
40/// let toolset = BrowserToolset::with_profile(browser, BrowserProfile::FormFilling);
41/// let tools = toolset.all_tools(); // 8 tools instead of 46
42/// ```
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub enum BrowserProfile {
45    /// 19 tools: navigation + interaction + extraction + wait + screenshot.
46    /// Best for simple browsing tasks.
47    Minimal,
48    /// 19 tools: same categories as Minimal (navigation + interaction + extraction + wait + screenshot).
49    /// Best for form-filling agents.
50    FormFilling,
51    /// 14 tools: navigation + extraction + screenshot + JS (scroll, hover, evaluate, alert).
52    /// Best for data extraction / scraping agents (no interaction tools).
53    Scraping,
54    /// All 46 tools. Use only when the agent needs full browser control.
55    Full,
56}
57
58/// A toolset that provides all browser automation tools.
59///
60/// Use this to add all browser tools to an agent at once, or use
61/// individual tools for more control.
62pub struct BrowserToolset {
63    resolver: SessionResolver,
64    /// Include navigation tools (navigate, back, forward, refresh)
65    include_navigation: bool,
66    /// Include interaction tools (click, type, select)
67    include_interaction: bool,
68    /// Include extraction tools (extract text, attributes, links, page info)
69    include_extraction: bool,
70    /// Include wait tools
71    include_wait: bool,
72    /// Include screenshot tool
73    include_screenshot: bool,
74    /// Include JavaScript evaluation tools
75    include_js: bool,
76    /// Include cookie management tools
77    include_cookies: bool,
78    /// Include window/tab management tools
79    include_windows: bool,
80    /// Include frame/iframe management tools
81    include_frames: bool,
82    /// Include advanced action tools (drag-drop, focus, file upload, etc.)
83    include_actions: bool,
84}
85
86impl BrowserToolset {
87    /// Create a new toolset with all tools enabled.
88    pub fn new(browser: Arc<BrowserSession>) -> Self {
89        Self {
90            resolver: SessionResolver::Fixed(browser),
91            include_navigation: true,
92            include_interaction: true,
93            include_extraction: true,
94            include_wait: true,
95            include_screenshot: true,
96            include_js: true,
97            include_cookies: true,
98            include_windows: true,
99            include_frames: true,
100            include_actions: true,
101        }
102    }
103
104    /// Create a pool-backed toolset with all tools enabled.
105    ///
106    /// Sessions are resolved per-user at runtime via `Toolset::tools(ctx)`.
107    /// The pool calls `get_or_create(ctx.user_id())` to obtain an isolated
108    /// browser session for each user.
109    ///
110    /// # Example
111    ///
112    /// ```rust,ignore
113    /// use adk_browser::{BrowserToolset, BrowserSessionPool, BrowserConfig};
114    /// use std::sync::Arc;
115    ///
116    /// let pool = Arc::new(BrowserSessionPool::new(BrowserConfig::default()));
117    /// let toolset = BrowserToolset::with_pool(pool);
118    /// ```
119    pub fn with_pool(pool: Arc<BrowserSessionPool>) -> Self {
120        Self {
121            resolver: SessionResolver::Pool(pool),
122            include_navigation: true,
123            include_interaction: true,
124            include_extraction: true,
125            include_wait: true,
126            include_screenshot: true,
127            include_js: true,
128            include_cookies: true,
129            include_windows: true,
130            include_frames: true,
131            include_actions: true,
132        }
133    }
134
135    /// Create a pool-backed toolset with a pre-configured profile.
136    ///
137    /// Combines pool-backed session resolution with profile-based tool
138    /// category selection.
139    ///
140    /// # Example
141    ///
142    /// ```rust,ignore
143    /// use adk_browser::{BrowserToolset, BrowserProfile, BrowserSessionPool, BrowserConfig};
144    /// use std::sync::Arc;
145    ///
146    /// let pool = Arc::new(BrowserSessionPool::new(BrowserConfig::default()));
147    /// let toolset = BrowserToolset::with_pool_and_profile(pool, BrowserProfile::Minimal);
148    /// ```
149    pub fn with_pool_and_profile(pool: Arc<BrowserSessionPool>, profile: BrowserProfile) -> Self {
150        match profile {
151            BrowserProfile::Minimal => Self {
152                resolver: SessionResolver::Pool(pool),
153                include_navigation: true,
154                include_interaction: true,
155                include_extraction: true,
156                include_wait: true,
157                include_screenshot: true,
158                include_js: false,
159                include_cookies: false,
160                include_windows: false,
161                include_frames: false,
162                include_actions: false,
163            },
164            BrowserProfile::FormFilling => Self {
165                resolver: SessionResolver::Pool(pool),
166                include_navigation: true,
167                include_interaction: true,
168                include_extraction: true,
169                include_wait: true,
170                include_screenshot: true,
171                include_js: false,
172                include_cookies: false,
173                include_windows: false,
174                include_frames: false,
175                include_actions: false,
176            },
177            BrowserProfile::Scraping => Self {
178                resolver: SessionResolver::Pool(pool),
179                include_navigation: true,
180                include_interaction: false,
181                include_extraction: true,
182                include_wait: false,
183                include_screenshot: true,
184                include_js: true,
185                include_cookies: false,
186                include_windows: false,
187                include_frames: false,
188                include_actions: false,
189            },
190            BrowserProfile::Full => Self::with_pool(pool),
191        }
192    }
193
194    /// Create a toolset with a pre-configured profile.
195    ///
196    /// This is the recommended way to create a toolset for most agents.
197    /// Using `BrowserProfile::Full` is equivalent to `BrowserToolset::new()`.
198    pub fn with_profile(browser: Arc<BrowserSession>, profile: BrowserProfile) -> Self {
199        match profile {
200            BrowserProfile::Minimal => Self {
201                resolver: SessionResolver::Fixed(browser),
202                include_navigation: true,
203                include_interaction: true,
204                include_extraction: true,
205                include_wait: true,
206                include_screenshot: true,
207                include_js: false,
208                include_cookies: false,
209                include_windows: false,
210                include_frames: false,
211                include_actions: false,
212            },
213            BrowserProfile::FormFilling => Self {
214                resolver: SessionResolver::Fixed(browser),
215                include_navigation: true,
216                include_interaction: true,
217                include_extraction: true,
218                include_wait: true,
219                include_screenshot: true,
220                include_js: false,
221                include_cookies: false,
222                include_windows: false,
223                include_frames: false,
224                include_actions: false,
225            },
226            BrowserProfile::Scraping => Self {
227                resolver: SessionResolver::Fixed(browser),
228                include_navigation: true,
229                include_interaction: false,
230                include_extraction: true,
231                include_wait: false,
232                include_screenshot: true,
233                include_js: true, // scroll only
234                include_cookies: false,
235                include_windows: false,
236                include_frames: false,
237                include_actions: false,
238            },
239            BrowserProfile::Full => Self::new(browser),
240        }
241    }
242
243    /// Enable or disable navigation tools.
244    pub fn with_navigation(mut self, enabled: bool) -> Self {
245        self.include_navigation = enabled;
246        self
247    }
248
249    /// Enable or disable interaction tools.
250    pub fn with_interaction(mut self, enabled: bool) -> Self {
251        self.include_interaction = enabled;
252        self
253    }
254
255    /// Enable or disable extraction tools.
256    pub fn with_extraction(mut self, enabled: bool) -> Self {
257        self.include_extraction = enabled;
258        self
259    }
260
261    /// Enable or disable wait tools.
262    pub fn with_wait(mut self, enabled: bool) -> Self {
263        self.include_wait = enabled;
264        self
265    }
266
267    /// Enable or disable screenshot tool.
268    pub fn with_screenshot(mut self, enabled: bool) -> Self {
269        self.include_screenshot = enabled;
270        self
271    }
272
273    /// Enable or disable JavaScript tools.
274    pub fn with_js(mut self, enabled: bool) -> Self {
275        self.include_js = enabled;
276        self
277    }
278
279    /// Enable or disable cookie management tools.
280    pub fn with_cookies(mut self, enabled: bool) -> Self {
281        self.include_cookies = enabled;
282        self
283    }
284
285    /// Enable or disable window/tab management tools.
286    pub fn with_windows(mut self, enabled: bool) -> Self {
287        self.include_windows = enabled;
288        self
289    }
290
291    /// Enable or disable frame/iframe management tools.
292    pub fn with_frames(mut self, enabled: bool) -> Self {
293        self.include_frames = enabled;
294        self
295    }
296
297    /// Enable or disable advanced action tools.
298    pub fn with_actions(mut self, enabled: bool) -> Self {
299        self.include_actions = enabled;
300        self
301    }
302
303    /// Get all tools as a vector (synchronous version).
304    ///
305    /// Works for fixed-session toolsets. For pool-backed toolsets,
306    /// returns an empty vec with a warning — use `Toolset::tools(ctx)` or
307    /// `try_all_tools()` instead.
308    pub fn all_tools(&self) -> Vec<Arc<dyn Tool>> {
309        match &self.resolver {
310            SessionResolver::Fixed(session) => self.build_tools(session.clone()),
311            SessionResolver::Pool(_) => {
312                tracing::warn!(
313                    "BrowserToolset::all_tools() called on a pool-backed toolset. \
314                     Returns empty vec. Use Toolset::tools(ctx) instead."
315                );
316                Vec::new()
317            }
318        }
319    }
320
321    /// Try to get all tools synchronously. Returns an error for pool-backed toolsets.
322    ///
323    /// Prefer `Toolset::tools(ctx)` for pool-backed toolsets.
324    pub fn try_all_tools(&self) -> Result<Vec<Arc<dyn Tool>>> {
325        match &self.resolver {
326            SessionResolver::Fixed(session) => Ok(self.build_tools(session.clone())),
327            SessionResolver::Pool(_) => Err(adk_core::AdkError::tool(
328                "Cannot resolve tools synchronously for a pool-backed BrowserToolset. \
329                 Use Toolset::tools(ctx) instead.",
330            )),
331        }
332    }
333
334    /// Internal: build tool instances from a resolved session.
335    fn build_tools(&self, browser: Arc<BrowserSession>) -> Vec<Arc<dyn Tool>> {
336        let mut tools: Vec<Arc<dyn Tool>> = Vec::new();
337
338        if self.include_navigation {
339            tools.push(Arc::new(NavigateTool::new(browser.clone())));
340            tools.push(Arc::new(BackTool::new(browser.clone())));
341            tools.push(Arc::new(ForwardTool::new(browser.clone())));
342            tools.push(Arc::new(RefreshTool::new(browser.clone())));
343        }
344
345        if self.include_interaction {
346            tools.push(Arc::new(ClickTool::new(browser.clone())));
347            tools.push(Arc::new(DoubleClickTool::new(browser.clone())));
348            tools.push(Arc::new(TypeTool::new(browser.clone())));
349            tools.push(Arc::new(ClearTool::new(browser.clone())));
350            tools.push(Arc::new(SelectTool::new(browser.clone())));
351        }
352
353        if self.include_extraction {
354            tools.push(Arc::new(ExtractTextTool::new(browser.clone())));
355            tools.push(Arc::new(ExtractAttributeTool::new(browser.clone())));
356            tools.push(Arc::new(ExtractLinksTool::new(browser.clone())));
357            tools.push(Arc::new(PageInfoTool::new(browser.clone())));
358            tools.push(Arc::new(PageSourceTool::new(browser.clone())));
359        }
360
361        if self.include_wait {
362            tools.push(Arc::new(WaitForElementTool::new(browser.clone())));
363            tools.push(Arc::new(WaitTool::new()));
364            tools.push(Arc::new(WaitForPageLoadTool::new(browser.clone())));
365            tools.push(Arc::new(WaitForTextTool::new(browser.clone())));
366        }
367
368        if self.include_screenshot {
369            tools.push(Arc::new(ScreenshotTool::new(browser.clone())));
370        }
371
372        if self.include_js {
373            tools.push(Arc::new(EvaluateJsTool::new(browser.clone())));
374            tools.push(Arc::new(ScrollTool::new(browser.clone())));
375            tools.push(Arc::new(HoverTool::new(browser.clone())));
376            tools.push(Arc::new(AlertTool::new(browser.clone())));
377        }
378
379        if self.include_cookies {
380            tools.push(Arc::new(GetCookiesTool::new(browser.clone())));
381            tools.push(Arc::new(GetCookieTool::new(browser.clone())));
382            tools.push(Arc::new(AddCookieTool::new(browser.clone())));
383            tools.push(Arc::new(DeleteCookieTool::new(browser.clone())));
384            tools.push(Arc::new(DeleteAllCookiesTool::new(browser.clone())));
385        }
386
387        if self.include_windows {
388            tools.push(Arc::new(ListWindowsTool::new(browser.clone())));
389            tools.push(Arc::new(NewTabTool::new(browser.clone())));
390            tools.push(Arc::new(NewWindowTool::new(browser.clone())));
391            tools.push(Arc::new(SwitchWindowTool::new(browser.clone())));
392            tools.push(Arc::new(CloseWindowTool::new(browser.clone())));
393            tools.push(Arc::new(MaximizeWindowTool::new(browser.clone())));
394            tools.push(Arc::new(MinimizeWindowTool::new(browser.clone())));
395            tools.push(Arc::new(SetWindowSizeTool::new(browser.clone())));
396        }
397
398        if self.include_frames {
399            tools.push(Arc::new(SwitchToFrameTool::new(browser.clone())));
400            tools.push(Arc::new(SwitchToParentFrameTool::new(browser.clone())));
401            tools.push(Arc::new(SwitchToDefaultContentTool::new(browser.clone())));
402        }
403
404        if self.include_actions {
405            tools.push(Arc::new(DragAndDropTool::new(browser.clone())));
406            tools.push(Arc::new(RightClickTool::new(browser.clone())));
407            tools.push(Arc::new(FocusTool::new(browser.clone())));
408            tools.push(Arc::new(ElementStateTool::new(browser.clone())));
409            tools.push(Arc::new(PressKeyTool::new(browser.clone())));
410            tools.push(Arc::new(FileUploadTool::new(browser.clone())));
411            tools.push(Arc::new(PrintToPdfTool::new(browser)));
412        }
413
414        tools
415    }
416}
417
418#[async_trait]
419impl Toolset for BrowserToolset {
420    fn name(&self) -> &str {
421        "browser"
422    }
423
424    async fn tools(&self, ctx: Arc<dyn ReadonlyContext>) -> Result<Vec<Arc<dyn Tool>>> {
425        let session = self.resolver.resolve(&ctx).await?;
426        Ok(self.build_tools(session))
427    }
428}
429
430/// Helper function to create a minimal browser toolset with only essential tools.
431pub fn minimal_browser_tools(browser: Arc<BrowserSession>) -> Vec<Arc<dyn Tool>> {
432    vec![
433        Arc::new(NavigateTool::new(browser.clone())),
434        Arc::new(ClickTool::new(browser.clone())),
435        Arc::new(TypeTool::new(browser.clone())),
436        Arc::new(ExtractTextTool::new(browser.clone())),
437        Arc::new(WaitForElementTool::new(browser.clone())),
438        Arc::new(ScreenshotTool::new(browser)),
439    ]
440}
441
442/// Helper function to create a read-only browser toolset (no interaction).
443pub fn readonly_browser_tools(browser: Arc<BrowserSession>) -> Vec<Arc<dyn Tool>> {
444    vec![
445        Arc::new(NavigateTool::new(browser.clone())),
446        Arc::new(ExtractTextTool::new(browser.clone())),
447        Arc::new(ExtractAttributeTool::new(browser.clone())),
448        Arc::new(ExtractLinksTool::new(browser.clone())),
449        Arc::new(PageInfoTool::new(browser.clone())),
450        Arc::new(ScreenshotTool::new(browser.clone())),
451        Arc::new(ScrollTool::new(browser)),
452    ]
453}
454
455#[cfg(test)]
456mod tests {
457    use super::*;
458    use crate::config::BrowserConfig;
459
460    #[test]
461    fn test_toolset_all_tools() {
462        let browser = Arc::new(BrowserSession::new(BrowserConfig::default()));
463        let toolset = BrowserToolset::new(browser);
464        let tools = toolset.all_tools();
465
466        // Should have 46 tools total
467        assert!(tools.len() > 40);
468
469        // Check some tool names exist
470        let tool_names: Vec<&str> = tools.iter().map(|t| t.name()).collect();
471        assert!(tool_names.contains(&"browser_navigate"));
472        assert!(tool_names.contains(&"browser_click"));
473        assert!(tool_names.contains(&"browser_type"));
474        assert!(tool_names.contains(&"browser_screenshot"));
475        // New tools
476        assert!(tool_names.contains(&"browser_get_cookies"));
477        assert!(tool_names.contains(&"browser_new_tab"));
478        assert!(tool_names.contains(&"browser_switch_to_frame"));
479        assert!(tool_names.contains(&"browser_drag_and_drop"));
480    }
481
482    #[test]
483    fn test_toolset_selective() {
484        let browser = Arc::new(BrowserSession::new(BrowserConfig::default()));
485        let toolset = BrowserToolset::new(browser)
486            .with_navigation(true)
487            .with_interaction(false)
488            .with_extraction(false)
489            .with_wait(false)
490            .with_screenshot(false)
491            .with_js(false)
492            .with_cookies(false)
493            .with_windows(false)
494            .with_frames(false)
495            .with_actions(false);
496
497        let tools = toolset.all_tools();
498
499        // Should only have navigation tools
500        assert_eq!(tools.len(), 4);
501    }
502
503    #[test]
504    fn test_minimal_tools() {
505        let browser = Arc::new(BrowserSession::new(BrowserConfig::default()));
506        let tools = minimal_browser_tools(browser);
507
508        assert_eq!(tools.len(), 6);
509    }
510
511    #[test]
512    fn test_profile_form_filling() {
513        let browser = Arc::new(BrowserSession::new(BrowserConfig::default()));
514        let toolset = BrowserToolset::with_profile(browser, BrowserProfile::FormFilling);
515        let tools = toolset.all_tools();
516
517        // Navigation (4) + Interaction (5) + Extraction (5) + Wait (4) + Screenshot (1) = 19
518        let tool_names: Vec<&str> = tools.iter().map(|t| t.name()).collect();
519        assert!(tool_names.contains(&"browser_navigate"));
520        assert!(tool_names.contains(&"browser_click"));
521        assert!(tool_names.contains(&"browser_type"));
522        assert!(tool_names.contains(&"browser_select"));
523        assert!(tool_names.contains(&"browser_clear"));
524        assert!(tool_names.contains(&"browser_screenshot"));
525        // Should NOT include cookies, windows, frames, actions
526        assert!(!tool_names.contains(&"browser_get_cookies"));
527        assert!(!tool_names.contains(&"browser_new_tab"));
528        assert!(!tool_names.contains(&"browser_switch_to_frame"));
529        assert!(!tool_names.contains(&"browser_drag_and_drop"));
530    }
531
532    #[test]
533    fn test_profile_scraping() {
534        let browser = Arc::new(BrowserSession::new(BrowserConfig::default()));
535        let toolset = BrowserToolset::with_profile(browser, BrowserProfile::Scraping);
536        let tools = toolset.all_tools();
537
538        let tool_names: Vec<&str> = tools.iter().map(|t| t.name()).collect();
539        assert!(tool_names.contains(&"browser_navigate"));
540        assert!(tool_names.contains(&"browser_extract_text"));
541        assert!(tool_names.contains(&"browser_extract_links"));
542        assert!(tool_names.contains(&"browser_screenshot"));
543        assert!(tool_names.contains(&"browser_scroll"));
544        // Should NOT include interaction tools
545        assert!(!tool_names.contains(&"browser_click"));
546        assert!(!tool_names.contains(&"browser_type"));
547    }
548
549    #[test]
550    fn test_profile_full_matches_new() {
551        let browser1 = Arc::new(BrowserSession::new(BrowserConfig::default()));
552        let browser2 = Arc::new(BrowserSession::new(BrowserConfig::default()));
553        let full = BrowserToolset::with_profile(browser1, BrowserProfile::Full);
554        let default = BrowserToolset::new(browser2);
555        assert_eq!(full.all_tools().len(), default.all_tools().len());
556    }
557}