playwright_cdp/cdp/
connection.rs

1//! The multiplexing CDP connection.
2//!
3//! A single browser-level WebSocket carries traffic for the browser and every
4//! attached target (page), each tagged with `sessionId`. This struct owns the
5//! request/response correlation (by `id`) and the per-session event fan-out.
6
7use crate::cdp::messages::{CdpEvent, CdpRequest, Message};
8use crate::cdp::transport::{self, WebSocketWriter};
9use crate::{Error, Result};
10use parking_lot::Mutex;
11use serde_json::Value;
12use std::collections::HashMap;
13use std::sync::atomic::{AtomicU32, Ordering};
14use std::sync::Arc;
15use tokio::sync::{broadcast, mpsc, oneshot};
16
17/// Broadcast capacity for per-session event channels.
18const EVENT_CHANNEL_CAPACITY: usize = 128;
19
20/// Key into the event-channel map: `None` = browser-level, `Some(sid)` = target.
21type SessionKey = Option<String>;
22
23/// The single shared CDP connection over one WebSocket.
24///
25/// Cheap to share via `Arc<CdpConnection>`. A background dispatch task runs
26/// until the socket closes; when it does, all pending callbacks resolve to
27/// [`Error::ChannelClosed`].
28pub struct CdpConnection {
29    next_id: AtomicU32,
30    callbacks: Mutex<HashMap<u32, oneshot::Sender<Result<Value>>>>,
31    event_channels: Mutex<HashMap<SessionKey, broadcast::Sender<CdpEvent>>>,
32    writer: tokio::sync::Mutex<WebSocketWriter>,
33}
34
35impl CdpConnection {
36    /// Connect to a CDP WebSocket endpoint and start the dispatch task.
37    ///
38    /// Returns the shared connection plus a receiver for browser-level events
39    /// (`session_id == None`).
40    pub async fn connect(
41        ws_url: &str,
42        headers: Option<HashMap<String, String>>,
43    ) -> Result<(Arc<Self>, broadcast::Receiver<CdpEvent>)> {
44        let (writer, message_rx) = transport::connect(ws_url, headers).await?;
45
46        let conn = Arc::new(Self {
47            next_id: AtomicU32::new(1),
48            callbacks: Mutex::new(HashMap::new()),
49            event_channels: Mutex::new(HashMap::new()),
50            writer: tokio::sync::Mutex::new(writer),
51        });
52
53        let browser_rx = conn.subscribe(None);
54
55        let conn_for_loop = Arc::clone(&conn);
56        tokio::spawn(async move {
57            conn_for_loop.dispatch_loop(message_rx).await;
58        });
59
60        Ok((conn, browser_rx))
61    }
62
63    /// Send a CDP command and await its response (no timeout — wrap at a
64    /// higher layer if needed). `session_id` is `None` for browser-level.
65    pub async fn send_command(
66        self: &Arc<Self>,
67        session_id: Option<&str>,
68        method: &str,
69        params: Value,
70    ) -> Result<Value> {
71        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
72        let (tx, rx) = oneshot::channel::<Result<Value>>();
73        self.callbacks.lock().insert(id, tx);
74
75        let request = CdpRequest {
76            id,
77            method: method.to_string(),
78            params,
79            session_id: session_id.map(|s| s.to_string()),
80        };
81        let value = serde_json::to_value(&request)?;
82
83        // Serialize whole-message writes so frames never interleave.
84        let send_result = self.writer.lock().await.send(&value).await;
85        if let Err(e) = send_result {
86            self.callbacks.lock().remove(&id);
87            return Err(e);
88        }
89
90        match rx.await {
91            Ok(res) => res,
92            Err(_) => {
93                self.callbacks.lock().remove(&id);
94                Err(Error::ChannelClosed)
95            }
96        }
97    }
98
99    /// Subscribe to events for a session. `None` = browser-level. New
100    /// subscribers only receive events emitted *after* this call — subscribe
101    /// before issuing the command whose events you want to observe.
102    pub fn subscribe(&self, session_id: SessionKey) -> broadcast::Receiver<CdpEvent> {
103        self.event_channels
104            .lock()
105            .entry(session_id)
106            .or_insert_with(|| broadcast::channel(EVENT_CHANNEL_CAPACITY).0)
107            .subscribe()
108    }
109
110    /// Best-effort graceful close of the underlying socket.
111    pub async fn close(&self) -> Result<()> {
112        self.writer.lock().await.close().await
113    }
114
115    async fn dispatch_loop(self: Arc<Self>, mut message_rx: mpsc::Receiver<Value>) {
116        while let Some(v) = message_rx.recv().await {
117            match serde_json::from_value::<Message>(v) {
118                Ok(Message::Response(r)) => {
119                    let result = match r.error {
120                        Some(e) => Err(Error::CdpError {
121                            code: e.code,
122                            message: e.message,
123                            data: e.data,
124                        }),
125                        None => Ok(r.result.unwrap_or(Value::Null)),
126                    };
127                    if let Some(cb) = self.callbacks.lock().remove(&r.id) {
128                        let _ = cb.send(result);
129                    } else {
130                        tracing::trace!(id = r.id, "orphan CDP response");
131                    }
132                }
133                Ok(Message::Event(e)) => {
134                    let key: SessionKey = e.session_id.clone();
135                    let channel = self
136                        .event_channels
137                        .lock()
138                        .entry(key)
139                        .or_insert_with(|| broadcast::channel(EVENT_CHANNEL_CAPACITY).0)
140                        .clone();
141                    // broadcast::send returns Err if there are no receivers;
142                    // that's fine — events with no subscribers are dropped.
143                    let _ = channel.send(e);
144                }
145                Err(err) => {
146                    tracing::error!(error = %err, "failed to parse CDP message");
147                }
148            }
149        }
150
151        tracing::debug!("CDP transport closed; failing pending callbacks");
152        // Dropping the callbacks resolves their `rx` with a RecvError → ChannelClosed.
153        self.callbacks.lock().clear();
154    }
155}
156
157// Unit tests for the dispatch logic live behind an integration test that uses
158// a real (or loopback) socket — `CdpConnection` cannot be constructed in a pure
159// unit test because `WebSocketWriter` requires a live WebSocket.
playwright_cdp/cdp/connection.rs

playwright_cdp/cdp/
connection.rs
