use std::collections::HashMap; use std::path::{Component, Path, PathBuf}; use std::sync::Arc; use std::time::{Duration, Instant}; use anyhow::{Result, anyhow, bail}; use tokio::sync::Mutex; use uuid::Uuid; use super::action::BrowserAction; use super::runner::AgentBrowserRunner; use super::security::validate_navigation; use crate::bus::MediaRef; use crate::config::{BrowserConfig, expand_path}; use crate::tools::{ToolResult, ToolResultWithMedia}; struct BrowserSession { agent_browser_id: String, gate: Mutex<()>, last_used: std::sync::Mutex, } pub(super) struct BrowserManager { runner: AgentBrowserRunner, sessions: Mutex>>, max_sessions: usize, idle_timeout: Duration, artifact_dir: PathBuf, allow_private_hosts: bool, allowed_domains: Vec, } impl BrowserManager { pub(super) fn new(config: &BrowserConfig, workspace_dir: PathBuf) -> Result { if config.max_sessions == 0 { bail!("browser.max_sessions must be greater than zero"); } if config.command.trim().is_empty() { bail!("browser.command cannot be empty"); } let artifact_dir = expand_path(&config.artifact_dir); let artifact_dir = if artifact_dir.is_absolute() { artifact_dir } else { workspace_dir.join(artifact_dir) }; Ok(Self { runner: AgentBrowserRunner::new(config, workspace_dir), sessions: Mutex::new(HashMap::new()), max_sessions: config.max_sessions, idle_timeout: Duration::from_secs(config.idle_timeout_secs.max(1)), artifact_dir, allow_private_hosts: config.allow_private_hosts, allowed_domains: config.allowed_domains.clone(), }) } pub(super) async fn execute( &self, picobot_session_id: &str, action: BrowserAction, ) -> Result { if let BrowserAction::Open { url } = &action { validate_navigation(url, self.allow_private_hosts, &self.allowed_domains) .await .map_err(anyhow::Error::msg)?; } if action.is_close() { return self.close(picobot_session_id).await; } let screenshot_path = match action.screenshot_filename() { Some(filename) => Some(self.prepare_screenshot_path(filename).await?), None => None, }; let (session, stale) = self.session_for(picobot_session_id).await?; for stale_session in stale { let _ = self .runner .run(&stale_session, &["close".to_string()]) .await; } let _gate = session.gate.lock().await; let path_string = screenshot_path .as_ref() .map(|path| path.to_string_lossy().into_owned()); let commands = action.commands(path_string.as_deref()); let mut last_response = None; for command in commands { last_response = Some(self.runner.run(&session.agent_browser_id, &command).await?); } *session .last_used .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) = Instant::now(); let response = last_response.ok_or_else(|| anyhow!("browser action produced no command"))?; let mut output = self.runner.render_response(&response); let mut media_refs = Vec::new(); if let Some(path) = screenshot_path { let metadata = tokio::fs::metadata(&path) .await .map_err(|error| anyhow!("agent-browser did not create screenshot: {error}"))?; if !metadata.is_file() || metadata.len() == 0 { bail!("agent-browser created an empty screenshot"); } let canonical = tokio::fs::canonicalize(&path).await.unwrap_or(path); let canonical = canonical.to_string_lossy().into_owned(); output = format!("Screenshot saved: {canonical}\n{output}"); media_refs.push(MediaRef { path: canonical, media_type: "image".to_string(), }); } Ok(ToolResultWithMedia { result: ToolResult { success: true, output, error: None, }, media_refs, }) } async fn session_for( &self, picobot_session_id: &str, ) -> Result<(Arc, Vec)> { let now = Instant::now(); let mut sessions = self.sessions.lock().await; if let Some(session) = sessions.get(picobot_session_id) { *session .last_used .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) = now; return Ok((session.clone(), Vec::new())); } let mut stale_ids = Vec::new(); sessions.retain(|_, session| { let last_used = *session .last_used .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); let idle = now.duration_since(last_used) >= self.idle_timeout; let keep = !idle || session.gate.try_lock().is_err(); if !keep { stale_ids.push(session.agent_browser_id.clone()); } keep }); if sessions.len() >= self.max_sessions { bail!( "browser session limit reached ({}); close another dialog browser or wait for idle cleanup", self.max_sessions ); } let session = Arc::new(BrowserSession { agent_browser_id: format!("picobot-{}", Uuid::new_v4().simple()), gate: Mutex::new(()), last_used: std::sync::Mutex::new(now), }); sessions.insert(picobot_session_id.to_string(), session.clone()); Ok((session, stale_ids)) } async fn close(&self, picobot_session_id: &str) -> Result { let session = self.sessions.lock().await.remove(picobot_session_id); let Some(session) = session else { return Ok(ToolResult { success: true, output: "Browser session is already closed.".to_string(), error: None, } .into()); }; let _gate = session.gate.lock().await; let response = self .runner .run(&session.agent_browser_id, &["close".to_string()]) .await?; Ok(ToolResult { success: true, output: self.runner.render_response(&response), error: None, } .into()) } async fn prepare_screenshot_path(&self, requested: Option<&str>) -> Result { tokio::fs::create_dir_all(&self.artifact_dir).await?; let filename = match requested { Some(requested) => { let path = Path::new(requested); if path.is_absolute() || path .components() .any(|component| !matches!(component, Component::Normal(_))) { bail!("screenshot path must be a filename without directory components"); } let filename = path .file_name() .and_then(|name| name.to_str()) .ok_or_else(|| anyhow!("invalid screenshot filename"))?; if !filename.to_ascii_lowercase().ends_with(".png") { bail!("screenshot filename must end in .png"); } filename.to_string() } None => format!( "picobot-screenshot-{}-{}.png", chrono::Utc::now().format("%Y%m%dT%H%M%S"), &Uuid::new_v4().simple().to_string()[..8] ), }; Ok(self.artifact_dir.join(filename)) } }