1use 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
10enum SessionResolver {
13 Fixed(Arc<BrowserSession>),
15 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub enum BrowserProfile {
47 Minimal,
50 FormFilling,
53 Scraping,
56 Full,
59}
60
61pub struct BrowserToolset {
75 resolver: SessionResolver,
76 include_navigation: bool,
78 include_interaction: bool,
80 include_extraction: bool,
82 include_wait: bool,
84 include_screenshot: bool,
86 include_js: bool,
88 include_evaluate_js: bool,
90 include_cookies: bool,
92 include_windows: bool,
94 include_frames: bool,
96 include_actions: bool,
98 allowed_schemes: Vec<String>,
100}
101
102impl BrowserToolset {
103 pub fn new(browser: Arc<BrowserSession>) -> Self {
107 Self::from_profile(SessionResolver::Fixed(browser), BrowserProfile::Full)
108 }
109
110 pub fn with_pool(pool: Arc<BrowserSessionPool>) -> Self {
126 Self::from_profile(SessionResolver::Pool(pool), BrowserProfile::Full)
127 }
128
129 pub fn with_pool_and_profile(pool: Arc<BrowserSessionPool>, profile: BrowserProfile) -> Self {
144 Self::from_profile(SessionResolver::Pool(pool), profile)
145 }
146
147 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 pub fn with_navigation(mut self, enabled: bool) -> Self {
180 self.include_navigation = enabled;
181 self
182 }
183
184 pub fn with_interaction(mut self, enabled: bool) -> Self {
186 self.include_interaction = enabled;
187 self
188 }
189
190 pub fn with_extraction(mut self, enabled: bool) -> Self {
192 self.include_extraction = enabled;
193 self
194 }
195
196 pub fn with_wait(mut self, enabled: bool) -> Self {
198 self.include_wait = enabled;
199 self
200 }
201
202 pub fn with_screenshot(mut self, enabled: bool) -> Self {
204 self.include_screenshot = enabled;
205 self
206 }
207
208 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 pub fn with_evaluate_js(mut self, enabled: bool) -> Self {
236 self.include_evaluate_js = enabled;
237 self
238 }
239
240 #[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 pub fn with_cookies(mut self, enabled: bool) -> Self {
270 self.include_cookies = enabled;
271 self
272 }
273
274 pub fn with_windows(mut self, enabled: bool) -> Self {
276 self.include_windows = enabled;
277 self
278 }
279
280 pub fn with_frames(mut self, enabled: bool) -> Self {
282 self.include_frames = enabled;
283 self
284 }
285
286 pub fn with_actions(mut self, enabled: bool) -> Self {
288 self.include_actions = enabled;
289 self
290 }
291
292 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 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 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
430pub 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
442pub 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 assert!(tools.len() > 40);
468
469 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 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 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 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 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 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}