Skip to content

fix: inline VS Code webviews in snapshots #1013

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zhanba wants to merge 1 commit into
base: main
Choose a base branch
from
+225 −35
Open

fix: inline VS Code webviews in snapshots #1013

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
145 changes: 113 additions & 32 deletions cli/src/native/actions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1624,6 +1624,74 @@ fn handle_cdp_url(state: &DaemonState) -> Result<Value, String> {
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<String> {
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<Option<String>, 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<Value, String> {
let mgr = state.browser.as_ref().ok_or("Browser not launched")?;

Expand Down Expand Up @@ -1771,8 +1839,15 @@ async fn handle_close(state: &mut DaemonState) -> Result<Value, String> {
// ---------------------------------------------------------------------------

async fn handle_snapshot(cmd: &Value, state: &mut DaemonState) -> Result<Value, String> {
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
Expand All @@ -1799,13 +1874,11 @@ async fn handle_snapshot(cmd: &Value, state: &mut DaemonState) -> Result<Value,
&session_id,
&options,
&mut state.ref_map,
state.active_frame_id.as_deref(),
frame_id.as_deref(),
&state.iframe_sessions,
)
.await?;

let url = mgr.get_url().await.unwrap_or_default();

let refs: serde_json::Map<String, Value> = state
.ref_map
.entries_sorted()
Expand Down Expand Up @@ -4150,6 +4223,7 @@ async fn handle_waitforfunction(cmd: &Value, state: &DaemonState) -> Result<Valu
// ---------------------------------------------------------------------------

async fn handle_frame(cmd: &Value, state: &mut DaemonState) -> Result<Value, String> {
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();

Expand All @@ -4166,33 +4240,6 @@ async fn handle_frame(cmd: &Value, state: &mut DaemonState) -> Result<Value, Str
.send_command_no_params("Page.getFrameTree", Some(&session_id))
.await?;

fn find_frame(tree: &Value, name: Option<&str>, url: Option<&str>) -> Option<String> {
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
Expand Down Expand Up @@ -4275,13 +4322,13 @@ async fn handle_frame(cmd: &Value, state: &mut DaemonState) -> Result<Value, Str
);
let result = mgr.evaluate(&js, None).await?;
let frame_name = result.as_str().ok_or("Could not find frame for selector")?;
if let Some(frame_id) = find_frame(frame_tree, Some(frame_name), None) {
if let Some(frame_id) = find_frame_id(frame_tree, Some(frame_name), None) {
state.active_frame_id = Some(frame_id);
return Ok(json!({ "frame": frame_name }));
}
}

if let Some(frame_id) = find_frame(frame_tree, name, url) {
if let Some(frame_id) = find_frame_id(frame_tree, name, url) {
let label = name.or(url).unwrap_or("frame");
state.active_frame_id = Some(frame_id);
return Ok(json!({ "frame": label }));
Expand Down Expand Up @@ -6763,6 +6810,40 @@ mod tests {
assert_eq!(resp["error"], "Something went wrong");
}

#[test]
fn test_find_frame_id_finds_vscode_active_frame() {
let tree = json!({
"frame": {
"id": "root",
"name": "",
"url": "vscode-file://vscode-app/Applications/Visual%20Studio%20Code.app/Contents/Resources/app/out/vs/code/electron-browser/workbench/workbench.html"
},
"childFrames": [
{
"frame": {
"id": "wrapper",
"name": "webview-host",
"url": "vscode-webview://1l4q6f4e6i2l8rb3q0r4d0v6kq2b6jgg/index.html?id=webview-element-uuid"
},
"childFrames": [
{
"frame": {
"id": "active",
"name": "active-frame",
"url": "vscode-webview://1l4q6f4e6i2l8rb3q0r4d0v6kq2b6jgg/index.html?id=webview-element-uuid&origin=tabby"
}
}
]
}
]
});

assert_eq!(
find_frame_id(&tree, Some("active-frame"), None),
Some("active".to_string())
);
}

#[tokio::test]
async fn test_daemon_state_new() {
let state = DaemonState::new();
Expand Down
40 changes: 38 additions & 2 deletions cli/src/native/browser.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,8 +95,14 @@ fn is_internal_chrome_target(url: &str) -> 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))
}

Expand Down Expand Up @@ -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)]
Expand Down Expand Up @@ -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 {
Expand Down
75 changes: 74 additions & 1 deletion cli/src/native/snapshot.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<String> {
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<Option<String>, 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,
Expand Down Expand Up @@ -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).
Expand All @@ -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",
Expand Down