adk_browser/
pool.rs

1//! Browser session pool for multi-user environments.
2//!
3//! Provides per-user session isolation by managing a pool of `BrowserSession`
4//! instances keyed by user ID. Sessions are created lazily on first access
5//! and can be released individually or cleaned up in bulk.
6
7use crate::config::BrowserConfig;
8use crate::session::BrowserSession;
9use adk_core::{AdkError, Result};
10use std::collections::HashMap;
11use std::sync::Arc;
12use tokio::sync::RwLock;
13
14/// A pool of browser sessions keyed by user ID.
15///
16/// Use this in multi-user agent platforms where each user needs an isolated
17/// browser instance. Sessions are created lazily via [`get_or_create`] and
18/// cleaned up via [`release`] or [`cleanup_all`].
19///
20/// # Example
21///
22/// ```rust,ignore
23/// use adk_browser::{BrowserConfig, BrowserSessionPool};
24///
25/// let pool = BrowserSessionPool::new(BrowserConfig::default(), 10);
26///
27/// // In a tool's execute(), resolve session from user context:
28/// let session = pool.get_or_create("user_123").await?;
29/// session.navigate("https://example.com").await?;
30///
31/// // On shutdown:
32/// pool.cleanup_all().await;
33/// ```
34pub struct BrowserSessionPool {
35    config: BrowserConfig,
36    sessions: RwLock<HashMap<String, Arc<BrowserSession>>>,
37    max_sessions: usize,
38}
39
40impl BrowserSessionPool {
41    /// Create a new session pool.
42    ///
43    /// `max_sessions` limits the number of concurrent browser sessions.
44    /// When the limit is reached, `get_or_create` will return an error.
45    pub fn new(config: BrowserConfig, max_sessions: usize) -> Self {
46        Self { config, sessions: RwLock::new(HashMap::new()), max_sessions }
47    }
48
49    /// Get an existing session for the user, or create a new one.
50    ///
51    /// The session is started automatically if newly created.
52    /// If the session exists but is stale, it will be reconnected.
53    pub async fn get_or_create(&self, user_id: &str) -> Result<Arc<BrowserSession>> {
54        // Fast path: check if session exists and is alive
55        {
56            let sessions = self.sessions.read().await;
57            if let Some(session) = sessions.get(user_id) {
58                if session.is_active().await {
59                    return Ok(session.clone());
60                }
61            }
62        }
63
64        // Slow path: create or replace session
65        let mut sessions = self.sessions.write().await;
66
67        // Double-check after acquiring write lock
68        if let Some(session) = sessions.get(user_id) {
69            if session.is_active().await {
70                return Ok(session.clone());
71            }
72            sessions.remove(user_id);
73        }
74
75        // Check capacity
76        if sessions.len() >= self.max_sessions {
77            return Err(AdkError::Tool(format!(
78                "Browser session pool full ({} sessions). Release unused sessions or increase max_sessions.",
79                self.max_sessions
80            )));
81        }
82
83        let session = Arc::new(BrowserSession::new(self.config.clone()));
84        session.start().await?;
85        sessions.insert(user_id.to_string(), session.clone());
86
87        Ok(session)
88    }
89
90    /// Release and stop a user's browser session.
91    pub async fn release(&self, user_id: &str) -> Result<()> {
92        let mut sessions = self.sessions.write().await;
93        if let Some(session) = sessions.remove(user_id) {
94            session.stop().await.ok(); // best-effort cleanup
95        }
96        Ok(())
97    }
98
99    /// Stop and remove all sessions. Call during graceful shutdown.
100    pub async fn cleanup_all(&self) {
101        let mut sessions = self.sessions.write().await;
102        for (_, session) in sessions.drain() {
103            session.stop().await.ok();
104        }
105    }
106
107    /// Number of active sessions in the pool.
108    pub async fn active_count(&self) -> usize {
109        self.sessions.read().await.len()
110    }
111
112    /// List all user IDs with active sessions.
113    pub async fn active_users(&self) -> Vec<String> {
114        self.sessions.read().await.keys().cloned().collect()
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    #[test]
123    fn test_pool_creation() {
124        let pool = BrowserSessionPool::new(BrowserConfig::default(), 5);
125        assert_eq!(pool.max_sessions, 5);
126    }
127
128    #[tokio::test]
129    async fn test_pool_active_count_starts_zero() {
130        let pool = BrowserSessionPool::new(BrowserConfig::default(), 5);
131        assert_eq!(pool.active_count().await, 0);
132    }
133
134    #[tokio::test]
135    async fn test_pool_active_users_starts_empty() {
136        let pool = BrowserSessionPool::new(BrowserConfig::default(), 5);
137        assert!(pool.active_users().await.is_empty());
138    }
139
140    #[tokio::test]
141    async fn test_pool_release_nonexistent_user() {
142        let pool = BrowserSessionPool::new(BrowserConfig::default(), 5);
143        let result = pool.release("nonexistent").await;
144        assert!(result.is_ok());
145    }
146
147    #[tokio::test]
148    async fn test_pool_cleanup_all_empty() {
149        let pool = BrowserSessionPool::new(BrowserConfig::default(), 5);
150        pool.cleanup_all().await;
151        assert_eq!(pool.active_count().await, 0);
152    }
153}
adk_browser/pool.rs

adk_browser/
pool.rs
