browser_control/session/
attach.rs

1//! Attach to a page target and expose engine-agnostic high-level operations.
2//!
3//! [`PageSession`] hides the CDP/BiDi split behind a single async API
4//! (`evaluate`, `navigate`, `screenshot`). The CLI subcommands instantiate
5//! a fresh session per call; the MCP server may pre-build a session backed
6//! by a long-lived BiDi client via [`PageSession::from_bidi_cache`].
7
8use std::sync::Arc;
9
10use anyhow::{anyhow, Result};
11use regex::Regex;
12use serde_json::{json, Value};
13
14use crate::bidi::BidiClient;
15use crate::cdp::CdpClient;
16use crate::detect::Engine;
17use crate::session::targets::{open_bidi, open_cdp};
18
19/// A bound page-level session. Variants are not constructed directly outside
20/// this module; use [`PageSession::attach`].
21pub enum PageSession {
22    Cdp(CdpPage),
23    /// A BiDi page session. The client is shared via `Arc` so the MCP server
24    /// can keep a single persistent BiDi session across many tool calls
25    /// (Firefox limits a browser to one BiDi session at a time).
26    Bidi(BidiPage),
27}
28
29pub struct CdpPage {
30    pub client: CdpClient,
31    pub session_id: String,
32    pub target_id: String,
33}
34
35pub struct BidiPage {
36    pub client: Arc<BidiClient>,
37    pub context: String,
38}
39
40impl PageSession {
41    /// Attach to a fresh page session over `engine`.
42    ///
43    /// If `url_regex` is `Some`, the first page target whose URL matches is
44    /// selected; otherwise the first page (or top-level browsing context) is
45    /// used.
46    pub async fn attach(endpoint: &str, engine: Engine, url_regex: Option<&str>) -> Result<Self> {
47        let pattern = url_regex.map(Regex::new).transpose()?;
48        match engine {
49            Engine::Cdp => {
50                let client = open_cdp(endpoint).await?;
51                let target_id = pick_cdp_page(&client, pattern.as_ref()).await?;
52                let session_id = client.attach_to_target(&target_id).await?;
53                Ok(PageSession::Cdp(CdpPage {
54                    client,
55                    session_id,
56                    target_id,
57                }))
58            }
59            Engine::Bidi => {
60                let client = Arc::new(open_bidi(endpoint).await?);
61                client.session_new().await?;
62                let context = pick_bidi_context(&client, pattern.as_ref()).await?;
63                Ok(PageSession::Bidi(BidiPage { client, context }))
64            }
65        }
66    }
67
68    /// Build a BiDi session from a pre-opened, possibly cached client.
69    ///
70    /// The MCP server uses this to share one BiDi client across tool calls;
71    /// `session.new` is invoked only when the client was freshly opened (the
72    /// caller is expected to have done so).
73    pub async fn from_bidi_cache(client: Arc<BidiClient>, url_regex: Option<&str>) -> Result<Self> {
74        let pattern = url_regex.map(Regex::new).transpose()?;
75        let context = pick_bidi_context(&client, pattern.as_ref()).await?;
76        Ok(PageSession::Bidi(BidiPage { client, context }))
77    }
78
79    /// Evaluate `expression` in the page's main world.
80    ///
81    /// `await_promise = true` mirrors `Runtime.evaluate({awaitPromise:true})`
82    /// and is appropriate for fetch / promise-returning code. The returned
83    /// value is the raw `result.value` from CDP / BiDi after `returnByValue`.
84    pub async fn evaluate(&self, expression: &str, await_promise: bool) -> Result<Value> {
85        match self {
86            PageSession::Cdp(p) => {
87                let v = p
88                    .client
89                    .send_with_session(
90                        "Runtime.evaluate",
91                        json!({
92                            "expression": expression,
93                            "returnByValue": true,
94                            "awaitPromise": await_promise,
95                        }),
96                        Some(&p.session_id),
97                    )
98                    .await?;
99                Ok(v["result"]["value"].clone())
100            }
101            PageSession::Bidi(p) => {
102                let _ = await_promise; // BiDi always awaits per script_evaluate
103                let v = p.client.script_evaluate(&p.context, expression).await?;
104                Ok(v["result"]["value"].clone())
105            }
106        }
107    }
108
109    /// Navigate the current page to `url`.
110    pub async fn navigate(&self, url: &str) -> Result<()> {
111        match self {
112            PageSession::Cdp(p) => {
113                p.client
114                    .send_with_session("Page.navigate", json!({"url": url}), Some(&p.session_id))
115                    .await?;
116                Ok(())
117            }
118            PageSession::Bidi(p) => {
119                p.client.browsing_context_navigate(&p.context, url).await?;
120                Ok(())
121            }
122        }
123    }
124
125    /// Capture a PNG screenshot of the current page; returns base64 data.
126    pub async fn screenshot(&self, full_page: bool) -> Result<String> {
127        match self {
128            PageSession::Cdp(p) => {
129                let v = p
130                    .client
131                    .send_with_session(
132                        "Page.captureScreenshot",
133                        json!({
134                            "format": "png",
135                            "captureBeyondViewport": full_page,
136                        }),
137                        Some(&p.session_id),
138                    )
139                    .await?;
140                v["data"]
141                    .as_str()
142                    .map(|s| s.to_string())
143                    .ok_or_else(|| anyhow!("no screenshot data"))
144            }
145            PageSession::Bidi(p) => {
146                let _ = full_page; // BiDi captures the viewport by default
147                p.client
148                    .browsing_context_capture_screenshot(&p.context)
149                    .await
150            }
151        }
152    }
153
154    /// Engine this session is bound to.
155    pub fn engine(&self) -> Engine {
156        match self {
157            PageSession::Cdp(_) => Engine::Cdp,
158            PageSession::Bidi(_) => Engine::Bidi,
159        }
160    }
161
162    /// Release the underlying CDP connection (no-op for BiDi, whose client
163    /// is shared via `Arc`).
164    pub async fn close(self) {
165        match self {
166            PageSession::Cdp(p) => p.client.close().await,
167            PageSession::Bidi(_) => {}
168        }
169    }
170}
171
172async fn pick_cdp_page(client: &CdpClient, pattern: Option<&Regex>) -> Result<String> {
173    let targets = client.list_targets().await?;
174    let mut pages = targets
175        .iter()
176        .filter(|t| t.get("type").and_then(|v| v.as_str()) == Some("page"));
177    let pick = if let Some(re) = pattern {
178        pages
179            .find(|t| {
180                t.get("url")
181                    .and_then(|v| v.as_str())
182                    .is_some_and(|u| re.is_match(u))
183            })
184            .ok_or_else(|| anyhow!("no CDP page target matched URL regex"))?
185    } else {
186        pages
187            .next()
188            .ok_or_else(|| anyhow!("no page target found"))?
189    };
190    pick.get("targetId")
191        .and_then(|v| v.as_str())
192        .map(|s| s.to_string())
193        .ok_or_else(|| anyhow!("targetId missing from page target"))
194}
195
196async fn pick_bidi_context(client: &BidiClient, pattern: Option<&Regex>) -> Result<String> {
197    let tree = client.send("browsingContext.getTree", json!({})).await?;
198    let contexts = tree
199        .get("contexts")
200        .and_then(|v| v.as_array())
201        .ok_or_else(|| anyhow!("no contexts in browsingContext.getTree"))?;
202    if let Some(re) = pattern {
203        for c in contexts {
204            let url = c.get("url").and_then(|v| v.as_str()).unwrap_or("");
205            if re.is_match(url) {
206                return c
207                    .get("context")
208                    .and_then(|v| v.as_str())
209                    .map(|s| s.to_string())
210                    .ok_or_else(|| anyhow!("no context id"));
211            }
212        }
213        Err(anyhow!("no BiDi context matched URL regex"))
214    } else {
215        contexts
216            .first()
217            .and_then(|c| c.get("context").and_then(|v| v.as_str()))
218            .map(|s| s.to_string())
219            .ok_or_else(|| anyhow!("no top-level browsing context"))
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226    use futures_util::{SinkExt, StreamExt};
227    use tokio_tungstenite::tungstenite::Message;
228
229    async fn spawn_cdp_mock(targets: Vec<Value>) -> String {
230        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
231        let addr = listener.local_addr().unwrap();
232        tokio::spawn(async move {
233            let (stream, _) = listener.accept().await.unwrap();
234            let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap();
235            while let Some(Ok(Message::Text(t))) = ws.next().await {
236                let req: Value = serde_json::from_str(&t).unwrap();
237                let id = req["id"].as_u64().unwrap();
238                let method = req["method"].as_str().unwrap_or("");
239                let result = match method {
240                    "Target.getTargets" => json!({"targetInfos": targets.clone()}),
241                    "Target.attachToTarget" => json!({"sessionId": "S1"}),
242                    "Runtime.evaluate" => json!({"result": {"value": "ok"}}),
243                    "Page.navigate" => json!({}),
244                    "Page.captureScreenshot" => json!({"data": "PNGDATA"}),
245                    _ => json!({}),
246                };
247                let resp = json!({"id": id, "result": result});
248                ws.send(Message::Text(resp.to_string())).await.unwrap();
249            }
250        });
251        format!("ws://{addr}")
252    }
253
254    #[tokio::test]
255    async fn attach_cdp_picks_first_page_when_no_regex() {
256        let url = spawn_cdp_mock(vec![
257            json!({"targetId":"a","type":"page","url":"https://example.com/"}),
258            json!({"targetId":"b","type":"page","url":"https://other.test/"}),
259        ])
260        .await;
261        let s = PageSession::attach(&url, Engine::Cdp, None).await.unwrap();
262        match s {
263            PageSession::Cdp(p) => {
264                assert_eq!(p.target_id, "a");
265                assert_eq!(p.session_id, "S1");
266            }
267            _ => panic!("expected CDP"),
268        }
269    }
270
271    #[tokio::test]
272    async fn attach_cdp_url_regex_selects_matching() {
273        let url = spawn_cdp_mock(vec![
274            json!({"targetId":"a","type":"page","url":"https://example.com/"}),
275            json!({"targetId":"b","type":"page","url":"https://other.test/"}),
276        ])
277        .await;
278        let s = PageSession::attach(&url, Engine::Cdp, Some(r"other"))
279            .await
280            .unwrap();
281        match s {
282            PageSession::Cdp(p) => assert_eq!(p.target_id, "b"),
283            _ => panic!("expected CDP"),
284        }
285    }
286
287    #[tokio::test]
288    async fn attach_cdp_url_regex_no_match_errors() {
289        let url = spawn_cdp_mock(vec![
290            json!({"targetId":"a","type":"page","url":"https://example.com/"}),
291        ])
292        .await;
293        let err = match PageSession::attach(&url, Engine::Cdp, Some("nomatch")).await {
294            Ok(_) => panic!("expected error"),
295            Err(e) => e,
296        };
297        assert!(err.to_string().contains("no CDP page target matched"));
298    }
299
300    #[tokio::test]
301    async fn evaluate_round_trip_cdp() {
302        let url = spawn_cdp_mock(vec![
303            json!({"targetId":"a","type":"page","url":"https://example.com/"}),
304        ])
305        .await;
306        let s = PageSession::attach(&url, Engine::Cdp, None).await.unwrap();
307        let v = s.evaluate("1+1", false).await.unwrap();
308        assert_eq!(v, json!("ok"));
309        s.close().await;
310    }
311
312    #[tokio::test]
313    async fn screenshot_round_trip_cdp() {
314        let url = spawn_cdp_mock(vec![
315            json!({"targetId":"a","type":"page","url":"https://example.com/"}),
316        ])
317        .await;
318        let s = PageSession::attach(&url, Engine::Cdp, None).await.unwrap();
319        let b64 = s.screenshot(false).await.unwrap();
320        assert_eq!(b64, "PNGDATA");
321        s.close().await;
322    }
323}
browser_control/session/attach.rs

browser_control/session/
attach.rs
