Speed up snapshots of large pages

cli/src/native/actions.rs

@@ -11238,7 +11238,7 @@ async fn handle_getbyrole(cmd: &Value, state: &mut DaemonState) -> Result<Value,
 
     let ax_tree: GetFullAXTreeResult = mgr
         .client
-        .send_command_typed(
+        .send_command_typed_raw(
             "Accessibility.getFullAXTree",
             &ax_params,
             Some(effective_session_id),

cli/src/native/cdp/client.rs

@@ -20,6 +20,12 @@ use tokio_tungstenite::Connector;
 use super::types::{CdpCommand, CdpError, CdpEvent, CdpMessage};
 
 type PendingMap = Arc<Mutex<HashMap<u64, oneshot::Sender<CdpMessage>>>>;
+
+/// Responses whose caller wants the raw JSON text rather than a parsed
+/// `serde_json::Value`. The accessibility tree of a large page is several
+/// megabytes of JSON, and building a `Value` from it only to deserialize it
+/// again costs more than the command itself, so those callers parse once.
+type RawPendingMap = Arc<Mutex<HashMap<u64, oneshot::Sender<String>>>>;
 type WsTx = Arc<
     Mutex<
         futures_util::stream::SplitSink<
@@ -155,6 +161,34 @@ fn extract_command_id(msg: &str) -> Option<u64> {
     envelope.id.filter(|id| *id > 0)
 }
 
+/// Response envelope for [`CdpClient::send_command_typed_raw`]. Unknown fields,
+/// `sessionId` among them, are ignored.
+#[derive(Deserialize)]
+struct RawCommandResponse<T> {
+    result: Option<T>,
+    error: Option<CdpError>,
+}
+
+/// Read the command id from a CDP response without parsing the message.
+///
+/// Responses start with `{"id":<number>,`, so the id is available from the
+/// first bytes even when the payload is megabytes of JSON. Returns `None` for
+/// anything that is not shaped like a response envelope, which sends the
+/// message down the normal parsed path.
+fn leading_response_id(msg: &str) -> Option<u64> {
+    let rest = msg.strip_prefix("{\"id\":")?;
+    let end = rest
+        .find(|c: char| !c.is_ascii_digit())
+        .unwrap_or(rest.len());
+    if end == 0 || end > 20 {
+        return None;
+    }
+    if !matches!(rest.as_bytes().get(end), None | Some(b',') | Some(b'}')) {
+        return None;
+    }
+    rest[..end].parse::<u64>().ok().filter(|id| *id > 0)
+}
+
 pub struct CdpClient {
     ws_tx: WsTx,
     socket: socket2::Socket,
@@ -164,6 +198,7 @@ pub struct CdpClient {
     cancel_tx: watch::Sender<bool>,
     next_id: AtomicU64,
     pending: PendingMap,
+    raw_pending: RawPendingMap,
     event_tx: broadcast::Sender<CdpEvent>,
     raw_tx: broadcast::Sender<RawCdpMessage>,
     reader_handle: Mutex<Option<JoinHandle<()>>>,
@@ -178,6 +213,7 @@ pub struct CdpClient {
 /// it via `done`.
 struct PendingGuard {
     pending: PendingMap,
+    raw_pending: RawPendingMap,
     id: u64,
     done: bool,
 }
@@ -188,10 +224,12 @@ impl Drop for PendingGuard {
             return;
         }
         let pending = self.pending.clone();
+        let raw_pending = self.raw_pending.clone();
         let id = self.id;
         if let Ok(handle) = tokio::runtime::Handle::try_current() {
             handle.spawn(async move {
                 pending.lock().await.remove(&id);
+                raw_pending.lock().await.remove(&id);
             });
         }
     }
@@ -258,6 +296,7 @@ impl CdpClient {
         let ws_tx = Arc::new(Mutex::new(ws_tx));
 
         let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
+        let raw_pending: RawPendingMap = Arc::new(Mutex::new(HashMap::new()));
         let (event_tx, _) = broadcast::channel(4096);
         let (raw_tx, _) = broadcast::channel(4096);
         let closed = Arc::new(AtomicBool::new(false));
@@ -266,6 +305,7 @@ impl CdpClient {
         let method_taps: MethodTaps = Arc::new(std::sync::Mutex::new(Vec::new()));
 
         let pending_clone = pending.clone();
+        let raw_pending_clone = raw_pending.clone();
         let event_tx_clone = event_tx.clone();
         let raw_tx_clone = raw_tx.clone();
         let closed_reader = closed.clone();
@@ -319,6 +359,17 @@ impl CdpClient {
                     });
                 }
 
+                // Responses registered for raw delivery skip the `Value` parse
+                // entirely: the caller deserializes the payload once, straight
+                // from this text (large accessibility trees).
+                if let Some(id) = leading_response_id(&msg) {
+                    let waiter = raw_pending_clone.lock().await.remove(&id);
+                    if let Some(tx) = waiter {
+                        let _ = tx.send(msg);
+                        continue;
+                    }
+                }
+
                 let parsed: CdpMessage = match serde_json::from_str(&msg) {
                     Ok(parsed) => parsed,
                     Err(_) => {
@@ -328,6 +379,13 @@ impl CdpClient {
                             Ok(parsed) => parsed,
                             Err(error) => {
                                 if let Some(id) = extract_command_id(candidate) {
+                                    // Hand the unparsable text to a raw waiter
+                                    // so it fails fast instead of timing out.
+                                    let raw_waiter = raw_pending_clone.lock().await.remove(&id);
+                                    if let Some(tx) = raw_waiter {
+                                        let _ = tx.send(candidate.to_string());
+                                        continue;
+                                    }
                                     let waiter = pending_clone.lock().await.remove(&id);
                                     if let Some(tx) = waiter {
                                         let _ = tx.send(CdpMessage {
@@ -354,6 +412,13 @@ impl CdpClient {
                 };
 
                 if let Some(id) = parsed.id {
+                    // Responses from servers that do not lead with the id still
+                    // reach a raw waiter, at the cost of the parse above.
+                    let raw_waiter = raw_pending_clone.lock().await.remove(&id);
+                    if let Some(tx) = raw_waiter {
+                        let _ = tx.send(msg);
+                        continue;
+                    }
                     // Response to a command
                     let mut pending = pending_clone.lock().await;
                     if let Some(tx) = pending.remove(&id) {
@@ -396,6 +461,7 @@ impl CdpClient {
             // instead of waiting for the 30-second timeout.
             closed_reader.store(true, Ordering::SeqCst);
             pending_clone.lock().await.clear();
+            raw_pending_clone.lock().await.clear();
 
             // Stop the keepalive task because the connection is gone.
             let _ = reader_cancel_tx.send(true);
@@ -433,6 +499,7 @@ impl CdpClient {
             cancel_tx,
             next_id: AtomicU64::new(1),
             pending,
+            raw_pending,
             event_tx,
             raw_tx,
             reader_handle: Mutex::new(Some(reader_handle)),
@@ -477,6 +544,7 @@ impl CdpClient {
         let _ = self.socket.shutdown(Shutdown::Both);
         Self::abort_task(&self.reader_handle).await;
         self.pending.lock().await.clear();
+        self.raw_pending.lock().await.clear();
         self.close_complete.store(true, Ordering::SeqCst);
     }
 
@@ -503,6 +571,7 @@ impl CdpClient {
         // Cleans up the pending entry if this future is cancelled mid-await (#1528).
         let mut guard = PendingGuard {
             pending: self.pending.clone(),
+            raw_pending: self.raw_pending.clone(),
             id,
             done: false,
         };
@@ -623,6 +692,91 @@ impl CdpClient {
         self.send_command(method, None, session_id).await
     }
 
+    /// Send a command and deserialize its `result` straight from the response
+    /// text, skipping the intermediate `serde_json::Value`.
+    ///
+    /// The reader hands the response over unparsed, so the payload is
+    /// traversed once instead of twice. Use this for commands whose responses
+    /// can be megabytes (the full accessibility tree of a large page).
+    pub async fn send_command_typed_raw<P: serde::Serialize, R: serde::de::DeserializeOwned>(
+        &self,
+        method: &str,
+        params: &P,
+        session_id: Option<&str>,
+    ) -> Result<R, String> {
+        let params_value = serde_json::to_value(params)
+            .map_err(|e| format!("Failed to serialize params: {}", e))?;
+        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
+        let cmd = CdpCommand {
+            id,
+            method: method.to_string(),
+            params: Some(params_value),
+            session_id: session_id.filter(|s| !s.is_empty()).map(|s| s.to_string()),
+        };
+        let json = serde_json::to_string(&cmd)
+            .map_err(|e| format!("Failed to serialize CDP command: {}", e))?;
+
+        let (tx, rx) = oneshot::channel();
+
+        // Cleans up the pending entry if this future is cancelled mid-await (#1528).
+        let mut guard = PendingGuard {
+            pending: self.pending.clone(),
+            raw_pending: self.raw_pending.clone(),
+            id,
+            done: false,
+        };
+
+        {
+            let mut ws_tx = self.ws_tx.lock().await;
+            self.ensure_open()?;
+            let mut raw_pending = self.raw_pending.lock().await;
+            raw_pending.insert(id, tx);
+            drop(raw_pending);
+            if let Err(error) = ws_tx.send(Message::Text(json)).await {
+                self.raw_pending.lock().await.remove(&id);
+                guard.done = true;
+                return Err(format!("Failed to send CDP command: {}", error));
+            }
+        }
+
+        let text = match tokio::time::timeout(std::time::Duration::from_secs(30), rx).await {
+            Ok(Ok(text)) => {
+                guard.done = true;
+                text
+            }
+            Ok(Err(_)) => {
+                guard.done = true;
+                return Err("CDP response channel closed".to_string());
+            }
+            Err(_) => {
+                guard.done = true;
+                self.raw_pending.lock().await.remove(&id);
+                return Err(format!("CDP command timed out: {}", method));
+            }
+        };
+
+        let parse = |text: &str| -> Result<RawCommandResponse<R>, String> {
+            serde_json::from_str(text)
+                .map_err(|e| format!("Failed to deserialize CDP response for {}: {}", method, e))
+        };
+        // Chrome can emit lone surrogate escapes; repair them the same way the
+        // shared reader does and retry once.
+        let response = parse(&text).or_else(|error| match sanitize_lone_surrogates(&text) {
+            Some(repaired) => parse(&repaired),
+            None => Err(error),
+        })?;
+
+        if let Some(error) = response.error {
+            return Err(format!("CDP error ({}): {}", method, error));
+        }
+        response.result.ok_or_else(|| {
+            format!(
+                "Failed to deserialize CDP response for {}: missing result",
+                method
+            )
+        })
+    }
+
     /// Send a CDP command without waiting for its response.
     ///
     /// This is useful for best-effort commands where Chrome may not emit a
@@ -778,6 +932,105 @@ mod tests {
         assert_eq!(extract_command_id(r#"{"id":0,"result":{}}"#), None);
     }
 
+    #[test]
+    fn reads_response_ids_without_parsing_the_payload() {
+        assert_eq!(
+            leading_response_id(r#"{"id":42,"result":{"nodes":[]}}"#),
+            Some(42)
+        );
+        assert_eq!(
+            leading_response_id(r#"{"id":42,"result":{"nodes":[]},"sessionId":"s1"}"#),
+            Some(42)
+        );
+        assert_eq!(leading_response_id(r#"{"id":42}"#), Some(42));
+        assert_eq!(leading_response_id(r#"{"result":{"id":42}}"#), None);
+        assert_eq!(
+            leading_response_id(r#"{"method":"Page.loadEventFired","params":{}}"#),
+            None
+        );
+        assert_eq!(leading_response_id(r#"{"id":0,"result":{}}"#), None);
+        assert_eq!(leading_response_id(r#"{"id":-1,"result":{}}"#), None);
+        assert_eq!(leading_response_id(r#"{"id":"42","result":{}}"#), None);
+        assert_eq!(leading_response_id(""), None);
+        assert_eq!(
+            leading_response_id(r#"{"id":12345678901234567890123,"result":{}}"#),
+            None
+        );
+    }
+
+    #[derive(Debug, PartialEq, serde::Deserialize)]
+    struct RawNodes {
+        nodes: Vec<u32>,
+    }
+
+    #[tokio::test]
+    async fn raw_command_parses_leading_id_responses_and_errors() {
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let url = format!("ws://{}", listener.local_addr().unwrap());
+
+        let server = tokio::spawn(async move {
+            let (stream, _) = listener.accept().await.unwrap();
+            let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap();
+            while let Some(Ok(message)) = ws.next().await {
+                let Ok(text) = message.to_text() else {
+                    continue;
+                };
+                let request: Value = serde_json::from_str(text).unwrap();
+                let id = request["id"].as_u64().unwrap();
+                let reply = match request["method"].as_str().unwrap() {
+                    "Raw.Ok" => format!(r#"{{"id":{id},"result":{{"nodes":[1,2,3]}}}}"#),
+                    "Raw.Error" => {
+                        format!(r#"{{"id":{id},"error":{{"code":-32000,"message":"boom"}}}}"#)
+                    }
+                    // A proxy that puts sessionId first still reaches the waiter.
+                    "Raw.LateId" => {
+                        format!(r#"{{"sessionId":"s1","id":{id},"result":{{"nodes":[4]}}}}"#)
+                    }
+                    "Raw.Repaired" => {
+                        format!(r#"{{"id":{id},"result":{{"nodes":[5],"name":"\ud800"}}}}"#)
+                    }
+                    _ => format!(r#"{{"id":{id},"result":{{}}}}"#),
+                };
+                ws.send(Message::Text(reply)).await.unwrap();
+            }
+        });
+
+        let client = CdpClient::connect(&url).await.unwrap();
+
+        let ok: RawNodes = client
+            .send_command_typed_raw("Raw.Ok", &json!({}), None)
+            .await
+            .unwrap();
+        assert_eq!(
+            ok,
+            RawNodes {
+                nodes: vec![1, 2, 3]
+            }
+        );
+
+        let error = client
+            .send_command_typed_raw::<_, RawNodes>("Raw.Error", &json!({}), None)
+            .await
+            .unwrap_err();
+        assert!(error.contains("boom"), "unexpected error: {error}");
+
+        let late: RawNodes = client
+            .send_command_typed_raw("Raw.LateId", &json!({}), None)
+            .await
+            .unwrap();
+        assert_eq!(late, RawNodes { nodes: vec![4] });
+
+        // A raw waiter whose payload needs surrogate repair still deserializes.
+        let repaired: RawNodes = client
+            .send_command_typed_raw("Raw.Repaired", &json!({}), None)
+            .await
+            .unwrap();
+        assert_eq!(repaired, RawNodes { nodes: vec![5] });
+
+        client.close().await;
+        server.abort();
+    }
+
     #[tokio::test]
     async fn repairs_response_and_event_then_processes_next_command() {
         let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();

cli/src/native/element.rs

@@ -891,7 +891,7 @@ async fn find_node_id_by_role_name(
     let (ax_params, effective_session_id) =
         resolve_ax_session(frame_id, session_id, iframe_sessions);
     let ax_tree: GetFullAXTreeResult = client
-        .send_command_typed(
+        .send_command_typed_raw(
             "Accessibility.getFullAXTree",
             &ax_params,
             Some(effective_session_id),