mod cdp_helpers;
mod frame_stitching;
mod options;
mod ref_resolution;
use std::collections::HashMap;
use futures::stream::{FuturesUnordered, StreamExt};
use tracing::{debug, instrument};
use viewpoint_js::js;
use self::frame_stitching::stitch_frame_content;
pub use self::options::SnapshotOptions;
pub(crate) use self::ref_resolution::apply_refs_to_snapshot;
use super::Page;
use super::locator::AriaSnapshot;
use super::locator::aria_js::aria_snapshot_with_refs_js;
use crate::error::PageError;
impl Page {
#[instrument(level = "debug", skip(self), fields(target_id = %self.target_id))]
pub async fn aria_snapshot_with_frames(&self) -> Result<AriaSnapshot, PageError> {
self.aria_snapshot_with_frames_and_options(SnapshotOptions::default())
.await
}
#[instrument(level = "debug", skip(self, options), fields(target_id = %self.target_id))]
pub async fn aria_snapshot_with_frames_and_options(
&self,
options: SnapshotOptions,
) -> Result<AriaSnapshot, PageError> {
if self.closed {
return Err(PageError::Closed);
}
let mut root_snapshot = self.capture_snapshot_with_refs(options.clone()).await?;
let frames = self.frames().await?;
let child_frames: Vec<_> = frames.iter().filter(|f| !f.is_main()).collect();
if child_frames.is_empty() {
return Ok(root_snapshot);
}
debug!(
frame_count = child_frames.len(),
"Capturing child frame snapshots in parallel"
);
let frame_futures: FuturesUnordered<_> = child_frames
.iter()
.map(|frame| {
let frame_id = frame.id().to_string();
let frame_url = frame.url().clone();
let frame_name = frame.name().clone();
let opts = options.clone();
async move {
match frame.capture_snapshot_with_refs(opts).await {
Ok((snapshot, ref_mappings)) => {
Some((frame_id, frame_url, frame_name, snapshot, ref_mappings))
}
Err(e) => {
tracing::warn!(
error = %e,
frame_id = %frame_id,
frame_url = %frame_url,
"Failed to capture frame snapshot, skipping"
);
None
}
}
}
})
.collect();
let results: Vec<_> = frame_futures.collect().await;
let mut frame_snapshots: HashMap<String, AriaSnapshot> = HashMap::new();
for result in results.into_iter().flatten() {
let (frame_id, frame_url, frame_name, snapshot, ref_mappings) = result;
for (ref_str, backend_node_id) in ref_mappings {
self.store_ref_mapping(ref_str, backend_node_id);
}
if !frame_url.is_empty() && frame_url != "about:blank" {
frame_snapshots.insert(frame_url, snapshot.clone());
}
if !frame_name.is_empty() {
frame_snapshots.insert(frame_name, snapshot.clone());
}
frame_snapshots.insert(frame_id, snapshot);
}
stitch_frame_content(&mut root_snapshot, &frame_snapshots, 0);
Ok(root_snapshot)
}
#[instrument(level = "debug", skip(self), fields(target_id = %self.target_id))]
pub async fn aria_snapshot(&self) -> Result<AriaSnapshot, PageError> {
self.aria_snapshot_with_options(SnapshotOptions::default())
.await
}
#[instrument(level = "debug", skip(self, options), fields(target_id = %self.target_id))]
pub async fn aria_snapshot_with_options(
&self,
options: SnapshotOptions,
) -> Result<AriaSnapshot, PageError> {
if self.closed {
return Err(PageError::Closed);
}
self.capture_snapshot_with_refs(options).await
}
#[instrument(level = "debug", skip(self, options), fields(target_id = %self.target_id))]
async fn capture_snapshot_with_refs(
&self,
options: SnapshotOptions,
) -> Result<AriaSnapshot, PageError> {
let snapshot_fn = aria_snapshot_with_refs_js();
let js_code = js! {
(function() {
const getSnapshotWithRefs = @{snapshot_fn};
return getSnapshotWithRefs(document.body);
})()
};
let result: viewpoint_cdp::protocol::runtime::EvaluateResult = self
.connection()
.send_command(
"Runtime.evaluate",
Some(viewpoint_cdp::protocol::runtime::EvaluateParams {
expression: js_code,
object_group: Some("viewpoint-snapshot".to_string()),
include_command_line_api: None,
silent: Some(true),
context_id: None,
return_by_value: Some(false), await_promise: Some(false),
}),
Some(self.session_id()),
)
.await?;
if let Some(exception) = result.exception_details {
return Err(PageError::EvaluationFailed(exception.text));
}
let result_object_id = result.result.object_id.ok_or_else(|| {
PageError::EvaluationFailed("No object ID from snapshot evaluation".to_string())
})?;
let snapshot_value = self
.get_property_value(&result_object_id, "snapshot")
.await?;
let mut snapshot: AriaSnapshot = serde_json::from_value(snapshot_value).map_err(|e| {
PageError::EvaluationFailed(format!("Failed to parse aria snapshot: {e}"))
})?;
self.clear_ref_map();
if options.include_refs {
let elements_result = self
.get_property_object(&result_object_id, "elements")
.await?;
if let Some(elements_object_id) = elements_result {
let element_object_ids = self.get_all_array_elements(&elements_object_id).await?;
let element_count = element_object_ids.len();
debug!(
element_count = element_count,
max_concurrency = options.max_concurrency,
"Resolving element refs in parallel"
);
let index_to_backend_id = self
.resolve_node_ids_parallel(element_object_ids, options.max_concurrency)
.await;
debug!(
resolved_count = index_to_backend_id.len(),
total_count = element_count,
"Completed parallel ref resolution"
);
let ref_to_backend_id = apply_refs_to_snapshot(
&mut snapshot,
&index_to_backend_id,
self.context_index,
self.page_index,
0, );
for (ref_str, backend_node_id) in ref_to_backend_id {
self.store_ref_mapping(ref_str, backend_node_id);
}
let _ = self.release_object(&elements_object_id).await;
}
}
let _ = self.release_object(&result_object_id).await;
Ok(snapshot)
}
}