From e225bdd64648c8f4aebb97e808ff7f399e8e9437 Mon Sep 17 00:00:00 2001 From: zhanba Date: Wed, 25 Mar 2026 13:22:44 +0800 Subject: [PATCH] fix: inline VS Code webviews in snapshots --- cli/src/native/actions.rs | 145 +++++++++++++++++++++++++++++-------- cli/src/native/browser.rs | 40 +++++++++- cli/src/native/snapshot.rs | 75 ++++++++++++++++++- 3 files changed, 225 insertions(+), 35 deletions(-) diff --git a/cli/src/native/actions.rs b/cli/src/native/actions.rs index 02c21747f..f09218d38 100644 --- a/cli/src/native/actions.rs +++ b/cli/src/native/actions.rs @@ -1624,6 +1624,74 @@ fn handle_cdp_url(state: &DaemonState) -> Result { Ok(json!({ "cdpUrl": mgr.get_cdp_url() })) } +fn is_vscode_webview_url(url: &str) -> bool { + url.starts_with("vscode-webview://") +} + +fn find_frame_id(tree: &Value, name: Option<&str>, url: Option<&str>) -> Option { + let frame = tree.get("frame")?; + let frame_name = frame.get("name").and_then(|v| v.as_str()).unwrap_or(""); + let frame_url = frame.get("url").and_then(|v| v.as_str()).unwrap_or(""); + let frame_id = frame.get("id").and_then(|v| v.as_str())?; + + if let Some(n) = name { + if frame_name == n { + return Some(frame_id.to_string()); + } + } + if let Some(u) = url { + if frame_url.contains(u) { + return Some(frame_id.to_string()); + } + } + + if let Some(children) = tree.get("childFrames").and_then(|v| v.as_array()) { + for child in children { + if let Some(id) = find_frame_id(child, name, url) { + return Some(id); + } + } + } + + None +} + +async fn resolve_vscode_active_frame_id( + mgr: &BrowserManager, + session_id: &str, + url: &str, +) -> Result, String> { + // VS Code webview targets expose the rendered extension UI inside a nested + // iframe named `active-frame`, not on the wrapper target itself. + if !is_vscode_webview_url(url) { + return Ok(None); + } + + let tree_result = mgr + .client + .send_command_no_params("Page.getFrameTree", Some(session_id)) + .await?; + Ok(find_frame_id( + &tree_result["frameTree"], + Some("active-frame"), + None, + )) +} + +fn sync_iframe_sessions_from_pages(state: &mut DaemonState) { + let Some(mgr) = state.browser.as_ref() else { + return; + }; + + for page in mgr.pages_list() { + if page.target_type == "iframe" { + state + .iframe_sessions + .insert(page.target_id, page.session_id); + } + } +} + async fn handle_inspect(state: &mut DaemonState) -> Result { let mgr = state.browser.as_ref().ok_or("Browser not launched")?; @@ -1771,8 +1839,15 @@ async fn handle_close(state: &mut DaemonState) -> Result { // --------------------------------------------------------------------------- async fn handle_snapshot(cmd: &Value, state: &mut DaemonState) -> Result { + sync_iframe_sessions_from_pages(state); let mgr = state.browser.as_ref().ok_or("Browser not launched")?; let session_id = mgr.active_session_id()?.to_string(); + let url = mgr.get_url().await.unwrap_or_default(); + let frame_id = if state.active_frame_id.is_some() { + state.active_frame_id.clone() + } else { + resolve_vscode_active_frame_id(mgr, &session_id, &url).await? + }; let options = SnapshotOptions { selector: cmd @@ -1799,13 +1874,11 @@ async fn handle_snapshot(cmd: &Value, state: &mut DaemonState) -> Result = state .ref_map .entries_sorted() @@ -4150,6 +4223,7 @@ async fn handle_waitforfunction(cmd: &Value, state: &DaemonState) -> Result Result { + sync_iframe_sessions_from_pages(state); let mgr = state.browser.as_mut().ok_or("Browser not launched")?; let session_id = mgr.active_session_id()?.to_string(); @@ -4166,33 +4240,6 @@ async fn handle_frame(cmd: &Value, state: &mut DaemonState) -> Result, url: Option<&str>) -> Option { - let frame = tree.get("frame")?; - let frame_name = frame.get("name").and_then(|v| v.as_str()).unwrap_or(""); - let frame_url = frame.get("url").and_then(|v| v.as_str()).unwrap_or(""); - let frame_id = frame.get("id").and_then(|v| v.as_str())?; - - if let Some(n) = name { - if frame_name == n { - return Some(frame_id.to_string()); - } - } - if let Some(u) = url { - if frame_url.contains(u) { - return Some(frame_id.to_string()); - } - } - - if let Some(children) = tree.get("childFrames").and_then(|v| v.as_array()) { - for child in children { - if let Some(id) = find_frame(child, name, url) { - return Some(id); - } - } - } - None - } - let frame_tree = &tree_result["frameTree"]; // If selector is a ref (@e1), resolve the iframe element from the ref map @@ -4275,13 +4322,13 @@ async fn handle_frame(cmd: &Value, state: &mut DaemonState) -> Result bool { || url.starts_with("devtools://") } +fn is_vscode_webview_target(url: &str) -> bool { + url.starts_with("vscode-webview://") +} + pub(crate) fn should_track_target(target: &TargetInfo) -> bool { - (target.target_type == "page" || target.target_type == "webview") + (target.target_type == "page" + || target.target_type == "webview" + || (target.target_type == "iframe" && is_vscode_webview_target(&target.url))) && (target.url.is_empty() || !is_internal_chrome_target(&target.url)) } @@ -141,7 +147,7 @@ pub struct PageInfo { pub session_id: String, pub url: String, pub title: String, - pub target_type: String, // "page" or "webview" + pub target_type: String, // "page", "webview", "iframe", or another CDP target type } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -1372,6 +1378,36 @@ mod tests { assert!(!should_track_target(&target)); } + #[test] + fn test_should_track_vscode_webview_iframe_target() { + let target = TargetInfo { + target_id: "vscode-webview-iframe".to_string(), + target_type: "iframe".to_string(), + title: "Example".to_string(), + url: "vscode-webview://123/index.html?id=webview-1&extensionId=publisher.example" + .to_string(), + attached: None, + browser_context_id: None, + }; + + assert!(should_track_target(&target)); + } + + #[test] + fn test_should_not_track_vscode_service_worker_target() { + let target = TargetInfo { + target_id: "vscode-webview-sw".to_string(), + target_type: "service_worker".to_string(), + title: "Service Worker".to_string(), + url: "vscode-webview://123/service-worker.js?v=4&extensionId=publisher.example" + .to_string(), + attached: None, + browser_context_id: None, + }; + + assert!(!should_track_target(&target)); + } + #[test] fn test_update_page_target_info_in_pages_updates_existing_page() { let mut pages = vec![PageInfo { diff --git a/cli/src/native/snapshot.rs b/cli/src/native/snapshot.rs index 77c554f44..e81269a8d 100644 --- a/cli/src/native/snapshot.rs +++ b/cli/src/native/snapshot.rs @@ -184,6 +184,72 @@ impl RoleNameTracker { } } +fn is_vscode_webview_url(url: &str) -> bool { + url.starts_with("vscode-webview://") +} + +fn find_frame_id(tree: &Value, name: Option<&str>, url: Option<&str>) -> Option { + let frame = tree.get("frame")?; + let frame_name = frame.get("name").and_then(|v| v.as_str()).unwrap_or(""); + let frame_url = frame.get("url").and_then(|v| v.as_str()).unwrap_or(""); + let frame_id = frame.get("id").and_then(|v| v.as_str())?; + + if let Some(n) = name { + if frame_name == n { + return Some(frame_id.to_string()); + } + } + if let Some(u) = url { + if frame_url.contains(u) { + return Some(frame_id.to_string()); + } + } + + if let Some(children) = tree.get("childFrames").and_then(|v| v.as_array()) { + for child in children { + if let Some(id) = find_frame_id(child, name, url) { + return Some(id); + } + } + } + + None +} + +async fn resolve_vscode_active_frame_id( + client: &CdpClient, + session_id: &str, +) -> Result, String> { + let location: EvaluateResult = client + .send_command_typed( + "Runtime.evaluate", + &EvaluateParams { + expression: "location.href".to_string(), + return_by_value: Some(true), + await_promise: Some(false), + }, + Some(session_id), + ) + .await?; + let url = location + .result + .value + .and_then(|v| v.as_str().map(|s| s.to_string())) + .unwrap_or_default(); + if !is_vscode_webview_url(&url) { + return Ok(None); + } + + let tree_result = client + .send_command_no_params("Page.getFrameTree", Some(session_id)) + .await?; + Ok(find_frame_id( + &tree_result["frameTree"], + Some("active-frame"), + None, + )) +} + pub async fn take_snapshot( client: &CdpClient, session_id: &str, @@ -254,7 +320,7 @@ pub async fn take_snapshot( None }; - let (ax_params, effective_session_id) = + let (mut ax_params, effective_session_id) = resolve_ax_session(frame_id, session_id, iframe_sessions); // Ensure domains are enabled on the iframe session (defensive fallback // in case the attach-time enable in execute_command was missed). @@ -266,6 +332,13 @@ pub async fn take_snapshot( .send_command_no_params("Accessibility.enable", Some(effective_session_id)) .await; } + if frame_id.is_none() || frame_id.is_some_and(|fid| iframe_sessions.contains_key(fid)) { + if let Some(active_frame_id) = + resolve_vscode_active_frame_id(client, effective_session_id).await? + { + ax_params = serde_json::json!({ "frameId": active_frame_id }); + } + } let ax_tree: GetFullAXTreeResult = client .send_command_typed( "Accessibility.getFullAXTree",