use std::sync::Arc;
use chrono::Utc;
use tokio::sync::RwLock;
use tracing::{debug, info, instrument};
use viewpoint_cdp::CdpConnection;
use viewpoint_cdp::protocol::tracing as cdp_tracing;
use crate::error::ContextError;
use crate::network::har::HarPage;
use crate::page::Page;
use super::action_handle::ActionHandle;
use super::capture;
use super::network;
use super::sources;
use super::types::{ActionEntry, SourceFileEntry, TracingOptions, TracingState};
use super::writer;
pub struct Tracing {
connection: Arc<CdpConnection>,
context_id: String,
pages: Arc<RwLock<Vec<Page>>>,
state: Arc<RwLock<TracingState>>,
}
impl std::fmt::Debug for Tracing {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Tracing")
.field("context_id", &self.context_id)
.finish_non_exhaustive()
}
}
impl Tracing {
pub(crate) fn new(
connection: Arc<CdpConnection>,
context_id: String,
pages: Arc<RwLock<Vec<Page>>>,
state: Arc<RwLock<TracingState>>,
) -> Self {
Self {
connection,
context_id,
pages,
state,
}
}
async fn get_session_ids(&self) -> Vec<String> {
let pages = self.pages.read().await;
pages
.iter()
.filter(|p| !p.session_id().is_empty())
.map(|p| p.session_id().to_string())
.collect()
}
#[instrument(level = "info", skip(self, options))]
pub async fn start(&self, options: TracingOptions) -> Result<(), ContextError> {
let mut state = self.state.write().await;
if state.is_recording {
return Err(ContextError::Internal(
"Tracing is already active".to_string(),
));
}
let session_ids = self.get_session_ids().await;
if session_ids.is_empty() {
return Err(ContextError::Internal(
"Cannot start tracing: no pages in context. Create a page first.".to_string(),
));
}
info!(
screenshots = options.screenshots,
snapshots = options.snapshots,
"Starting trace"
);
let categories = [
"devtools.timeline",
"disabled-by-default-devtools.timeline",
"disabled-by-default-devtools.timeline.frame",
];
for session_id in session_ids {
let params = cdp_tracing::StartParams {
categories: Some(categories.join(",")),
transfer_mode: Some(cdp_tracing::TransferMode::ReturnAsStream),
..Default::default()
};
self.connection
.send_command::<_, serde_json::Value>(
"Tracing.start",
Some(params),
Some(&session_id),
)
.await?;
self.connection
.send_command::<_, serde_json::Value>(
"Network.enable",
Some(serde_json::json!({})),
Some(&session_id),
)
.await?;
}
state.is_recording = true;
state.options = options;
state.actions.clear();
state.events.clear();
state.screenshots.clear();
state.snapshots.clear();
state.pending_requests.clear();
state.network_entries.clear();
state.har_pages.clear();
state.source_files.clear();
drop(state); network::start_network_listener(
self.connection.clone(),
self.state.clone(),
self.pages.clone(),
);
Ok(())
}
#[instrument(level = "info", skip(self), fields(path = %path.as_ref().display()))]
pub async fn stop(&self, path: impl AsRef<std::path::Path>) -> Result<(), ContextError> {
let path = path.as_ref();
let mut state = self.state.write().await;
if !state.is_recording {
return Err(ContextError::Internal("Tracing is not active".to_string()));
}
info!("Stopping trace and saving");
for session_id in self.get_session_ids().await {
let _ = self
.connection
.send_command::<_, serde_json::Value>("Tracing.end", None::<()>, Some(&session_id))
.await;
}
state.is_recording = false;
writer::write_trace_zip(path, &state)?;
Ok(())
}
#[instrument(level = "info", skip(self))]
pub async fn stop_discard(&self) -> Result<(), ContextError> {
let mut state = self.state.write().await;
if !state.is_recording {
return Err(ContextError::Internal("Tracing is not active".to_string()));
}
info!("Stopping trace and discarding");
for session_id in self.get_session_ids().await {
let _ = self
.connection
.send_command::<_, serde_json::Value>("Tracing.end", None::<()>, Some(&session_id))
.await;
}
state.is_recording = false;
state.actions.clear();
state.events.clear();
state.screenshots.clear();
state.snapshots.clear();
state.pending_requests.clear();
state.network_entries.clear();
state.har_pages.clear();
state.source_files.clear();
Ok(())
}
#[instrument(level = "debug", skip(self))]
pub async fn start_chunk(&self) -> Result<(), ContextError> {
let state = self.state.read().await;
if !state.is_recording {
return Err(ContextError::Internal("Tracing is not active".to_string()));
}
debug!("Starting new trace chunk");
Ok(())
}
#[instrument(level = "debug", skip(self), fields(path = %path.as_ref().display()))]
pub async fn stop_chunk(&self, path: impl AsRef<std::path::Path>) -> Result<(), ContextError> {
let path = path.as_ref();
let state = self.state.read().await;
if !state.is_recording {
return Err(ContextError::Internal("Tracing is not active".to_string()));
}
debug!("Stopping trace chunk and saving");
writer::write_trace_zip(path, &state)?;
Ok(())
}
pub async fn is_recording(&self) -> bool {
self.state.read().await.is_recording
}
pub async fn add_source_file(&self, path: impl Into<String>, content: impl Into<String>) {
let mut state = self.state.write().await;
state.source_files.push(SourceFileEntry {
path: path.into(),
content: content.into(),
});
}
pub async fn collect_sources(
&self,
dir: impl AsRef<std::path::Path>,
extensions: &[&str],
) -> Result<(), ContextError> {
let files = sources::collect_sources_from_dir(dir.as_ref(), extensions)?;
let mut state = self.state.write().await;
for (path, content) in files {
state.source_files.push(SourceFileEntry { path, content });
}
Ok(())
}
pub(crate) async fn record_action(
&self,
action_type: &str,
selector: Option<&str>,
page_id: Option<&str>,
) -> ActionHandle {
let start_time = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs_f64()
* 1000.0;
let action = ActionEntry {
action_type: action_type.to_string(),
selector: selector.map(ToString::to_string),
page_id: page_id.map(ToString::to_string),
start_time,
end_time: None,
result: None,
value: None,
url: None,
screenshot: None,
snapshot: None,
};
let mut state = self.state.write().await;
let index = state.actions.len();
state.actions.push(action);
ActionHandle::new(self.state.clone(), index)
}
pub(crate) async fn record_page(&self, page_id: &str, title: &str) {
let mut state = self.state.write().await;
let started_date_time = Utc::now().to_rfc3339();
let page = HarPage::new(page_id, title, &started_date_time);
state.har_pages.push(page);
state.current_page_id = Some(page_id.to_string());
}
pub(crate) async fn capture_screenshot(
&self,
session_id: &str,
name: Option<&str>,
) -> Result<(), ContextError> {
capture::capture_screenshot(&self.connection, &self.state, session_id, name).await
}
pub(crate) async fn capture_dom_snapshot(&self, session_id: &str) -> Result<(), ContextError> {
capture::capture_dom_snapshot(&self.connection, &self.state, session_id).await
}
pub(crate) async fn capture_action_context(
&self,
session_id: &str,
action_name: Option<&str>,
) -> Result<(), ContextError> {
capture::capture_action_context(&self.connection, &self.state, session_id, action_name)
.await
}
}