fix: restore WebSocket streaming in native daemon (#826)

* fix: restore WebSocket streaming in native daemon

The v0.20.0 Rust rewrite broke WebSocket streaming — connections opened
but received zero messages before closing. Multiple issues contributed:

1. StreamServer was dropped immediately after creation in daemon.rs,
   closing the broadcast channel and killing all WS connections.

2. Screencast frames were only processed during command polling
   (drain_cdp_events) instead of in real-time, unlike the 0.19.0
   TypeScript cdp.on('Page.screencastFrame') callback.

3. Auto-start/stop screencast on WS client connect/disconnect was
   missing from the Rust implementation.

4. Screencast CDP commands used the wrong session ID (daemon session
   name instead of the CDP page session from Target.attachToTarget).

5. Broadcast channel Lagged errors killed WS connections instead of
   being handled gracefully.

The fix adds a background CDP event loop in StreamServer that subscribes
to Chrome events and broadcasts screencast frames in real-time, properly
tracks the CDP page session ID, restores auto-screencast lifecycle, and
keeps the StreamServer alive in DaemonState.

Fixes #820

* fix: use actual CDP session ID for input dispatch in stream WebSocket

Pass the real cdp_session_id (from Target.attachToTarget) through to
handle_ws_client instead of an empty string. Previously, input commands
(mouse, keyboard, touch) were sent with `"sessionId": ""` which Chrome
silently rejects. Now the correct page session ID is read at dispatch
time, and when no session ID is set yet (before browser launch),
the field is omitted entirely via `None` so Chrome uses browser-level
dispatch.

---------

Co-authored-by: ctate <366502+ctate@users.noreply.github.com>

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<String>,
     /// Shared slot for stream server to receive CDP client when browser launches.
     pub stream_client: Option<Arc<RwLock<Option<Arc<CdpClient>>>>>,
+    /// Stream server instance kept alive so the broadcast channel remains open.
+    pub stream_server: Option<Arc<StreamServer>>,
 }
 
 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<Arc<RwLock<Option<Arc<CdpClient>>>>>,
+        stream_server: Option<Arc<StreamServer>>,
     ) -> 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,12 +393,17 @@ impl DaemonState {
                             }
                         }
                         "Page.screencastFrame" => {
+                            // 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" => {
                             let request_id = event
                                 .params
@@ -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<Value, String
     stream::stop_screencast(&mgr.client, session_id).await?;
     state.screencasting = false;
 
+    if let Some(ref server) = state.stream_server {
+        server.broadcast_status(true, false, 0, 0);
+    }
+
     Ok(json!({ "stopped": true }))
 }
 

cli/src/native/daemon.rs

@@ -40,6 +40,7 @@ pub async fn run_daemon(session: &str) {
     }
 
     let mut stream_client: Option<Arc<RwLock<Option<Arc<CdpClient>>>>> = None;
+    let mut stream_server_instance: Option<Arc<StreamServer>> = None;
     if let Ok(port_str) = env::var("AGENT_BROWSER_STREAM_PORT") {
         if let Ok(port) = port_str.parse::<u16>() {
             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::<u64>().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<Arc<RwLock<Option<Arc<CdpClient>>>>>,
+    stream_server: Option<Arc<StreamServer>>,
     idle_timeout_ms: Option<u64>,
 ) -> 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<tokio::sync::Mutex<DaemonState>> = 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<Arc<RwLock<Option<Arc<CdpClient>>>>>,
+    stream_server: Option<Arc<StreamServer>>,
     idle_timeout_ms: Option<u64>,
 ) -> 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<tokio::sync::Mutex<DaemonState>> = 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);

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<String>,
     client_count: Arc<Mutex<usize>>,
+    client_slot: Arc<RwLock<Option<Arc<CdpClient>>>>,
+    /// The active CDP page session ID (from Target.attachToTarget).
+    cdp_session_id: Arc<RwLock<Option<String>>>,
+    client_notify: Arc<Notify>,
+    screencasting: Arc<Mutex<bool>>,
 }
 
 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<String>) {
+        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<RwLock<Option<Arc<CdpClient>>>>,
-        session_id: String,
+        _session_id: String,
     ) -> Result<(Self, Arc<RwLock<Option<Arc<CdpClient>>>>), 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::<String>(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::<String>));
 
         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<String>,
     client_count: Arc<Mutex<usize>>,
     client_slot: Arc<RwLock<Option<Arc<CdpClient>>>>,
-    session_id: String,
+    client_notify: Arc<Notify>,
+    screencasting: Arc<Mutex<bool>>,
+    cdp_session_id: Arc<RwLock<Option<String>>>,
 ) {
     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<String>,
     client_count: Arc<Mutex<usize>>,
     client_slot: Arc<RwLock<Option<Arc<CdpClient>>>>,
-    session_id: String,
+    client_notify: Arc<Notify>,
+    screencasting: Arc<Mutex<bool>>,
+    cdp_session_id: Arc<RwLock<Option<String>>>,
 ) {
     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<String>,
+    client_slot: Arc<RwLock<Option<Arc<CdpClient>>>>,
+    client_notify: Arc<Notify>,
+    screencasting: Arc<Mutex<bool>>,
+    client_count: Arc<Mutex<usize>>,
+    cdp_session_id: Arc<RwLock<Option<String>>>,
+) {
+    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;
         }