refactor(agent): address review comments for window recording support

devolutions-session/Cargo.toml

@@ -44,8 +44,7 @@ win-api-wrappers = { path = "../crates/win-api-wrappers", optional = true }
 
 [dependencies.now-proto-pdu]
 optional = true
-git = "https://github.com/Devolutions/now-proto.git"
-branch = "feat/window-monitoring"
+version = "0.4.2"
 features = ["std"]
 
 [target.'cfg(windows)'.build-dependencies]

devolutions-session/src/dvc/task.rs

@@ -74,7 +74,7 @@ impl Task for DvcIoTask {
 
         // Spawning thread is relatively short operation, so it could be executed synchronously.
         let io_thread = std::thread::spawn(move || {
-            let io_thread_result: Result<(), anyhow::Error> = run_dvc_io(write_rx, read_tx, cloned_shutdown_event);
+            let io_thread_result = run_dvc_io(write_rx, read_tx, cloned_shutdown_event);
 
             if let Err(error) = io_thread_result {
                 error!(%error, "DVC IO thread failed");
@@ -298,6 +298,8 @@ struct MessageProcessor {
     sessions: HashMap<u32, WinApiProcess>,
     /// Shutdown signal sender for window monitoring task.
     window_monitor_shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
+    /// Handle for the window monitor task.
+    window_monitor_handle: Option<tokio::task::JoinHandle<()>>,
 }
 
 impl MessageProcessor {
@@ -312,6 +314,7 @@ impl MessageProcessor {
             capabilities,
             sessions: HashMap::new(),
             window_monitor_shutdown_tx: None,
+            window_monitor_handle: None,
         }
     }
 
@@ -482,7 +485,7 @@ impl MessageProcessor {
                 }
             }
             NowMessage::Session(NowSessionMessage::WindowRecStop(_stop_msg)) => {
-                self.stop_window_recording();
+                self.stop_window_recording().await;
             }
             NowMessage::System(NowSystemMessage::Shutdown(shutdown_msg)) => {
                 let mut current_process_token = win_api_wrappers::process::Process::current_process()
@@ -763,7 +766,7 @@ impl MessageProcessor {
 
     async fn start_window_recording(&mut self, start_msg: NowSessionWindowRecStartMsg) -> anyhow::Result<()> {
         // Stop any existing window recording first.
-        self.stop_window_recording();
+        self.stop_window_recording().await;
 
         info!("Starting window recording");
 
@@ -783,21 +786,30 @@ impl MessageProcessor {
 
         // Spawn window monitor task.
         let event_tx = self.io_notification_tx.clone();
-        tokio::task::spawn(async move {
+        let window_monitor_handle = tokio::task::spawn(async move {
             let config = WindowMonitorConfig::new(event_tx, track_title_changes, shutdown_rx)
                 .with_poll_interval_ms(poll_interval_ms);
 
             run_window_monitor(config).await;
         });
 
+        self.window_monitor_handle = Some(window_monitor_handle);
+
         Ok(())
     }
 
-    fn stop_window_recording(&mut self) {
+    async fn stop_window_recording(&mut self) {
         if let Some(shutdown_tx) = self.window_monitor_shutdown_tx.take() {
             info!("Stopping window recording");
             // Send shutdown signal (ignore errors if receiver was already dropped).
             let _ = shutdown_tx.send(());
+
+            // Wait for the task to finish.
+            if let Some(handle) = self.window_monitor_handle.take() {
+                if let Err(error) = handle.await {
+                    error!(%error, "Window monitor task panicked");
+                }
+            }
         }
     }
 }

devolutions-session/src/dvc/window_monitor.rs

@@ -5,7 +5,8 @@
 //! and timestamp (UTC).
 //!
 //! Uses Windows Event Hooks (SetWinEventHook) to receive EVENT_SYSTEM_FOREGROUND
-//! notifications whenever the foreground window changes, avoiding the need for polling.
+//! notifications whenever the foreground window changes. Additionally supports optional
+//! polling for detecting title changes within the same window.
 //!
 //! The module provides a callback-based interface for integrating with other systems
 //! (e.g., DVC protocol for transmitting window change events).
@@ -194,6 +195,14 @@ fn capture_foreground_window() -> Result<WindowSnapshot> {
     capture_window_snapshot(foreground_window)
 }
 
+/// Gets the current timestamp as seconds since Unix epoch.
+fn get_current_timestamp() -> u64 {
+    SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .map(|d| d.as_secs())
+        .unwrap_or(0)
+}
+
 /// Thread-local context for the event hook callback.
 ///
 /// Windows event hook callbacks must be plain C functions, so we use thread-local
@@ -394,18 +403,15 @@ pub async fn run_window_monitor(config: WindowMonitorConfig) {
     });
 
     // Wait for hook thread to send its thread ID.
-    let hook_thread_id = thread_id_rx.await.expect("Hook thread should send thread ID");
+    let hook_thread_id = thread_id_rx.await.expect("Failed to receive thread ID from hook thread; the thread may have panicked or exited unexpectedly during initialization");
 
     // Track last known window state to detect changes.
     let mut last_snapshot: Option<WindowSnapshot> = None;
 
     // Capture and notify about initial foreground window.
     match capture_foreground_window() {
         Ok(snapshot) => {
-            let timestamp = SystemTime::now()
-                .duration_since(std::time::UNIX_EPOCH)
-                .map(|d| d.as_secs())
-                .unwrap_or(0);
+            let timestamp = get_current_timestamp();
 
             info!(
                 process_id = snapshot.process_id,
@@ -441,10 +447,7 @@ pub async fn run_window_monitor(config: WindowMonitorConfig) {
         Err(error) => {
             debug!(%error, "No initial active window");
 
-            let timestamp = SystemTime::now()
-                .duration_since(std::time::UNIX_EPOCH)
-                .map(|d| d.as_secs())
-                .unwrap_or(0);
+            let timestamp = get_current_timestamp();
 
             // Send "no active window" event.
             let message = NowSessionWindowRecEventMsg::no_active_window(timestamp);
@@ -486,10 +489,7 @@ pub async fn run_window_monitor(config: WindowMonitorConfig) {
 
                 // Check if this is actually a change.
                 if last_snapshot.as_ref() != Some(&snapshot) {
-                    let timestamp = SystemTime::now()
-                        .duration_since(std::time::UNIX_EPOCH)
-                        .map(|d| d.as_secs())
-                        .unwrap_or(0);
+                    let timestamp = get_current_timestamp();
 
                     info!(
                         process_id = snapshot.process_id,
@@ -532,10 +532,7 @@ pub async fn run_window_monitor(config: WindowMonitorConfig) {
                     Ok(snapshot) => {
                         // Check if title or window changed.
                         if last_snapshot.as_ref() != Some(&snapshot) {
-                            let timestamp = SystemTime::now()
-                                .duration_since(std::time::UNIX_EPOCH)
-                                .map(|d| d.as_secs())
-                                .unwrap_or(0);
+                            let timestamp = get_current_timestamp();
 
                             // Determine if only the title changed for the same process.
                             let is_title_change = last_snapshot.as_ref()
@@ -544,7 +541,16 @@ pub async fn run_window_monitor(config: WindowMonitorConfig) {
 
                             // Skip title changes if tracking is disabled.
                             if is_title_change && !config.track_title_changes {
-                                last_snapshot = Some(snapshot);
+                                // Only update process_id and exe_path, keep the previous title
+                                // to avoid missing process/exe_path changes.
+                                let prev_title = last_snapshot
+                                    .as_ref()
+                                    .map_or_else(String::new, |s| s.title.clone());
+                                last_snapshot = Some(WindowSnapshot {
+                                    process_id: snapshot.process_id,
+                                    exe_path: snapshot.exe_path.clone(),
+                                    title: prev_title,
+                                });
                             } else {
                                 let message_result = if is_title_change {
                                     debug!(
@@ -590,10 +596,7 @@ pub async fn run_window_monitor(config: WindowMonitorConfig) {
 
                         // If we previously had an active window, send "no active window" event.
                         if last_snapshot.is_some() {
-                            let timestamp = SystemTime::now()
-                                .duration_since(std::time::UNIX_EPOCH)
-                                .map(|d| d.as_secs())
-                                .unwrap_or(0);
+                            let timestamp = get_current_timestamp();
 
                             let message = NowSessionWindowRecEventMsg::no_active_window(timestamp);
                             if config.event_tx.send(ServerChannelEvent::WindowRecordingEvent { message }).await.is_err() {