browsing 0.1.7

Browser automation: navigate, click, extract, screenshot. Standalone browser control via CDP.
Documentation
//! Web worker management via CDP Target domain

use crate::browser::views::WorkerInfo;
use crate::error::Result;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;

/// Manages web worker targets (web workers, service workers, shared workers)
#[derive(Debug, Clone)]
pub struct WorkerManager {
    client: Arc<crate::browser::cdp::CdpClient>,
    workers: Arc<Mutex<HashMap<String, WorkerInfo>>>,
}

impl WorkerManager {
    /// Create a new worker manager
    pub fn new(client: Arc<crate::browser::cdp::CdpClient>) -> Self {
        Self {
            client,
            workers: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Start discovering worker targets.
    /// Enables CDP Target domain and subscribes to targetCreated/targetDestroyed events.
    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(())
    }

    /// Stop discovering targets
    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(())
    }

    /// Attach to a worker target and return the session ID.
    /// After attaching, you can execute JS in the worker context.
    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()))?;

        // Update stored worker info
        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())
    }

    /// Detach from a worker target
    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(())
    }

    /// Execute JavaScript in a worker context.
    /// Requires `attach()` to have been called first for the target.
    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)
    }

    /// List all discovered workers
    pub async fn list_workers(&self) -> Vec<WorkerInfo> {
        self.workers.lock().await.values().cloned().collect()
    }

    /// Get a specific worker by target ID
    pub async fn get_worker(&self, target_id: &str) -> Option<WorkerInfo> {
        self.workers.lock().await.get(target_id).cloned()
    }

    /// Check if any workers are currently tracked
    pub async fn has_workers(&self) -> bool {
        !self.workers.lock().await.is_empty()
    }

    /// Get count of tracked workers
    pub async fn worker_count(&self) -> usize {
        self.workers.lock().await.len()
    }

    /// Clear tracked workers without detaching
    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"));
    }
}