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 every tool (which overwhelms LLM context windows),
31/// select a profile that matches your agent's task. No profile includes
32/// `browser_evaluate_js`; enable it with
33/// [`BrowserToolset::with_evaluate_js`].
34///
35/// # Example
36///
37/// ```rust,ignore
38/// use adk_browser::{BrowserToolset, BrowserProfile, BrowserSession, BrowserConfig};
39/// use std::sync::Arc;
40///
41/// let browser = Arc::new(BrowserSession::new(BrowserConfig::default()));
42/// let toolset = BrowserToolset::with_profile(browser, BrowserProfile::FormFilling);
43/// let tools = toolset.all_tools(); // 19 tools instead of 45
44/// ```
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub enum BrowserProfile {
47    /// 19 tools: navigation + interaction + extraction + wait + screenshot.
48    /// Best for simple browsing tasks.
49    Minimal,
50    /// 19 tools: same categories as Minimal (navigation + interaction + extraction + wait + screenshot).
51    /// Best for form-filling agents.
52    FormFilling,
53    /// 13 tools: navigation + extraction + screenshot + scroll, hover, and alert handling.
54    /// Best for data extraction / scraping agents (no interaction tools).
55    Scraping,
56    /// 45 tools: every tool except `browser_evaluate_js`. Use only when the agent
57    /// needs full browser control.
58    Full,
59}
60
61/// A toolset that provides all browser automation tools.
62///
63/// Use this to add all browser tools to an agent at once, or use
64/// individual tools for more control.
65///
66/// Two defaults limit what a model can reach through the browser:
67///
68/// - `browser_evaluate_js`, which runs model-written JavaScript in the page, is
69///   off until [`with_evaluate_js`](Self::with_evaluate_js) or
70///   [`with_js`](Self::with_js) enables it.
71/// - Navigation accepts only `http` and `https` URLs until
72///   [`with_allowed_schemes`](Self::with_allowed_schemes) says otherwise, so
73///   `file:`, `javascript:`, `data:`, and browser-internal URLs are refused.
74pub struct BrowserToolset {
75    resolver: SessionResolver,
76    /// Include navigation tools (navigate, back, forward, refresh)
77    include_navigation: bool,
78    /// Include interaction tools (click, type, select)
79    include_interaction: bool,
80    /// Include extraction tools (extract text, attributes, links, page info)
81    include_extraction: bool,
82    /// Include wait tools
83    include_wait: bool,
84    /// Include screenshot tool
85    include_screenshot: bool,
86    /// Include the JavaScript helper tools (scroll, hover, alert handling)
87    include_js: bool,
88    /// Include `browser_evaluate_js`, which runs arbitrary model-written JavaScript
89    include_evaluate_js: bool,
90    /// Include cookie management tools
91    include_cookies: bool,
92    /// Include window/tab management tools
93    include_windows: bool,
94    /// Include frame/iframe management tools
95    include_frames: bool,
96    /// Include advanced action tools (drag-drop, focus, file upload, etc.)
97    include_actions: bool,
98    /// URL schemes the navigation and new-tab/window tools accept
99    allowed_schemes: Vec<String>,
100}
101
102impl BrowserToolset {
103    /// Create a new toolset with every tool enabled except `browser_evaluate_js`.
104    ///
105    /// Equivalent to [`with_profile`](Self::with_profile) with [`BrowserProfile::Full`].
106    pub fn new(browser: Arc<BrowserSession>) -> Self {
107        Self::from_profile(SessionResolver::Fixed(browser), BrowserProfile::Full)
108    }
109
110    /// Create a pool-backed toolset with every tool enabled except `browser_evaluate_js`.
111    ///
112    /// Sessions are resolved per-user at runtime via `Toolset::tools(ctx)`.
113    /// The pool calls `get_or_create(ctx.user_id())` to obtain an isolated
114    /// browser session for each user.
115    ///
116    /// # Example
117    ///
118    /// ```rust,ignore
119    /// use adk_browser::{BrowserToolset, BrowserSessionPool, BrowserConfig};
120    /// use std::sync::Arc;
121    ///
122    /// let pool = Arc::new(BrowserSessionPool::new(BrowserConfig::default()));
123    /// let toolset = BrowserToolset::with_pool(pool);
124    /// ```
125    pub fn with_pool(pool: Arc<BrowserSessionPool>) -> Self {
126        Self::from_profile(SessionResolver::Pool(pool), BrowserProfile::Full)
127    }
128
129    /// Create a pool-backed toolset with a pre-configured profile.
130    ///
131    /// Combines pool-backed session resolution with profile-based tool
132    /// category selection.
133    ///
134    /// # Example
135    ///
136    /// ```rust,ignore
137    /// use adk_browser::{BrowserToolset, BrowserProfile, BrowserSessionPool, BrowserConfig};
138    /// use std::sync::Arc;
139    ///
140    /// let pool = Arc::new(BrowserSessionPool::new(BrowserConfig::default()));
141    /// let toolset = BrowserToolset::with_pool_and_profile(pool, BrowserProfile::Minimal);
142    /// ```
143    pub fn with_pool_and_profile(pool: Arc<BrowserSessionPool>, profile: BrowserProfile) -> Self {
144        Self::from_profile(SessionResolver::Pool(pool), profile)
145    }
146
147    /// Create a toolset with a pre-configured profile.
148    ///
149    /// This is the recommended way to create a toolset for most agents.
150    /// Using `BrowserProfile::Full` is equivalent to `BrowserToolset::new()`.
151    pub fn with_profile(browser: Arc<BrowserSession>, profile: BrowserProfile) -> Self {
152        Self::from_profile(SessionResolver::Fixed(browser), profile)
153    }
154
155    fn from_profile(resolver: SessionResolver, profile: BrowserProfile) -> Self {
156        let (interaction, wait, js, extended) = match profile {
157            BrowserProfile::Minimal | BrowserProfile::FormFilling => (true, true, false, false),
158            BrowserProfile::Scraping => (false, false, true, false),
159            BrowserProfile::Full => (true, true, true, true),
160        };
161        Self {
162            resolver,
163            include_navigation: true,
164            include_interaction: interaction,
165            include_extraction: true,
166            include_wait: wait,
167            include_screenshot: true,
168            include_js: js,
169            include_evaluate_js: false,
170            include_cookies: extended,
171            include_windows: extended,
172            include_frames: extended,
173            include_actions: extended,
174            allowed_schemes: default_allowed_schemes(),
175        }
176    }
177
178    /// Enable or disable navigation tools.
179    pub fn with_navigation(mut self, enabled: bool) -> Self {
180        self.include_navigation = enabled;
181        self
182    }
183
184    /// Enable or disable interaction tools.
185    pub fn with_interaction(mut self, enabled: bool) -> Self {
186        self.include_interaction = enabled;
187        self
188    }
189
190    /// Enable or disable extraction tools.
191    pub fn with_extraction(mut self, enabled: bool) -> Self {
192        self.include_extraction = enabled;
193        self
194    }
195
196    /// Enable or disable wait tools.
197    pub fn with_wait(mut self, enabled: bool) -> Self {
198        self.include_wait = enabled;
199        self
200    }
201
202    /// Enable or disable screenshot tool.
203    pub fn with_screenshot(mut self, enabled: bool) -> Self {
204        self.include_screenshot = enabled;
205        self
206    }
207
208    /// Enable or disable every JavaScript tool: scroll, hover, alert handling, and
209    /// `browser_evaluate_js`.
210    ///
211    /// `with_js(true)` is an explicit opt-in to `browser_evaluate_js`, which runs
212    /// model-written JavaScript in the page. To keep the helpers without it, follow
213    /// with `.with_evaluate_js(false)`.
214    pub fn with_js(mut self, enabled: bool) -> Self {
215        self.include_js = enabled;
216        self.include_evaluate_js = enabled;
217        self
218    }
219
220    /// Enable or disable `browser_evaluate_js`.
221    ///
222    /// Off by default. The tool runs arbitrary model-written JavaScript in the
223    /// current page, with the page's cookies and session, so a prompt-injected page
224    /// can steer it. Enable it only for agents that need it.
225    ///
226    /// # Example
227    ///
228    /// ```rust,ignore
229    /// use adk_browser::{BrowserConfig, BrowserSession, BrowserToolset};
230    /// use std::sync::Arc;
231    ///
232    /// let browser = Arc::new(BrowserSession::new(BrowserConfig::default()));
233    /// let toolset = BrowserToolset::new(browser).with_evaluate_js(true);
234    /// ```
235    pub fn with_evaluate_js(mut self, enabled: bool) -> Self {
236        self.include_evaluate_js = enabled;
237        self
238    }
239
240    /// Replace the URL schemes `browser_navigate`, `browser_new_tab`, and
241    /// `browser_new_window` accept.
242    ///
243    /// Defaults to [`DEFAULT_ALLOWED_SCHEMES`](crate::DEFAULT_ALLOWED_SCHEMES)
244    /// (`http`, `https`). Schemes are compared case-insensitively. Allowing `file`
245    /// exposes the host filesystem to the model, and allowing `javascript` or `data`
246    /// lets it run script without `browser_evaluate_js`.
247    ///
248    /// # Example
249    ///
250    /// ```rust,ignore
251    /// use adk_browser::{BrowserConfig, BrowserSession, BrowserToolset};
252    /// use std::sync::Arc;
253    ///
254    /// let browser = Arc::new(BrowserSession::new(BrowserConfig::default()));
255    /// // HTTPS only.
256    /// let toolset = BrowserToolset::new(browser).with_allowed_schemes(["https"]);
257    /// ```
258    #[must_use]
259    pub fn with_allowed_schemes<I, S>(mut self, schemes: I) -> Self
260    where
261        I: IntoIterator<Item = S>,
262        S: Into<String>,
263    {
264        self.allowed_schemes = schemes.into_iter().map(Into::into).collect();
265        self
266    }
267
268    /// Enable or disable cookie management tools.
269    pub fn with_cookies(mut self, enabled: bool) -> Self {
270        self.include_cookies = enabled;
271        self
272    }
273
274    /// Enable or disable window/tab management tools.
275    pub fn with_windows(mut self, enabled: bool) -> Self {
276        self.include_windows = enabled;
277        self
278    }
279
280    /// Enable or disable frame/iframe management tools.
281    pub fn with_frames(mut self, enabled: bool) -> Self {
282        self.include_frames = enabled;
283        self
284    }
285
286    /// Enable or disable advanced action tools.
287    pub fn with_actions(mut self, enabled: bool) -> Self {
288        self.include_actions = enabled;
289        self
290    }
291
292    /// Get all tools as a vector (synchronous version).
293    ///
294    /// Works for fixed-session toolsets. For pool-backed toolsets,
295    /// returns an empty vec with a warning — use `Toolset::tools(ctx)` or
296    /// `try_all_tools()` instead.
297    pub fn all_tools(&self) -> Vec<Arc<dyn Tool>> {
298        match &self.resolver {
299            SessionResolver::Fixed(session) => self.build_tools(session.clone()),
300            SessionResolver::Pool(_) => {
301                tracing::warn!(
302                    "BrowserToolset::all_tools() called on a pool-backed toolset. \
303                     Returns empty vec. Use Toolset::tools(ctx) instead."
304                );
305                Vec::new()
306            }
307        }
308    }
309
310    /// Try to get all tools synchronously. Returns an error for pool-backed toolsets.
311    ///
312    /// Prefer `Toolset::tools(ctx)` for pool-backed toolsets.
313    pub fn try_all_tools(&self) -> Result<Vec<Arc<dyn Tool>>> {
314        match &self.resolver {
315            SessionResolver::Fixed(session) => Ok(self.build_tools(session.clone())),
316            SessionResolver::Pool(_) => Err(adk_core::AdkError::tool(
317                "Cannot resolve tools synchronously for a pool-backed BrowserToolset. \
318                 Use Toolset::tools(ctx) instead.",
319            )),
320        }
321    }
322
323    /// Internal: build tool instances from a resolved session.
324    fn build_tools(&self, browser: Arc<BrowserSession>) -> Vec<Arc<dyn Tool>> {
325        let mut tools: Vec<Arc<dyn Tool>> = Vec::new();
326
327        if self.include_navigation {
328            tools.push(Arc::new(
329                NavigateTool::new(browser.clone())
330                    .with_allowed_schemes(self.allowed_schemes.clone()),
331            ));
332            tools.push(Arc::new(BackTool::new(browser.clone())));
333            tools.push(Arc::new(ForwardTool::new(browser.clone())));
334            tools.push(Arc::new(RefreshTool::new(browser.clone())));
335        }
336
337        if self.include_interaction {
338            tools.push(Arc::new(ClickTool::new(browser.clone())));
339            tools.push(Arc::new(DoubleClickTool::new(browser.clone())));
340            tools.push(Arc::new(TypeTool::new(browser.clone())));
341            tools.push(Arc::new(ClearTool::new(browser.clone())));
342            tools.push(Arc::new(SelectTool::new(browser.clone())));
343        }
344
345        if self.include_extraction {
346            tools.push(Arc::new(ExtractTextTool::new(browser.clone())));
347            tools.push(Arc::new(ExtractAttributeTool::new(browser.clone())));
348            tools.push(Arc::new(ExtractLinksTool::new(browser.clone())));
349            tools.push(Arc::new(PageInfoTool::new(browser.clone())));
350            tools.push(Arc::new(PageSourceTool::new(browser.clone())));
351        }
352
353        if self.include_wait {
354            tools.push(Arc::new(WaitForElementTool::new(browser.clone())));
355            tools.push(Arc::new(WaitTool::new()));
356            tools.push(Arc::new(WaitForPageLoadTool::new(browser.clone())));
357            tools.push(Arc::new(WaitForTextTool::new(browser.clone())));
358        }
359
360        if self.include_screenshot {
361            tools.push(Arc::new(ScreenshotTool::new(browser.clone())));
362        }
363
364        if self.include_evaluate_js {
365            tools.push(Arc::new(EvaluateJsTool::new(browser.clone())));
366        }
367
368        if self.include_js {
369            tools.push(Arc::new(ScrollTool::new(browser.clone())));
370            tools.push(Arc::new(HoverTool::new(browser.clone())));
371            tools.push(Arc::new(AlertTool::new(browser.clone())));
372        }
373
374        if self.include_cookies {
375            tools.push(Arc::new(GetCookiesTool::new(browser.clone())));
376            tools.push(Arc::new(GetCookieTool::new(browser.clone())));
377            tools.push(Arc::new(AddCookieTool::new(browser.clone())));
378            tools.push(Arc::new(DeleteCookieTool::new(browser.clone())));
379            tools.push(Arc::new(DeleteAllCookiesTool::new(browser.clone())));
380        }
381
382        if self.include_windows {
383            tools.push(Arc::new(ListWindowsTool::new(browser.clone())));
384            tools.push(Arc::new(
385                NewTabTool::new(browser.clone()).with_allowed_schemes(self.allowed_schemes.clone()),
386            ));
387            tools.push(Arc::new(
388                NewWindowTool::new(browser.clone())
389                    .with_allowed_schemes(self.allowed_schemes.clone()),
390            ));
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        // Every tool except browser_evaluate_js (45)
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}