use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::path::Path;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};
use uuid::Uuid;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BrowserSession {
pub id: String,
pub browser_type: BrowserType,
pub agent_id: String,
pub status: SessionStatus,
pub created_at: chrono::DateTime<chrono::Utc>,
pub last_activity: chrono::DateTime<chrono::Utc>,
pub url: Option<String>,
pub title: Option<String>,
pub headless_chrome_handle: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum BrowserType {
Chrome,
Firefox,
Safari,
Edge,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum SessionStatus {
Starting,
Running,
Idle,
Error,
Closed,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BrowserCommand {
pub action: BrowserAction,
pub parameters: HashMap<String, serde_json::Value>,
pub timeout_ms: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum BrowserAction {
Navigate { url: String },
Click { selector: String },
Type { selector: String, text: String },
Screenshot { path: Option<String> },
ExecuteScript { script: String },
WaitForElement { selector: String, timeout_ms: u32 },
GetTitle {},
GetUrl {},
Refresh {},
Back {},
Forward {},
GetText { selector: String },
GetAttribute { selector: String, attribute: String },
ScrollTo { selector: Option<String> },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CommandResult {
pub success: bool,
pub data: serde_json::Value,
pub error: Option<String>,
pub execution_time_ms: u64,
}
pub struct BrowserManager {
sessions: Arc<RwLock<HashMap<String, BrowserSession>>>,
}
impl Default for BrowserManager {
fn default() -> Self {
Self::new()
}
}
impl BrowserManager {
pub fn new() -> Self {
Self {
sessions: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn init_browsers(&self) -> Result<(), anyhow::Error> {
info!("Initializing headless_chrome browsers...");
info!("Headless_chrome browsers initialized successfully");
Ok(())
}
pub async fn create_session(
&self,
agent_id: String,
browser_type: BrowserType,
) -> Result<String, anyhow::Error> {
let session_id = Uuid::new_v4().to_string();
info!(
"Creating browser session {} for agent {} with {:?}",
session_id, agent_id, browser_type
);
let session = BrowserSession {
id: session_id.clone(),
browser_type,
agent_id,
status: SessionStatus::Starting,
created_at: chrono::Utc::now(),
last_activity: chrono::Utc::now(),
url: None,
title: None,
headless_chrome_handle: Some(format!("headless_chrome_{}", session_id)),
};
{
let mut sessions = self.sessions.write();
sessions.insert(session_id.clone(), session);
}
info!("Browser session {} created successfully", session_id);
Ok(session_id)
}
pub async fn get_session(&self, session_id: &str) -> Option<BrowserSession> {
let sessions = self.sessions.read().clone();
sessions.get(session_id).cloned()
}
pub async fn list_sessions(&self, agent_id: Option<&str>) -> Vec<BrowserSession> {
let sessions = self.sessions.read().clone();
sessions
.values()
.filter(|session| agent_id.is_none_or(|id| session.agent_id == id))
.cloned()
.collect()
}
pub async fn execute_command(
&self,
session_id: &str,
command: BrowserCommand,
) -> Result<CommandResult, anyhow::Error> {
let start_time = std::time::Instant::now();
let sessions = self.sessions.read().clone();
if let Some(session) = sessions.get(session_id) {
if session.status != SessionStatus::Running {
return Err(anyhow::anyhow!("Session {} is not running", session_id));
}
} else {
return Err(anyhow::anyhow!("Session {} not found", session_id));
}
debug!(
"Executing command {:?} on session {}",
command.action, session_id
);
let result = match &command.action {
BrowserAction::Screenshot { path } => {
let screenshot_path = path.unwrap_or_else(|| format!("screenshot_{}.png", session_id));
Ok(serde_json::json!({
"action": "screenshot",
"path": screenshot_path,
"result": "success"
}))
}
BrowserAction::Navigate { url } => {
Ok(serde_json::json!({
"url": url,
"title": "Mock Title",
"status": "loaded"
}))
}
_ => {
Ok(serde_json::json!({
"action": format!("{:?}", command.action),
"result": "success"
}))
}
};
let execution_time = start_time.elapsed().as_millis() as u64;
{
let mut sessions = self.sessions.write();
if let Some(session) = sessions.get_mut(session_id) {
session.last_activity = chrono::Utc::now();
if let BrowserAction::Navigate { url } = &command.action {
session.url = Some(url.clone());
session.title = Some("Mock Title".to_string());
}
}
}
let command_result = CommandResult {
success: true,
data: result,
error: None,
execution_time_ms: execution_time,
};
Ok(command_result)
}
pub async fn close_session(&self, session_id: &str) -> Result<(), anyhow::Error> {
info!("Closing browser session: {}", session_id);
{
let mut sessions = self.sessions.write();
sessions.remove(session_id);
}
info!("Browser session {} closed successfully", session_id);
Ok(())
}
async fn update_session_activity(&self, session_id: &str) {
let mut sessions = self.sessions.write();
if let Some(session) = sessions.get_mut(session_id) {
session.last_activity = chrono::Utc::now();
}
}
async fn cleanup_idle_sessions(
&self,
idle_timeout_minutes: u64,
) -> Result<usize, anyhow::Error> {
let now = chrono::Utc::now();
let timeout = chrono::Duration::minutes(idle_timeout_minutes as i64);
let mut sessions_to_close: Vec<String> = Vec::new();
{
let sessions = self.sessions.read().clone();
for (session_id, session) in sessions.iter() {
let duration_since_activity = now.signed_duration_since(session.last_activity);
if duration_since_activity > timeout && session.status == SessionStatus::Running
{
sessions_to_close.push(session_id.clone());
}
}
}
for session_id in &sessions_to_close {
if let Err(e) = self.close_session(session_id).await {
error!("Failed to close idle session {}: {}", session_id, e);
} else {
info!("Closed idle session: {}", session_id);
}
}
Ok(sessions_to_close.len())
}
pub async fn get_session_stats(&self) -> Result<SessionStats, anyhow::Error> {
let sessions = self.sessions.read().clone();
let mut browser_type_counts = HashMap::new();
let mut running_count = 0;
let mut idle_count = 0;
let mut error_count = 0;
for session in sessions.values() {
let browser_type_str = format!("{:?}", session.browser_type);
*browser_type_counts.entry(browser_type_str).or_insert(0) += 1;
match session.status {
SessionStatus::Running => running_count += 1,
SessionStatus::Idle => idle_count += 1,
SessionStatus::Error => error_count += 1,
_ => {}
}
}
Ok(SessionStats {
total_sessions: sessions.len(),
running_sessions: running_count,
idle_sessions: idle_count,
error_sessions: error_count,
browser_type_counts,
})
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionStats {
pub total_sessions: usize,
pub running_sessions: usize,
pub idle_sessions: usize,
pub error_sessions: usize,
pub browser_type_counts: HashMap<String, usize>,
}