diff --git a/cli/src/native/actions.rs b/cli/src/native/actions.rs index 9ea57cc..c84b6c3 100644 --- a/cli/src/native/actions.rs +++ b/cli/src/native/actions.rs @@ -26,7 +26,7 @@ use super::screenshot::{self, ScreenshotOptions}; use super::snapshot::{self, SnapshotOptions}; use super::state; use super::storage; -use super::stream; +use super::stream::{self, StreamServer}; use super::tracing::{self as native_tracing, TracingState}; use super::webdriver::appium::AppiumManager; use super::webdriver::backend::{BrowserBackend, WebDriverBackend, WEBDRIVER_UNSUPPORTED_ACTIONS}; @@ -108,6 +108,8 @@ pub struct DaemonState { pub active_frame_id: Option, /// Shared slot for stream server to receive CDP client when browser launches. pub stream_client: Option>>>>, + /// Stream server instance kept alive so the broadcast channel remains open. + pub stream_server: Option>, } impl DaemonState { @@ -141,15 +143,19 @@ impl DaemonState { request_tracking: false, active_frame_id: None, stream_client: None, + stream_server: None, } } - /// Create state with an optional stream client slot (for daemon startup with stream server). - pub fn new_with_stream_client( + /// Create state with an optional stream client slot and server instance + /// (for daemon startup with stream server). + pub fn new_with_stream( stream_client: Option>>>>, + stream_server: Option>, ) -> Self { let mut s = Self::new(); s.stream_client = stream_client; + s.stream_server = stream_server; s } @@ -165,6 +171,21 @@ impl DaemonState { let mut guard = slot.write().await; *guard = self.browser.as_ref().map(|m| Arc::clone(&m.client)); } + if let Some(ref server) = self.stream_server { + // Update the CDP page session ID so screencast commands target the right page + let session_id = self + .browser + .as_ref() + .and_then(|m| m.active_session_id().ok().map(|s| s.to_string())); + server.set_cdp_session_id(session_id).await; + + // Broadcast connection status change to WebSocket clients + let connected = self.browser.is_some(); + let sc = server.is_screencasting().await; + server.broadcast_status(connected, sc, 1280, 720); + // Notify the background CDP event loop that the client changed + server.notify_client_changed(); + } } /// Spawn a background task that polls screenshots and pipes them to ffmpeg. @@ -372,10 +393,15 @@ impl DaemonState { } } "Page.screencastFrame" => { - if let Some(sid) = - event.params.get("sessionId").and_then(|v| v.as_i64()) - { - pending_acks.push(sid); + // Frame broadcasting and acks are handled in real-time by the + // stream server's background CDP event loop. Here we just + // collect acks as a fallback for non-streaming mode. + if self.stream_server.is_none() { + if let Some(sid) = + event.params.get("sessionId").and_then(|v| v.as_i64()) + { + pending_acks.push(sid); + } } } "Fetch.requestPaused" => { @@ -3305,6 +3331,10 @@ async fn handle_screencast_start(cmd: &Value, state: &mut DaemonState) -> Result .await?; state.screencasting = true; + if let Some(ref server) = state.stream_server { + server.broadcast_status(true, true, max_width as u32, max_height as u32); + } + Ok(json!({ "started": true })) } @@ -3319,6 +3349,10 @@ async fn handle_screencast_stop(state: &mut DaemonState) -> Result>>>> = None; + let mut stream_server_instance: Option> = None; if let Ok(port_str) = env::var("AGENT_BROWSER_STREAM_PORT") { if let Ok(port) = port_str.parse::() { if port > 0 { @@ -51,6 +52,7 @@ pub async fn run_daemon(session: &str) { let _ = writeln!(std::io::stderr(), "Failed to write .stream file: {}", e); } + stream_server_instance = Some(Arc::new(stream_server)); } Err(e) => { let _ = writeln!(std::io::stderr(), "Stream server failed to start: {}", e); @@ -67,7 +69,14 @@ pub async fn run_daemon(session: &str) { .and_then(|s| s.parse::().ok()) .filter(|&ms| ms > 0); - let result = run_socket_server(&socket_path, session, stream_client, idle_timeout_ms).await; + let result = run_socket_server( + &socket_path, + session, + stream_client, + stream_server_instance, + idle_timeout_ms, + ) + .await; let _ = fs::remove_file(&socket_path); let _ = fs::remove_file(&pid_path); @@ -85,6 +94,7 @@ async fn run_socket_server( socket_path: &PathBuf, _session: &str, stream_client: Option>>>>, + stream_server: Option>, idle_timeout_ms: Option, ) -> Result<(), String> { use tokio::net::UnixListener; @@ -93,7 +103,7 @@ async fn run_socket_server( UnixListener::bind(socket_path).map_err(|e| format!("Failed to bind socket: {}", e))?; let state: std::sync::Arc> = std::sync::Arc::new( - tokio::sync::Mutex::new(DaemonState::new_with_stream_client(stream_client)), + tokio::sync::Mutex::new(DaemonState::new_with_stream(stream_client, stream_server)), ); let (reset_tx, mut reset_rx) = mpsc::channel::<()>(64); @@ -152,6 +162,7 @@ async fn run_socket_server( socket_path: &PathBuf, session: &str, stream_client: Option>>>>, + stream_server: Option>, idle_timeout_ms: Option, ) -> Result<(), String> { use tokio::net::TcpListener; @@ -166,7 +177,7 @@ async fn run_socket_server( let _ = fs::write(&port_path, port.to_string()); let state: std::sync::Arc> = std::sync::Arc::new( - tokio::sync::Mutex::new(DaemonState::new_with_stream_client(stream_client)), + tokio::sync::Mutex::new(DaemonState::new_with_stream(stream_client, stream_server)), ); let (reset_tx, mut reset_rx) = mpsc::channel::<()>(64); diff --git a/cli/src/native/stream.rs b/cli/src/native/stream.rs index 95694c1..939dbfc 100644 --- a/cli/src/native/stream.rs +++ b/cli/src/native/stream.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use futures_util::{SinkExt, StreamExt}; use tokio::net::TcpListener; -use tokio::sync::{broadcast, Mutex, RwLock}; +use tokio::sync::{broadcast, Mutex, Notify, RwLock}; use tokio_tungstenite::tungstenite::Message; use super::cdp::client::CdpClient; @@ -39,6 +39,11 @@ pub struct StreamServer { port: u16, frame_tx: broadcast::Sender, client_count: Arc>, + client_slot: Arc>>>, + /// The active CDP page session ID (from Target.attachToTarget). + cdp_session_id: Arc>>, + client_notify: Arc, + screencasting: Arc>, } impl StreamServer { @@ -63,10 +68,26 @@ impl StreamServer { Self::start_inner(preferred_port, client_slot, session_id).await } + /// Notify the background CDP listener that the client has changed (browser launched/closed). + pub fn notify_client_changed(&self) { + self.client_notify.notify_one(); + } + + /// Update the active CDP page session ID used for screencast commands. + pub async fn set_cdp_session_id(&self, session_id: Option) { + let mut guard = self.cdp_session_id.write().await; + *guard = session_id; + } + + /// Check whether the server currently has active screencast running. + pub async fn is_screencasting(&self) -> bool { + *self.screencasting.lock().await + } + async fn start_inner( preferred_port: u16, client_slot: Arc>>>, - session_id: String, + _session_id: String, ) -> Result<(Self, Arc>>>), String> { let addr = format!("127.0.0.1:{}", preferred_port); let listener = TcpListener::bind(&addr) @@ -80,18 +101,46 @@ impl StreamServer { let (frame_tx, _) = broadcast::channel::(64); let client_count = Arc::new(Mutex::new(0usize)); + let client_notify = Arc::new(Notify::new()); + let screencasting = Arc::new(Mutex::new(false)); + let cdp_session_id = Arc::new(RwLock::new(None::)); let frame_tx_clone = frame_tx.clone(); let client_count_clone = client_count.clone(); let client_slot_clone = client_slot.clone(); + let notify_clone = client_notify.clone(); + let screencasting_clone = screencasting.clone(); + let cdp_session_clone = cdp_session_id.clone(); + // WebSocket accept loop tokio::spawn(async move { accept_loop( listener, frame_tx_clone, client_count_clone, client_slot_clone, - session_id, + notify_clone, + screencasting_clone, + cdp_session_clone, + ) + .await; + }); + + // Background CDP event listener for real-time frame broadcasting + let frame_tx_bg = frame_tx.clone(); + let client_slot_bg = client_slot.clone(); + let client_notify_bg = client_notify.clone(); + let screencasting_bg = screencasting.clone(); + let client_count_bg = client_count.clone(); + let cdp_session_bg = cdp_session_id.clone(); + tokio::spawn(async move { + cdp_event_loop( + frame_tx_bg, + client_slot_bg, + client_notify_bg, + screencasting_bg, + client_count_bg, + cdp_session_bg, ) .await; }); @@ -101,6 +150,10 @@ impl StreamServer { port, frame_tx, client_count, + client_slot: client_slot.clone(), + cdp_session_id, + client_notify, + screencasting, }, client_slot, )) @@ -161,33 +214,50 @@ impl StreamServer { } } +#[allow(clippy::too_many_arguments)] async fn accept_loop( listener: TcpListener, frame_tx: broadcast::Sender, client_count: Arc>, client_slot: Arc>>>, - session_id: String, + client_notify: Arc, + screencasting: Arc>, + cdp_session_id: Arc>>, ) { while let Ok((stream, addr)) = listener.accept().await { let frame_rx = frame_tx.subscribe(); let client_count = client_count.clone(); let client_slot = client_slot.clone(); - let sid = session_id.clone(); + let client_notify = client_notify.clone(); + let screencasting = screencasting.clone(); + let cdp_session_id = cdp_session_id.clone(); tokio::spawn(async move { - handle_ws_client(stream, addr, frame_rx, client_count, client_slot, sid).await; + handle_ws_client( + stream, + addr, + frame_rx, + client_count, + client_slot, + client_notify, + screencasting, + cdp_session_id, + ) + .await; }); } } -#[allow(clippy::result_large_err)] +#[allow(clippy::result_large_err, clippy::too_many_arguments)] async fn handle_ws_client( stream: tokio::net::TcpStream, _addr: SocketAddr, mut frame_rx: broadcast::Receiver, client_count: Arc>, client_slot: Arc>>>, - session_id: String, + client_notify: Arc, + screencasting: Arc>, + cdp_session_id: Arc>>, ) { let callback = |req: &tokio_tungstenite::tungstenite::handshake::server::Request, @@ -220,6 +290,24 @@ async fn handle_ws_client( let (mut ws_tx, mut ws_rx) = ws_stream.split(); + // Send initial status (screencasting:false initially, matching 0.19.0) + { + let guard = client_slot.read().await; + let connected = guard.is_some(); + let sc = *screencasting.lock().await; + let status = json!({ + "type": "status", + "connected": connected, + "screencasting": sc, + "viewportWidth": 1280, + "viewportHeight": 720, + }); + let _ = ws_tx.send(Message::Text(status.to_string())).await; + } + + // Notify the CDP event loop that a client connected (may trigger auto-start screencast) + client_notify.notify_one(); + loop { tokio::select! { frame = frame_rx.recv() => { @@ -229,7 +317,11 @@ async fn handle_ws_client( break; } } - Err(_) => break, + Err(broadcast::error::RecvError::Lagged(_)) => { + // Slow consumer; skip missed frames and continue + continue; + } + Err(broadcast::error::RecvError::Closed) => break, } } msg = ws_rx.next() => { @@ -237,7 +329,8 @@ async fn handle_ws_client( Some(Ok(Message::Text(text))) => { let guard = client_slot.read().await; if let Some(ref client) = *guard { - handle_client_message(&text, client.as_ref(), &session_id).await; + let sid = cdp_session_id.read().await; + handle_client_message(&text, client.as_ref(), sid.as_deref()).await; } } Some(Ok(Message::Close(_))) | None => break, @@ -251,9 +344,166 @@ async fn handle_ws_client( let mut count = client_count.lock().await; *count = count.saturating_sub(1); } + + // Notify the CDP event loop that a client disconnected (may trigger auto-stop screencast) + client_notify.notify_one(); } -async fn handle_client_message(msg: &str, client: &CdpClient, session_id: &str) { +/// Background task that subscribes to CDP events and broadcasts screencast frames in real-time. +/// Also handles auto-start/stop of screencast based on WebSocket client count. +async fn cdp_event_loop( + frame_tx: broadcast::Sender, + client_slot: Arc>>>, + client_notify: Arc, + screencasting: Arc>, + client_count: Arc>, + cdp_session_id: Arc>>, +) { + loop { + // Wait until we're notified of a client/connection change + client_notify.notified().await; + + // Check if we have WS clients and a CDP client + let count = *client_count.lock().await; + let guard = client_slot.read().await; + + if count > 0 { + if let Some(ref client) = *guard { + // We have WS clients and a CDP client — start screencast and listen for frames + let mut event_rx = client.subscribe(); + let client_arc = Arc::clone(client); + drop(guard); + + // Get the CDP page session ID for targeted commands + let session_id = cdp_session_id.read().await.clone(); + + let _ = client_arc + .send_command( + "Page.startScreencast", + Some(json!({ + "format": "jpeg", + "quality": 80, + "maxWidth": 1280, + "maxHeight": 720, + "everyNthFrame": 1, + })), + session_id.as_deref(), + ) + .await; + + { + let mut sc = screencasting.lock().await; + *sc = true; + } + + // Broadcast screencasting:true status (matching 0.19.0 two-status sequence) + let status = json!({ + "type": "status", + "connected": true, + "screencasting": true, + "viewportWidth": 1280, + "viewportHeight": 720, + }); + let _ = frame_tx.send(status.to_string()); + + // Process CDP events in real-time until client disconnects or CDP closes + loop { + tokio::select! { + event = event_rx.recv() => { + match event { + Ok(evt) => { + if evt.method == "Page.screencastFrame" { + // Ack immediately (like 0.19.0) + if let Some(sid) = evt.params.get("sessionId").and_then(|v| v.as_i64()) { + let _ = client_arc.send_command( + "Page.screencastFrameAck", + Some(json!({ "sessionId": sid })), + evt.session_id.as_deref(), + ).await; + } + + // Broadcast frame to WS clients + if let Some(data) = evt.params.get("data").and_then(|v| v.as_str()) { + let meta = evt.params.get("metadata"); + let msg = json!({ + "type": "frame", + "data": data, + "metadata": { + "offsetTop": meta.and_then(|m| m.get("offsetTop")).and_then(|v| v.as_f64()).unwrap_or(0.0), + "pageScaleFactor": meta.and_then(|m| m.get("pageScaleFactor")).and_then(|v| v.as_f64()).unwrap_or(1.0), + "deviceWidth": meta.and_then(|m| m.get("deviceWidth")).and_then(|v| v.as_u64()).unwrap_or(1280), + "deviceHeight": meta.and_then(|m| m.get("deviceHeight")).and_then(|v| v.as_u64()).unwrap_or(720), + "scrollOffsetX": meta.and_then(|m| m.get("scrollOffsetX")).and_then(|v| v.as_f64()).unwrap_or(0.0), + "scrollOffsetY": meta.and_then(|m| m.get("scrollOffsetY")).and_then(|v| v.as_f64()).unwrap_or(0.0), + "timestamp": meta.and_then(|m| m.get("timestamp")).and_then(|v| v.as_u64()).unwrap_or(0), + } + }); + let _ = frame_tx.send(msg.to_string()); + } + } + } + Err(broadcast::error::RecvError::Lagged(_)) => continue, + Err(broadcast::error::RecvError::Closed) => break, + } + } + // Also check for notify (client count change or CDP client change) + _ = client_notify.notified() => { + let count = *client_count.lock().await; + let session_id = cdp_session_id.read().await.clone(); + if count == 0 { + // All WS clients gone — stop screencast + let _ = client_arc + .send_command_no_params("Page.stopScreencast", session_id.as_deref()) + .await; + let mut sc = screencasting.lock().await; + *sc = false; + break; + } + // Check if CDP client changed (browser closed/relaunched) + let client_changed = { + let guard = client_slot.read().await; + let same = guard + .as_ref() + .is_some_and(|c| Arc::ptr_eq(c, &client_arc)); + !same + }; + if client_changed { + // CDP client changed — stop our screencast and restart loop + let _ = client_arc + .send_command_no_params("Page.stopScreencast", session_id.as_deref()) + .await; + let mut sc = screencasting.lock().await; + *sc = false; + // Re-notify so we pick up the new client in the outer loop + client_notify.notify_one(); + break; + } + } + } + } + } else { + drop(guard); + // No CDP client yet — wait for next notification + } + } else { + // No WS clients — if screencasting, stop it + let was_screencasting = *screencasting.lock().await; + if was_screencasting { + if let Some(ref client) = *guard { + let session_id = cdp_session_id.read().await.clone(); + let _ = client + .send_command_no_params("Page.stopScreencast", session_id.as_deref()) + .await; + } + let mut sc = screencasting.lock().await; + *sc = false; + } + drop(guard); + } + } +} + +async fn handle_client_message(msg: &str, client: &CdpClient, session_id: Option<&str>) { let parsed: Value = match serde_json::from_str(msg) { Ok(v) => v, Err(_) => return, @@ -276,7 +526,7 @@ async fn handle_client_message(msg: &str, client: &CdpClient, session_id: &str) "deltaY": parsed.get("deltaY").and_then(|v| v.as_f64()).unwrap_or(0.0), "modifiers": parsed.get("modifiers").and_then(|v| v.as_i64()).unwrap_or(0), })), - Some(session_id), + session_id, ) .await; } @@ -291,7 +541,7 @@ async fn handle_client_message(msg: &str, client: &CdpClient, session_id: &str) "text": parsed.get("text"), "modifiers": parsed.get("modifiers").and_then(|v| v.as_i64()).unwrap_or(0), })), - Some(session_id), + session_id, ) .await; } @@ -304,7 +554,7 @@ async fn handle_client_message(msg: &str, client: &CdpClient, session_id: &str) "touchPoints": parsed.get("touchPoints").unwrap_or(&json!([])), "modifiers": parsed.get("modifiers").and_then(|v| v.as_i64()).unwrap_or(0), })), - Some(session_id), + session_id, ) .await; }