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)]
44pub enum BrowserProfile {
45 Minimal,
48 FormFilling,
51 Scraping,
54 Full,
56}
57
58pub struct BrowserToolset {
63 resolver: SessionResolver,
64 include_navigation: bool,
66 include_interaction: bool,
68 include_extraction: bool,
70 include_wait: bool,
72 include_screenshot: bool,
74 include_js: bool,
76 include_cookies: bool,
78 include_windows: bool,
80 include_frames: bool,
82 include_actions: bool,
84}
85
86impl BrowserToolset {
87 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 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 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 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, 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 pub fn with_navigation(mut self, enabled: bool) -> Self {
245 self.include_navigation = enabled;
246 self
247 }
248
249 pub fn with_interaction(mut self, enabled: bool) -> Self {
251 self.include_interaction = enabled;
252 self
253 }
254
255 pub fn with_extraction(mut self, enabled: bool) -> Self {
257 self.include_extraction = enabled;
258 self
259 }
260
261 pub fn with_wait(mut self, enabled: bool) -> Self {
263 self.include_wait = enabled;
264 self
265 }
266
267 pub fn with_screenshot(mut self, enabled: bool) -> Self {
269 self.include_screenshot = enabled;
270 self
271 }
272
273 pub fn with_js(mut self, enabled: bool) -> Self {
275 self.include_js = enabled;
276 self
277 }
278
279 pub fn with_cookies(mut self, enabled: bool) -> Self {
281 self.include_cookies = enabled;
282 self
283 }
284
285 pub fn with_windows(mut self, enabled: bool) -> Self {
287 self.include_windows = enabled;
288 self
289 }
290
291 pub fn with_frames(mut self, enabled: bool) -> Self {
293 self.include_frames = enabled;
294 self
295 }
296
297 pub fn with_actions(mut self, enabled: bool) -> Self {
299 self.include_actions = enabled;
300 self
301 }
302
303 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 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 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
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}