use crate::browser::views::WorkerInfo;
use crate::error::Result;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
#[derive(Debug, Clone)]
pub struct WorkerManager {
client: Arc<crate::browser::cdp::CdpClient>,
workers: Arc<Mutex<HashMap<String, WorkerInfo>>>,
}
impl WorkerManager {
pub fn new(client: Arc<crate::browser::cdp::CdpClient>) -> Self {
Self {
client,
workers: Arc::new(Mutex::new(HashMap::new())),
}
}
pub async fn start_discovery(&self) -> Result<()> {
self.client
.send_command("Target.setDiscoverTargets", serde_json::json!({ "discover": true }))
.await?;
let workers = Arc::clone(&self.workers);
let mut rx = self.client.subscribe_events().await?;
tokio::spawn(async move {
while let Ok(event) = rx.recv().await {
match event.method.as_str() {
"Target.targetCreated" => {
if let Some(target_info) = event.params.get("targetInfo") {
let target_type = target_info
.get("type")
.and_then(|v| v.as_str())
.unwrap_or("");
if target_type == "worker" || target_type == "service_worker" {
let target_id = target_info
.get("targetId")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let url = target_info
.get("url")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let info = WorkerInfo {
target_id: target_id.clone(),
url,
worker_type: target_type.to_string(),
session_id: None,
};
workers.lock().await.insert(target_id, info);
}
}
}
"Target.targetDestroyed" => {
if let Some(target_id) = event.params.get("targetId").and_then(|v| v.as_str()) {
workers.lock().await.remove(target_id);
}
}
_ => {}
}
}
});
Ok(())
}
pub async fn stop_discovery(&self) -> Result<()> {
self.client
.send_command("Target.setDiscoverTargets", serde_json::json!({ "discover": false }))
.await?;
self.workers.lock().await.clear();
Ok(())
}
pub async fn attach(&self, target_id: &str) -> Result<String> {
let result = self
.client
.send_command(
"Target.attachToTarget",
serde_json::json!({
"targetId": target_id,
"flatten": true,
}),
)
.await?;
let session_id = result
.get("sessionId")
.and_then(|v| v.as_str())
.ok_or_else(|| crate::error::BrowsingError::Browser("No sessionId in attach response".to_string()))?;
let mut guard = self.workers.lock().await;
if let Some(info) = guard.get_mut(target_id) {
info.session_id = Some(session_id.to_string());
}
Ok(session_id.to_string())
}
pub async fn detach(&self, target_id: &str) -> Result<()> {
let _ = self
.client
.send_command(
"Target.detachFromTarget",
serde_json::json!({ "targetId": target_id }),
)
.await;
let mut guard = self.workers.lock().await;
if let Some(info) = guard.get_mut(target_id) {
info.session_id = None;
}
Ok(())
}
pub async fn evaluate_in_worker(
&self,
target_id: &str,
expression: &str,
) -> Result<serde_json::Value> {
let session_id = {
let guard = self.workers.lock().await;
guard
.get(target_id)
.and_then(|w| w.session_id.clone())
.ok_or_else(|| crate::error::BrowsingError::Browser("Worker not attached".to_string()))?
};
let result = self
.client
.send_command_with_session(
"Runtime.evaluate",
serde_json::json!({
"expression": expression,
"returnByValue": true,
}),
Some(&session_id),
)
.await?;
Ok(result)
}
pub async fn list_workers(&self) -> Vec<WorkerInfo> {
self.workers.lock().await.values().cloned().collect()
}
pub async fn get_worker(&self, target_id: &str) -> Option<WorkerInfo> {
self.workers.lock().await.get(target_id).cloned()
}
pub async fn has_workers(&self) -> bool {
!self.workers.lock().await.is_empty()
}
pub async fn worker_count(&self) -> usize {
self.workers.lock().await.len()
}
pub async fn clear(&self) {
self.workers.lock().await.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_worker_info_creation() {
let info = WorkerInfo {
target_id: "worker-1".to_string(),
url: "https://example.com/worker.js".to_string(),
worker_type: "worker".to_string(),
session_id: Some("session-1".to_string()),
};
assert_eq!(info.target_id, "worker-1");
assert_eq!(info.url, "https://example.com/worker.js");
}
#[test]
fn test_worker_info_serialization() {
let info = WorkerInfo {
target_id: "w-1".to_string(),
url: "worker.js".to_string(),
worker_type: "service_worker".to_string(),
session_id: None,
};
let json = serde_json::to_string(&info).unwrap();
assert!(json.contains("w-1"));
assert!(json.contains("service_worker"));
}
}