cdp_html_shot/
tab.rs

1use crate::element::Element;
2use crate::transport::{Transport, TransportResponse, next_id};
3use crate::types::{CaptureOptions, ImageFormat, Viewport};
4use crate::utils::{self, send_and_get_msg};
5use anyhow::{Context, Result, anyhow};
6use serde_json::{Value, json};
7use std::sync::Arc;
8use std::time::Duration;
9use tokio::time;
10
11/// Represents a CDP browser tab (target) session.
12pub struct Tab {
13    pub(crate) transport: Arc<Transport>,
14    pub(crate) session_id: String,
15    pub(crate) target_id: String,
16}
17
18impl Tab {
19    pub(crate) async fn new(transport: Arc<Transport>) -> Result<Self> {
20        let TransportResponse::Response(res_create) = transport
21            .send(json!({ "id": next_id(), "method": "Target.createTarget", "params": { "url": "about:blank" } }))
22            .await? else { return Err(anyhow!("Invalid response type")); };
23
24        let target_id = res_create.result["targetId"]
25            .as_str()
26            .context("No targetId")?
27            .to_string();
28
29        let TransportResponse::Response(res_attach) = transport
30            .send(json!({ "id": next_id(), "method": "Target.attachToTarget", "params": { "targetId": target_id } }))
31            .await? else { return Err(anyhow!("Invalid response type")); };
32
33        let session_id = res_attach.result["sessionId"]
34            .as_str()
35            .context("No sessionId")?
36            .to_string();
37
38        Ok(Self {
39            transport,
40            session_id,
41            target_id,
42        })
43    }
44
45    pub(crate) async fn send_cmd(&self, method: &str, params: serde_json::Value) -> Result<Value> {
46        let msg_id = next_id();
47        let msg = json!({
48            "id": msg_id,
49            "method": method,
50            "params": params
51        })
52        .to_string();
53        let res = send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg).await?;
54        utils::serde_msg(&res)
55    }
56
57    pub async fn set_viewport(&self, viewport: &Viewport) -> Result<&Self> {
58        let screen_orientation = if viewport.is_landscape {
59            json!({"type": "landscapePrimary", "angle": 90})
60        } else {
61            json!({"type": "portraitPrimary", "angle": 0})
62        };
63
64        self.send_cmd(
65            "Emulation.setDeviceMetricsOverride",
66            json!({
67                "width": viewport.width,
68                "height": viewport.height,
69                "deviceScaleFactor": viewport.device_scale_factor,
70                "mobile": viewport.is_mobile,
71                "screenOrientation": screen_orientation
72            }),
73        )
74        .await?;
75
76        if viewport.has_touch {
77            self.send_cmd(
78                "Emulation.setTouchEmulationEnabled",
79                json!({
80                    "enabled": true,
81                    "maxTouchPoints": 5
82                }),
83            )
84            .await?;
85        }
86
87        Ok(self)
88    }
89
90    pub async fn clear_viewport(&self) -> Result<&Self> {
91        self.send_cmd("Emulation.clearDeviceMetricsOverride", json!({}))
92            .await?;
93        Ok(self)
94    }
95
96    pub async fn set_content(&self, content: &str) -> Result<&Self> {
97        self.send_cmd("Page.enable", json!({})).await?;
98
99        // Register listener BEFORE triggering the event to avoid race conditions
100        let event_rx = self
101            .transport
102            .listen_for_event(&self.session_id, "Page.loadEventFired")
103            .await?;
104
105        let js_write = format!(
106            r#"document.open(); document.write({}); document.close();"#,
107            serde_json::to_string(content)?
108        );
109
110        self.send_cmd(
111            "Runtime.evaluate",
112            json!({
113                "expression": js_write,
114                "awaitPromise": true
115            }),
116        )
117        .await?;
118
119        time::timeout(Duration::from_secs(30), event_rx)
120            .await
121            .map_err(|_| anyhow!("Timeout waiting for event Page.loadEventFired"))?
122            .map_err(|_| anyhow!("Event channel closed"))?;
123
124        Ok(self)
125    }
126
127    pub async fn evaluate(&self, expression: &str) -> Result<Value> {
128        let result = self
129            .send_cmd(
130                "Runtime.evaluate",
131                json!({
132                    "expression": expression,
133                    "returnByValue": true,
134                    "awaitPromise": true
135                }),
136            )
137            .await?;
138        Ok(result["result"]["result"]["value"].clone())
139    }
140
141    pub async fn evaluate_as_string(&self, expression: &str) -> Result<String> {
142        let value = self.evaluate(expression).await?;
143        value
144            .as_str()
145            .map(|s| s.to_string())
146            .or_else(|| Some(value.to_string()))
147            .context("Failed to convert result to string")
148    }
149
150    pub async fn find_element(&self, selector: &str) -> Result<Element<'_>> {
151        let msg_id = next_id();
152        let msg_doc =
153            json!({ "id": msg_id, "method": "DOM.getDocument", "params": {} }).to_string();
154        let res_doc =
155            send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg_doc).await?;
156        let data_doc = utils::serde_msg(&res_doc)?;
157        let root_node_id = data_doc["result"]["root"]["nodeId"]
158            .as_u64()
159            .context("No root node")?;
160
161        let msg_id = next_id();
162        let msg_sel = json!({
163            "id": msg_id,
164            "method": "DOM.querySelector",
165            "params": { "nodeId": root_node_id, "selector": selector }
166        })
167        .to_string();
168        let res_sel =
169            send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg_sel).await?;
170        let data_sel = utils::serde_msg(&res_sel)?;
171        let node_id = data_sel["result"]["nodeId"]
172            .as_u64()
173            .context("Element not found")?;
174
175        Element::new(self, node_id).await
176    }
177
178    pub async fn wait_for_selector(&self, selector: &str, timeout_ms: u64) -> Result<Element<'_>> {
179        let start = std::time::Instant::now();
180        let timeout = std::time::Duration::from_millis(timeout_ms);
181
182        loop {
183            match self.find_element(selector).await {
184                Ok(element) => return Ok(element),
185                Err(_) if start.elapsed() < timeout => {
186                    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
187                }
188                Err(e) => return Err(e),
189            }
190        }
191    }
192
193    pub async fn screenshot(&self, opts: CaptureOptions) -> Result<String> {
194        if let Some(ref viewport) = opts.viewport {
195            self.set_viewport(viewport).await?;
196        }
197
198        let mut params = json!({
199            "format": opts.format.as_str(),
200            "fromSurface": true,
201            "captureBeyondViewport": opts.full_page,
202        });
203
204        if matches!(opts.format, ImageFormat::Jpeg | ImageFormat::WebP) {
205            params["quality"] = json!(opts.quality.unwrap_or(90));
206        }
207
208        if let Some(ref clip) = opts.clip {
209            params["clip"] = json!({
210                "x": clip.x,
211                "y": clip.y,
212                "width": clip.width,
213                "height": clip.height,
214                "scale": clip.scale
215            });
216        }
217
218        if opts.omit_background && matches!(opts.format, ImageFormat::Png) {
219            self.send_cmd(
220                "Emulation.setDefaultBackgroundColorOverride",
221                json!({ "color": { "r": 0, "g": 0, "b": 0, "a": 0 } }),
222            )
223            .await?;
224        }
225
226        self.activate().await?;
227
228        let result = self.send_cmd("Page.captureScreenshot", params).await?;
229
230        if opts.omit_background && matches!(opts.format, ImageFormat::Png) {
231            let _ = self
232                .send_cmd("Emulation.setDefaultBackgroundColorOverride", json!({}))
233                .await;
234        }
235
236        result["result"]["data"]
237            .as_str()
238            .map(|s| s.to_string())
239            .context("No image data received")
240    }
241
242    pub async fn activate(&self) -> Result<&Self> {
243        let msg_id = next_id();
244        let msg = json!({ "id": msg_id, "method": "Target.activateTarget", "params": { "targetId": self.target_id } }).to_string();
245        send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg).await?;
246        Ok(self)
247    }
248
249    pub async fn goto(&self, url: &str) -> Result<&Self> {
250        self.send_cmd("Page.enable", json!({})).await?;
251
252        let event_rx = self
253            .transport
254            .listen_for_event(&self.session_id, "Page.loadEventFired")
255            .await?;
256
257        let msg_id = next_id();
258        let msg = json!({ "id": msg_id, "method": "Page.navigate", "params": { "url": url } })
259            .to_string();
260        send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg).await?;
261
262        time::timeout(Duration::from_secs(30), event_rx)
263            .await
264            .map_err(|_| anyhow!("Timeout waiting for event Page.loadEventFired"))?
265            .map_err(|_| anyhow!("Event channel closed"))?;
266
267        Ok(self)
268    }
269
270    pub async fn goto_no_wait(&self, url: &str) -> Result<&Self> {
271        let msg_id = next_id();
272        let msg = json!({ "id": msg_id, "method": "Page.navigate", "params": { "url": url } })
273            .to_string();
274        send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg).await?;
275        Ok(self)
276    }
277
278    pub async fn reload(&self) -> Result<&Self> {
279        self.send_cmd("Page.enable", json!({})).await?;
280
281        let event_rx = self
282            .transport
283            .listen_for_event(&self.session_id, "Page.loadEventFired")
284            .await?;
285
286        self.send_cmd("Page.reload", json!({})).await?;
287
288        time::timeout(Duration::from_secs(30), event_rx)
289            .await
290            .map_err(|_| anyhow!("Timeout waiting for event Page.loadEventFired"))?
291            .map_err(|_| anyhow!("Event channel closed"))?;
292
293        Ok(self)
294    }
295
296    pub async fn url(&self) -> Result<String> {
297        self.evaluate_as_string("window.location.href").await
298    }
299
300    pub async fn title(&self) -> Result<String> {
301        self.evaluate_as_string("document.title").await
302    }
303
304    pub fn session_id(&self) -> &str {
305        &self.session_id
306    }
307
308    pub fn target_id(&self) -> &str {
309        &self.target_id
310    }
311
312    pub async fn close(&self) -> Result<()> {
313        let msg_id = next_id();
314        let msg = json!({ "id": msg_id, "method": "Target.closeTarget", "params": { "targetId": self.target_id } }).to_string();
315        send_and_get_msg(self.transport.clone(), msg_id, &self.session_id, msg).await?;
316        Ok(())
317    }
318}