The file C:\Users\MikeSwanson\Claude\guru-connect\server\src\session\mod.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
     1→//! Session management for GuruConnect
     2→//!
     3→//! Manages active remote desktop sessions, tracking which agents
     4→//! are connected and which viewers are watching them.
     5→
     6→use std::collections::{HashMap, HashSet};
     7→use std::sync::Arc;
     8→use std::time::Instant;
     9→use tokio::sync::{broadcast, RwLock};
    10→use uuid::Uuid;
    11→
    12→/// Unique identifier for a session
    13→pub type SessionId = Uuid;
    14→
    15→/// Unique identifier for an agent
    16→pub type AgentId = String;
    17→
    18→/// Unique identifier for a viewer
    19→pub type ViewerId = String;
    20→
    21→/// Heartbeat timeout (90 seconds - 3x the agent's 30 second interval)
    22→const HEARTBEAT_TIMEOUT_SECS: u64 = 90;
    23→
    24→/// Session state
    25→#[derive(Debug, Clone)]
    26→pub struct Session {
    27→    pub id: SessionId,
    28→    pub agent_id: AgentId,
    29→    pub agent_name: String,
    30→    pub started_at: chrono::DateTime<chrono::Utc>,
    31→    pub viewer_count: usize,
    32→    pub is_streaming: bool,
    33→    pub last_heartbeat: chrono::DateTime<chrono::Utc>,
    34→    // Agent status info
    35→    pub os_version: Option<String>,
    36→    pub is_elevated: bool,
    37→    pub uptime_secs: i64,
    38→    pub display_count: i32,
    39→}
    40→
    41→/// Channel for sending frames from agent to viewers
    42→pub type FrameSender = broadcast::Sender<Vec<u8>>;
    43→pub type FrameReceiver = broadcast::Receiver<Vec<u8>>;
    44→
    45→/// Channel for sending input events from viewer to agent
    46→pub type InputSender = tokio::sync::mpsc::Sender<Vec<u8>>;
    47→pub type InputReceiver = tokio::sync::mpsc::Receiver<Vec<u8>>;
    48→
    49→/// Internal session data with channels
    50→struct SessionData {
    51→    info: Session,
    52→    /// Channel for video frames (agent -> viewers)
    53→    frame_tx: FrameSender,
    54→    /// Channel for input events (viewer -> agent)
    55→    input_tx: InputSender,
    56→    input_rx: Option<InputReceiver>,
    57→    /// Set of connected viewer IDs
    58→    viewers: HashSet<ViewerId>,
    59→    /// Instant for heartbeat tracking
    60→    last_heartbeat_instant: Instant,
    61→}
    62→
    63→/// Manages all active sessions
    64→#[derive(Clone)]
    65→pub struct SessionManager {
    66→    sessions: Arc<RwLock<HashMap<SessionId, SessionData>>>,
    67→    agents: Arc<RwLock<HashMap<AgentId, SessionId>>>,
    68→}
    69→
    70→impl SessionManager {
    71→    pub fn new() -> Self {
    72→        Self {
    73→            sessions: Arc::new(RwLock::new(HashMap::new())),
    74→            agents: Arc::new(RwLock::new(HashMap::new())),
    75→        }
    76→    }
    77→
    78→    /// Register a new agent and create a session
    79→    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String) -> (SessionId, FrameSender, InputReceiver) {
    80→        let session_id = Uuid::new_v4();
    81→
    82→        // Create channels
    83→        let (frame_tx, _) = broadcast::channel(16); // Buffer 16 frames
    84→        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64); // Buffer 64 input events
    85→
    86→        let now = chrono::Utc::now();
    87→        let session = Session {
    88→            id: session_id,
    89→            agent_id: agent_id.clone(),
    90→            agent_name,
    91→            started_at: now,
    92→            viewer_count: 0,
    93→            is_streaming: false,
    94→            last_heartbeat: now,
    95→            os_version: None,
    96→            is_elevated: false,
    97→            uptime_secs: 0,
    98→            display_count: 1,
    99→        };
   100→
   101→        let session_data = SessionData {
   102→            info: session,
   103→            frame_tx: frame_tx.clone(),
   104→            input_tx,
   105→            input_rx: None,
   106→            viewers: HashSet::new(),
   107→            last_heartbeat_instant: Instant::now(),
   108→        };
   109→
   110→        let mut sessions = self.sessions.write().await;
   111→        sessions.insert(session_id, session_data);
   112→
   113→        let mut agents = self.agents.write().await;
   114→        agents.insert(agent_id, session_id);
   115→
   116→        (session_id, frame_tx, input_rx)
   117→    }
   118→
   119→    /// Update agent status from heartbeat or status message
   120→    pub async fn update_agent_status(
   121→        &self,
   122→        session_id: SessionId,
   123→        os_version: Option<String>,
   124→        is_elevated: bool,
   125→        uptime_secs: i64,
   126→        display_count: i32,
   127→        is_streaming: bool,
   128→    ) {
   129→        let mut sessions = self.sessions.write().await;
   130→        if let Some(session_data) = sessions.get_mut(&session_id) {
   131→            session_data.info.last_heartbeat = chrono::Utc::now();
   132→            session_data.last_heartbeat_instant = Instant::now();
   133→            session_data.info.is_streaming = is_streaming;
   134→            if let Some(os) = os_version {
   135→                session_data.info.os_version = Some(os);
   136→            }
   137→            session_data.info.is_elevated = is_elevated;
   138→            session_data.info.uptime_secs = uptime_secs;
   139→            session_data.info.display_count = display_count;
   140→        }
   141→    }
   142→
   143→    /// Update heartbeat timestamp
   144→    pub async fn update_heartbeat(&self, session_id: SessionId) {
   145→        let mut sessions = self.sessions.write().await;
   146→        if let Some(session_data) = sessions.get_mut(&session_id) {
   147→            session_data.info.last_heartbeat = chrono::Utc::now();
   148→            session_data.last_heartbeat_instant = Instant::now();
   149→        }
   150→    }
   151→
   152→    /// Check if a session has timed out (no heartbeat for too long)
   153→    pub async fn is_session_timed_out(&self, session_id: SessionId) -> bool {
   154→        let sessions = self.sessions.read().await;
   155→        if let Some(session_data) = sessions.get(&session_id) {
   156→            session_data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS
   157→        } else {
   158→            true // Non-existent sessions are considered timed out
   159→        }
   160→    }
   161→
   162→    /// Get sessions that have timed out
   163→    pub async fn get_timed_out_sessions(&self) -> Vec<SessionId> {
   164→        let sessions = self.sessions.read().await;
   165→        sessions
   166→            .iter()
   167→            .filter(|(_, data)| data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS)
   168→            .map(|(id, _)| *id)
   169→            .collect()
   170→    }
   171→
   172→    /// Get a session by agent ID
   173→    pub async fn get_session_by_agent(&self, agent_id: &str) -> Option<Session> {
   174→        let agents = self.agents.read().await;
   175→        let session_id = agents.get(agent_id)?;
   176→
   177→        let sessions = self.sessions.read().await;
   178→        sessions.get(session_id).map(|s| s.info.clone())
   179→    }
   180→
   181→    /// Get a session by session ID
   182→    pub async fn get_session(&self, session_id: SessionId) -> Option<Session> {
   183→        let sessions = self.sessions.read().await;
   184→        sessions.get(&session_id).map(|s| s.info.clone())
   185→    }
   186→
   187→    /// Join a session as a viewer, returns channels and sends StartStream to agent
   188→    pub async fn join_session(&self, session_id: SessionId, viewer_id: ViewerId) -> Option<(FrameReceiver, InputSender)> {
   189→        let mut sessions = self.sessions.write().await;
   190→        let session_data = sessions.get_mut(&session_id)?;
   191→
   192→        let was_empty = session_data.viewers.is_empty();
   193→        session_data.viewers.insert(viewer_id.clone());
   194→        session_data.info.viewer_count = session_data.viewers.len();
   195→
   196→        let frame_rx = session_data.frame_tx.subscribe();
   197→        let input_tx = session_data.input_tx.clone();
   198→
   199→        // If this is the first viewer, send StartStream to agent
   200→        if was_empty {
   201→            tracing::info!("First viewer {} joined session {}, sending StartStream", viewer_id, session_id);
   202→            self.send_start_stream_internal(session_data, &viewer_id).await;
   203→        }
   204→
   205→        Some((frame_rx, input_tx))
   206→    }
   207→
   208→    /// Internal helper to send StartStream message
   209→    async fn send_start_stream_internal(session_data: &SessionData, viewer_id: &str) {
   210→        use crate::proto;
   211→        use prost::Message;
   212→
   213→        let start_stream = proto::Message {
   214→            payload: Some(proto::message::Payload::StartStream(proto::StartStream {
   215→                viewer_id: viewer_id.to_string(),
   216→                display_id: 0, // Primary display
   217→            })),
   218→        };
   219→
   220→        let mut buf = Vec::new();
   221→        if start_stream.encode(&mut buf).is_ok() {
   222→            let _ = session_data.input_tx.send(buf).await;
   223→        }
   224→    }
   225→
   226→    /// Leave a session as a viewer, sends StopStream if no viewers left
   227→    pub async fn leave_session(&self, session_id: SessionId, viewer_id: &ViewerId) {
   228→        let mut sessions = self.sessions.write().await;
   229→        if let Some(session_data) = sessions.get_mut(&session_id) {
   230→            session_data.viewers.remove(viewer_id);
   231→            session_data.info.viewer_count = session_data.viewers.len();
   232→
   233→            // If no more viewers, send StopStream to agent
   234→            if session_data.viewers.is_empty() {
   235→                tracing::info!("Last viewer {} left session {}, sending StopStream", viewer_id, session_id);
   236→                self.send_stop_stream_internal(session_data, viewer_id).await;
   237→            }
   238→        }
   239→    }
   240→
   241→    /// Internal helper to send StopStream message
   242→    async fn send_stop_stream_internal(session_data: &SessionData, viewer_id: &str) {
   243→        use crate::proto;
   244→        use prost::Message;
   245→
   246→        let stop_stream = proto::Message {
   247→            payload: Some(proto::message::Payload::StopStream(proto::StopStream {
   248→                viewer_id: viewer_id.to_string(),
   249→            })),
   250→        };
   251→
   252→        let mut buf = Vec::new();
   253→        if stop_stream.encode(&mut buf).is_ok() {
   254→            let _ = session_data.input_tx.send(buf).await;
   255→        }
   256→    }
   257→
   258→    /// Remove a session (when agent disconnects)
   259→    pub async fn remove_session(&self, session_id: SessionId) {
   260→        let mut sessions = self.sessions.write().await;
   261→        if let Some(session_data) = sessions.remove(&session_id) {
   262→            let mut agents = self.agents.write().await;
   263→            agents.remove(&session_data.info.agent_id);
   264→        }
   265→    }
   266→
   267→    /// Disconnect a session by sending a disconnect message to the agent
   268→    /// Returns true if the message was sent successfully
   269→    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {
   270→        let sessions = self.sessions.read().await;
   271→        if let Some(session_data) = sessions.get(&session_id) {
   272→            // Create disconnect message
   273→            use crate::proto;
   274→            use prost::Message;
   275→
   276→            let disconnect_msg = proto::Message {
   277→                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
   278→                    reason: reason.to_string(),
   279→                })),
   280→            };
   281→
   282→            let mut buf = Vec::new();
   283→            if disconnect_msg.encode(&mut buf).is_ok() {
   284→                // Send via input channel (will be forwarded to agent's WebSocket)
   285→                if session_data.input_tx.send(buf).await.is_ok() {
   286→                    return true;
   287→                }
   288→            }
   289→        }
   290→        false
   291→    }
   292→
   293→    /// List all active sessions
   294→    pub async fn list_sessions(&self) -> Vec<Session> {
   295→        let sessions = self.sessions.read().await;
   296→        sessions.values().map(|s| s.info.clone()).collect()
   297→    }
   298→}
   299→
   300→impl Default for SessionManager {
   301→    fn default() -> Self {
   302→        Self::new()
   303→    }
   304→}
   305→

//! Session management for GuruConnect
//!
//! Manages active remote desktop sessions, tracking which agents
//! are connected and which viewers are watching them.

use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::{broadcast, RwLock};
use uuid::Uuid;

/// Unique identifier for a session
pub type SessionId = Uuid;

/// Unique identifier for an agent
pub type AgentId = String;

/// Unique identifier for a viewer
pub type ViewerId = String;

/// Heartbeat timeout (90 seconds - 3x the agent's 30 second interval)
const HEARTBEAT_TIMEOUT_SECS: u64 = 90;

/// Session state
#[derive(Debug, Clone)]
pub struct Session {
    pub id: SessionId,
    pub agent_id: AgentId,
    pub agent_name: String,
    pub started_at: chrono::DateTime<chrono::Utc>,
    pub viewer_count: usize,
    pub is_streaming: bool,
    pub last_heartbeat: chrono::DateTime<chrono::Utc>,
    // Agent status info
    pub os_version: Option<String>,
    pub is_elevated: bool,
    pub uptime_secs: i64,
    pub display_count: i32,
}

/// Channel for sending frames from agent to viewers
pub type FrameSender = broadcast::Sender<Vec<u8>>;
pub type FrameReceiver = broadcast::Receiver<Vec<u8>>;

/// Channel for sending input events from viewer to agent
pub type InputSender = tokio::sync::mpsc::Sender<Vec<u8>>;
pub type InputReceiver = tokio::sync::mpsc::Receiver<Vec<u8>>;

/// Internal session data with channels
struct SessionData {
    info: Session,
    /// Channel for video frames (agent -> viewers)
    frame_tx: FrameSender,
    /// Channel for input events (viewer -> agent)
    input_tx: InputSender,
    input_rx: Option<InputReceiver>,
    /// Set of connected viewer IDs
    viewers: HashSet<ViewerId>,
    /// Instant for heartbeat tracking
    last_heartbeat_instant: Instant,
}

/// Manages all active sessions
#[derive(Clone)]
pub struct SessionManager {
    sessions: Arc<RwLock<HashMap<SessionId, SessionData>>>,
    agents: Arc<RwLock<HashMap<AgentId, SessionId>>>,
}

impl SessionManager {
    pub fn new() -> Self {
        Self {
            sessions: Arc::new(RwLock::new(HashMap::new())),
            agents: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Register a new agent and create a session
    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String) -> (SessionId, FrameSender, InputReceiver) {
        let session_id = Uuid::new_v4();

        // Create channels
        let (frame_tx, _) = broadcast::channel(16); // Buffer 16 frames
        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64); // Buffer 64 input events

        let now = chrono::Utc::now();
        let session = Session {
            id: session_id,
            agent_id: agent_id.clone(),
            agent_name,
            started_at: now,
            viewer_count: 0,
            is_streaming: false,
            last_heartbeat: now,
            os_version: None,
            is_elevated: false,
            uptime_secs: 0,
            display_count: 1,
        };

        let session_data = SessionData {
            info: session,
            frame_tx: frame_tx.clone(),
            input_tx,
            input_rx: None,
            viewers: HashSet::new(),
            last_heartbeat_instant: Instant::now(),
        };

        let mut sessions = self.sessions.write().await;
        sessions.insert(session_id, session_data);

        let mut agents = self.agents.write().await;
        agents.insert(agent_id, session_id);

        (session_id, frame_tx, input_rx)
    }

    /// Update agent status from heartbeat or status message
    pub async fn update_agent_status(
        &self,
        session_id: SessionId,
        os_version: Option<String>,
        is_elevated: bool,
        uptime_secs: i64,
        display_count: i32,
        is_streaming: bool,
    ) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            session_data.info.last_heartbeat = chrono::Utc::now();
            session_data.last_heartbeat_instant = Instant::now();
            session_data.info.is_streaming = is_streaming;
            if let Some(os) = os_version {
                session_data.info.os_version = Some(os);
            }
            session_data.info.is_elevated = is_elevated;
            session_data.info.uptime_secs = uptime_secs;
            session_data.info.display_count = display_count;
        }
    }

    /// Update heartbeat timestamp
    pub async fn update_heartbeat(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            session_data.info.last_heartbeat = chrono::Utc::now();
            session_data.last_heartbeat_instant = Instant::now();
        }
    }

    /// Check if a session has timed out (no heartbeat for too long)
    pub async fn is_session_timed_out(&self, session_id: SessionId) -> bool {
        let sessions = self.sessions.read().await;
        if let Some(session_data) = sessions.get(&session_id) {
            session_data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS
        } else {
            true // Non-existent sessions are considered timed out
        }
    }

    /// Get sessions that have timed out
    pub async fn get_timed_out_sessions(&self) -> Vec<SessionId> {
        let sessions = self.sessions.read().await;
        sessions
            .iter()
            .filter(|(_, data)| data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS)
            .map(|(id, _)| *id)
            .collect()
    }

    /// Get a session by agent ID
    pub async fn get_session_by_agent(&self, agent_id: &str) -> Option<Session> {
        let agents = self.agents.read().await;
        let session_id = agents.get(agent_id)?;

        let sessions = self.sessions.read().await;
        sessions.get(session_id).map(|s| s.info.clone())
    }

    /// Get a session by session ID
    pub async fn get_session(&self, session_id: SessionId) -> Option<Session> {
        let sessions = self.sessions.read().await;
        sessions.get(&session_id).map(|s| s.info.clone())
    }

    /// Join a session as a viewer, returns channels and sends StartStream to agent
    pub async fn join_session(&self, session_id: SessionId, viewer_id: ViewerId) -> Option<(FrameReceiver, InputSender)> {
        let mut sessions = self.sessions.write().await;
        let session_data = sessions.get_mut(&session_id)?;

        let was_empty = session_data.viewers.is_empty();
        session_data.viewers.insert(viewer_id.clone());
        session_data.info.viewer_count = session_data.viewers.len();

        let frame_rx = session_data.frame_tx.subscribe();
        let input_tx = session_data.input_tx.clone();

        // If this is the first viewer, send StartStream to agent
        if was_empty {
            tracing::info!("First viewer {} joined session {}, sending StartStream", viewer_id, session_id);
            self.send_start_stream_internal(session_data, &viewer_id).await;
        }

        Some((frame_rx, input_tx))
    }

    /// Internal helper to send StartStream message
    async fn send_start_stream_internal(session_data: &SessionData, viewer_id: &str) {
        use crate::proto;
        use prost::Message;

        let start_stream = proto::Message {
            payload: Some(proto::message::Payload::StartStream(proto::StartStream {
                viewer_id: viewer_id.to_string(),
                display_id: 0, // Primary display
            })),
        };

        let mut buf = Vec::new();
        if start_stream.encode(&mut buf).is_ok() {
            let _ = session_data.input_tx.send(buf).await;
        }
    }

    /// Leave a session as a viewer, sends StopStream if no viewers left
    pub async fn leave_session(&self, session_id: SessionId, viewer_id: &ViewerId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            session_data.viewers.remove(viewer_id);
            session_data.info.viewer_count = session_data.viewers.len();

            // If no more viewers, send StopStream to agent
            if session_data.viewers.is_empty() {
                tracing::info!("Last viewer {} left session {}, sending StopStream", viewer_id, session_id);
                self.send_stop_stream_internal(session_data, viewer_id).await;
            }
        }
    }

    /// Internal helper to send StopStream message
    async fn send_stop_stream_internal(session_data: &SessionData, viewer_id: &str) {
        use crate::proto;
        use prost::Message;

        let stop_stream = proto::Message {
            payload: Some(proto::message::Payload::StopStream(proto::StopStream {
                viewer_id: viewer_id.to_string(),
            })),
        };

        let mut buf = Vec::new();
        if stop_stream.encode(&mut buf).is_ok() {
            let _ = session_data.input_tx.send(buf).await;
        }
    }

    /// Remove a session (when agent disconnects)
    pub async fn remove_session(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.remove(&session_id) {
            let mut agents = self.agents.write().await;
            agents.remove(&session_data.info.agent_id);
        }
    }

    /// Disconnect a session by sending a disconnect message to the agent
    /// Returns true if the message was sent successfully
    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {
        let sessions = self.sessions.read().await;
        if let Some(session_data) = sessions.get(&session_id) {
            // Create disconnect message
            use crate::proto;
            use prost::Message;

            let disconnect_msg = proto::Message {
                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
                    reason: reason.to_string(),
                })),
            };

            let mut buf = Vec::new();
            if disconnect_msg.encode(&mut buf).is_ok() {
                // Send via input channel (will be forwarded to agent's WebSocket)
                if session_data.input_tx.send(buf).await.is_ok() {
                    return true;
                }
            }
        }
        false
    }

    /// List all active sessions
    pub async fn list_sessions(&self) -> Vec<Session> {
        let sessions = self.sessions.read().await;
        sessions.values().map(|s| s.info.clone()).collect()
    }
}

impl Default for SessionManager {
    fn default() -> Self {
        Self::new()
    }
}

-//! Manages active remote desktop sessions, tracking which agents

-//! are connected and which viewers are watching them.

-    pub started_at: chrono::DateTime<chrono::Utc>,

-/// Channel for sending frames from agent to viewers

-pub type FrameSender = broadcast::Sender<Vec<u8>>;

-pub type FrameReceiver = broadcast::Receiver<Vec<u8>>;

-/// Channel for sending input events from viewer to agent

-pub type InputSender = tokio::sync::mpsc::Sender<Vec<u8>>;

-pub type InputReceiver = tokio::sync::mpsc::Receiver<Vec<u8>>;

-/// Internal session data with channels

-    /// Channel for video frames (agent -> viewers)

-    /// Channel for input events (viewer -> agent)

-    sessions: Arc<RwLock<HashMap<SessionId, SessionData>>>,

-    agents: Arc<RwLock<HashMap<AgentId, SessionId>>>,

-            sessions: Arc::new(RwLock::new(HashMap::new())),

-            agents: Arc::new(RwLock::new(HashMap::new())),

-    /// Register a new agent and create a session

-    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String) -> (SessionId, FrameSender, InputReceiver) {

-        let session_id = Uuid::new_v4();

-        let (frame_tx, _) = broadcast::channel(16); // Buffer 16 frames

-        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64); // Buffer 64 input events

-            agent_id: agent_id.clone(),

-            started_at: chrono::Utc::now(),

-        let session_data = SessionData {

-            frame_tx: frame_tx.clone(),

-            input_rx: None, // Will be taken by the agent handler

-        let mut sessions = self.sessions.write().await;

-        sessions.insert(session_id, session_data);

-        let mut agents = self.agents.write().await;

-        agents.insert(agent_id, session_id);

-        (session_id, frame_tx, input_rx)

-    pub async fn get_session_by_agent(&self, agent_id: &str) -> Option<Session> {

-        let agents = self.agents.read().await;

-        let session_id = agents.get(agent_id)?;

-        let sessions = self.sessions.read().await;

-        sessions.get(session_id).map(|s| s.info.clone())

-    pub async fn get_session(&self, session_id: SessionId) -> Option<Session> {

-        sessions.get(&session_id).map(|s| s.info.clone())

-    pub async fn join_session(&self, session_id: SessionId) -> Option<(FrameReceiver, InputSender)> {

-        let session_data = sessions.get_mut(&session_id)?;

-        session_data.info.viewer_count += 1;

-        let frame_rx = session_data.frame_tx.subscribe();

-        let input_tx = session_data.input_tx.clone();

-    pub async fn leave_session(&self, session_id: SessionId) {

-        let mut sessions = self.sessions.write().await;

-        if let Some(session_data) = sessions.get_mut(&session_id) {

-            session_data.info.viewer_count = session_data.info.viewer_count.saturating_sub(1);

-    /// Remove a session (when agent disconnects)

-    pub async fn remove_session(&self, session_id: SessionId) {

-        if let Some(session_data) = sessions.remove(&session_id) {

-            let mut agents = self.agents.write().await;

-            agents.remove(&session_data.info.agent_id);

-    /// Disconnect a session by sending a disconnect message to the agent

-    /// Returns true if the message was sent successfully

-    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {

-        let sessions = self.sessions.read().await;

-        if let Some(session_data) = sessions.get(&session_id) {

-            // Create disconnect message

-            let disconnect_msg = proto::Message {

-                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {

-                    reason: reason.to_string(),

-            if disconnect_msg.encode(&mut buf).is_ok() {

-                // Send via input channel (will be forwarded to agent's WebSocket)

-                if session_data.input_tx.send(buf).await.is_ok() {

-    pub async fn list_sessions(&self) -> Vec<Session> {

-        sessions.values().map(|s| s.info.clone()).collect()

+//! Manages active remote desktop sessions, tracking which agents

+//! are connected and which viewers are watching them.

+/// Heartbeat timeout (90 seconds - 3x the agent's 30 second interval)

+const HEARTBEAT_TIMEOUT_SECS: u64 = 90;

+    pub started_at: chrono::DateTime<chrono::Utc>,

+    pub last_heartbeat: chrono::DateTime<chrono::Utc>,

+/// Channel for sending frames from agent to viewers

+pub type FrameSender = broadcast::Sender<Vec<u8>>;

+pub type FrameReceiver = broadcast::Receiver<Vec<u8>>;

+/// Channel for sending input events from viewer to agent

+pub type InputSender = tokio::sync::mpsc::Sender<Vec<u8>>;

+pub type InputReceiver = tokio::sync::mpsc::Receiver<Vec<u8>>;

+/// Internal session data with channels

+    /// Channel for video frames (agent -> viewers)

+    /// Channel for input events (viewer -> agent)

+    sessions: Arc<RwLock<HashMap<SessionId, SessionData>>>,

+    agents: Arc<RwLock<HashMap<AgentId, SessionId>>>,

+            sessions: Arc::new(RwLock::new(HashMap::new())),

+            agents: Arc::new(RwLock::new(HashMap::new())),

+    /// Register a new agent and create a session

+    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String) -> (SessionId, FrameSender, InputReceiver) {

+        let session_id = Uuid::new_v4();

+        let (frame_tx, _) = broadcast::channel(16); // Buffer 16 frames

+        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64); // Buffer 64 input events

+            agent_id: agent_id.clone(),

+        let session_data = SessionData {

+            frame_tx: frame_tx.clone(),

+            last_heartbeat_instant: Instant::now(),

+        let mut sessions = self.sessions.write().await;

+        sessions.insert(session_id, session_data);

+        let mut agents = self.agents.write().await;

+        agents.insert(agent_id, session_id);

+        (session_id, frame_tx, input_rx)

+    /// Update agent status from heartbeat or status message

+        if let Some(session_data) = sessions.get_mut(&session_id) {

+            session_data.info.last_heartbeat = chrono::Utc::now();

+            session_data.last_heartbeat_instant = Instant::now();

+            session_data.info.is_streaming = is_streaming;

+            if let Some(os) = os_version {

+                session_data.info.os_version = Some(os);

+            session_data.info.is_elevated = is_elevated;

+            session_data.info.uptime_secs = uptime_secs;

+            session_data.info.display_count = display_count;

+    pub async fn update_heartbeat(&self, session_id: SessionId) {

+    /// Check if a session has timed out (no heartbeat for too long)

+    pub async fn is_session_timed_out(&self, session_id: SessionId) -> bool {

+        let sessions = self.sessions.read().await;

+        if let Some(session_data) = sessions.get(&session_id) {

+            session_data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS

+            true // Non-existent sessions are considered timed out

+    /// Get sessions that have timed out

+    pub async fn get_timed_out_sessions(&self) -> Vec<SessionId> {

+            .filter(|(_, data)| data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS)

+    pub async fn get_session_by_agent(&self, agent_id: &str) -> Option<Session> {

+        let agents = self.agents.read().await;

+        let session_id = agents.get(agent_id)?;

+        sessions.get(session_id).map(|s| s.info.clone())

+    pub async fn get_session(&self, session_id: SessionId) -> Option<Session> {

+        sessions.get(&session_id).map(|s| s.info.clone())

+    /// Join a session as a viewer, returns channels and sends StartStream to agent

+    pub async fn join_session(&self, session_id: SessionId, viewer_id: ViewerId) -> Option<(FrameReceiver, InputSender)> {

+        let mut sessions = self.sessions.write().await;

+        let session_data = sessions.get_mut(&session_id)?;

+        let was_empty = session_data.viewers.is_empty();

+        session_data.viewers.insert(viewer_id.clone());

+        session_data.info.viewer_count = session_data.viewers.len();

+        let frame_rx = session_data.frame_tx.subscribe();

+        let input_tx = session_data.input_tx.clone();

+        // If this is the first viewer, send StartStream to agent

+            tracing::info!("First viewer {} joined session {}, sending StartStream", viewer_id, session_id);

+            self.send_start_stream_internal(session_data, &viewer_id).await;

+    /// Internal helper to send StartStream message

+    async fn send_start_stream_internal(session_data: &SessionData, viewer_id: &str) {

+        let start_stream = proto::Message {

+            payload: Some(proto::message::Payload::StartStream(proto::StartStream {

+                viewer_id: viewer_id.to_string(),

+                display_id: 0, // Primary display

+        if start_stream.encode(&mut buf).is_ok() {

+            let _ = session_data.input_tx.send(buf).await;

+    /// Leave a session as a viewer, sends StopStream if no viewers left

+    pub async fn leave_session(&self, session_id: SessionId, viewer_id: &ViewerId) {

+        let mut sessions = self.sessions.write().await;

+            session_data.viewers.remove(viewer_id);

+            // If no more viewers, send StopStream to agent

+            if session_data.viewers.is_empty() {

+                tracing::info!("Last viewer {} left session {}, sending StopStream", viewer_id, session_id);

+                self.send_stop_stream_internal(session_data, viewer_id).await;

+    /// Internal helper to send StopStream message

+    async fn send_stop_stream_internal(session_data: &SessionData, viewer_id: &str) {

+        let stop_stream = proto::Message {

+            payload: Some(proto::message::Payload::StopStream(proto::StopStream {

+                viewer_id: viewer_id.to_string(),

+        if stop_stream.encode(&mut buf).is_ok() {

+            let _ = session_data.input_tx.send(buf).await;

+    /// Remove a session (when agent disconnects)

+    pub async fn remove_session(&self, session_id: SessionId) {

+        let mut sessions = self.sessions.write().await;

+        if let Some(session_data) = sessions.remove(&session_id) {

+            let mut agents = self.agents.write().await;

+            agents.remove(&session_data.info.agent_id);

+    /// Disconnect a session by sending a disconnect message to the agent

+    /// Returns true if the message was sent successfully

+    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {

+        let sessions = self.sessions.read().await;

+        if let Some(session_data) = sessions.get(&session_id) {

+            // Create disconnect message

+            let disconnect_msg = proto::Message {

+                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {

+                    reason: reason.to_string(),

+            if disconnect_msg.encode(&mut buf).is_ok() {

+                // Send via input channel (will be forwarded to agent's WebSocket)

+                if session_data.input_tx.send(buf).await.is_ok() {

+    pub async fn list_sessions(&self) -> Vec<Session> {

+        sessions.values().map(|s| s.info.clone()).collect()

//! Session management for GuruConnect
//!
//! Manages active remote desktop sessions, tracking which agents
//! are connected and which viewers are watching them.

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{broadcast, RwLock};
use uuid::Uuid;

/// Unique identifier for a session
pub type SessionId = Uuid;

/// Unique identifier for an agent
pub type AgentId = String;

/// Session state
#[derive(Debug, Clone)]
pub struct Session {
    pub id: SessionId,
    pub agent_id: AgentId,
    pub agent_name: String,
    pub started_at: chrono::DateTime<chrono::Utc>,
    pub viewer_count: usize,
}

/// Channel for sending frames from agent to viewers
pub type FrameSender = broadcast::Sender<Vec<u8>>;
pub type FrameReceiver = broadcast::Receiver<Vec<u8>>;

/// Channel for sending input events from viewer to agent
pub type InputSender = tokio::sync::mpsc::Sender<Vec<u8>>;
pub type InputReceiver = tokio::sync::mpsc::Receiver<Vec<u8>>;

/// Internal session data with channels
struct SessionData {
    info: Session,
    /// Channel for video frames (agent -> viewers)
    frame_tx: FrameSender,
    /// Channel for input events (viewer -> agent)
    input_tx: InputSender,
    input_rx: Option<InputReceiver>,
}

/// Manages all active sessions
#[derive(Clone)]
pub struct SessionManager {
    sessions: Arc<RwLock<HashMap<SessionId, SessionData>>>,
    agents: Arc<RwLock<HashMap<AgentId, SessionId>>>,
}

impl SessionManager {
    pub fn new() -> Self {
        Self {
            sessions: Arc::new(RwLock::new(HashMap::new())),
            agents: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Register a new agent and create a session
    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String) -> (SessionId, FrameSender, InputReceiver) {
        let session_id = Uuid::new_v4();

        // Create channels
        let (frame_tx, _) = broadcast::channel(16); // Buffer 16 frames
        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64); // Buffer 64 input events

        let session = Session {
            id: session_id,
            agent_id: agent_id.clone(),
            agent_name,
            started_at: chrono::Utc::now(),
            viewer_count: 0,
        };

        let session_data = SessionData {
            info: session,
            frame_tx: frame_tx.clone(),
            input_tx,
            input_rx: None, // Will be taken by the agent handler
        };

        let mut sessions = self.sessions.write().await;
        sessions.insert(session_id, session_data);

        let mut agents = self.agents.write().await;
        agents.insert(agent_id, session_id);

        (session_id, frame_tx, input_rx)
    }

    /// Get a session by agent ID
    pub async fn get_session_by_agent(&self, agent_id: &str) -> Option<Session> {
        let agents = self.agents.read().await;
        let session_id = agents.get(agent_id)?;

        let sessions = self.sessions.read().await;
        sessions.get(session_id).map(|s| s.info.clone())
    }

    /// Get a session by session ID
    pub async fn get_session(&self, session_id: SessionId) -> Option<Session> {
        let sessions = self.sessions.read().await;
        sessions.get(&session_id).map(|s| s.info.clone())
    }

    /// Join a session as a viewer
    pub async fn join_session(&self, session_id: SessionId) -> Option<(FrameReceiver, InputSender)> {
        let mut sessions = self.sessions.write().await;
        let session_data = sessions.get_mut(&session_id)?;

        session_data.info.viewer_count += 1;

        let frame_rx = session_data.frame_tx.subscribe();
        let input_tx = session_data.input_tx.clone();

        Some((frame_rx, input_tx))
    }

    /// Leave a session as a viewer
    pub async fn leave_session(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            session_data.info.viewer_count = session_data.info.viewer_count.saturating_sub(1);
        }
    }

    /// Remove a session (when agent disconnects)
    pub async fn remove_session(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.remove(&session_id) {
            let mut agents = self.agents.write().await;
            agents.remove(&session_data.info.agent_id);
        }
    }

    /// Disconnect a session by sending a disconnect message to the agent
    /// Returns true if the message was sent successfully
    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {
        let sessions = self.sessions.read().await;
        if let Some(session_data) = sessions.get(&session_id) {
            // Create disconnect message
            use crate::proto;
            use prost::Message;

            let disconnect_msg = proto::Message {
                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
                    reason: reason.to_string(),
                })),
            };

            let mut buf = Vec::new();
            if disconnect_msg.encode(&mut buf).is_ok() {
                // Send via input channel (will be forwarded to agent's WebSocket)
                if session_data.input_tx.send(buf).await.is_ok() {
                    return true;
                }
            }
        }
        false
    }

    /// List all active sessions
    pub async fn list_sessions(&self) -> Vec<Session> {
        let sessions = self.sessions.read().await;
        sessions.values().map(|s| s.info.clone()).collect()
    }
}

impl Default for SessionManager {
    fn default() -> Self {
        Self::new()
    }
}


Now I need to update the relay module to:
1. Handle AgentStatus and Heartbeat messages from agent
2. Pass viewer_id to join_session and leave_session
3. Generate viewer IDs for connections

Let me update the relay module:


Now update the relay module to handle the new messages and pass viewer IDs:


Write

//! WebSocket relay handlers
//!
//! Handles WebSocket connections from agents and viewers,
//! relaying video frames and input events between them.

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        Query, State,
    },
    response::IntoResponse,
};
use futures_util::{SinkExt, StreamExt};
use prost::Message as ProstMessage;
use serde::Deserialize;
use tracing::{error, info, warn};
use uuid::Uuid;

use crate::proto;
use crate::session::SessionManager;
use crate::AppState;

#[derive(Debug, Deserialize)]
pub struct AgentParams {
    agent_id: String,
    #[serde(default)]
    agent_name: Option<String>,
    #[serde(default)]
    support_code: Option<String>,
    #[serde(default)]
    hostname: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct ViewerParams {
    session_id: String,
}

/// WebSocket handler for agent connections
pub async fn agent_ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(params): Query<AgentParams>,
) -> impl IntoResponse {
    let agent_id = params.agent_id;
    let agent_name = params.hostname.or(params.agent_name).unwrap_or_else(|| agent_id.clone());
    let support_code = params.support_code;
    let sessions = state.sessions.clone();
    let support_codes = state.support_codes.clone();

    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, agent_id, agent_name, support_code))
}

/// WebSocket handler for viewer connections
pub async fn viewer_ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(params): Query<ViewerParams>,
) -> impl IntoResponse {
    let session_id = params.session_id;
    let sessions = state.sessions.clone();

    ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, session_id))
}

/// Handle an agent WebSocket connection
async fn handle_agent_connection(
    socket: WebSocket,
    sessions: SessionManager,
    support_codes: crate::support_codes::SupportCodeManager,
    agent_id: String,
    agent_name: String,
    support_code: Option<String>,
) {
    info!("Agent connected: {} ({})", agent_name, agent_id);

    let (mut ws_sender, mut ws_receiver) = socket.split();

    // If a support code was provided, check if it's valid
    if let Some(ref code) = support_code {
        // Check if the code is cancelled or invalid
        if support_codes.is_cancelled(code).await {
            warn!("Agent tried to connect with cancelled code: {}", code);
            // Send disconnect message to agent
            let disconnect_msg = proto::Message {
                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
                    reason: "Support session was cancelled by technician".to_string(),
                })),
            };
            let mut buf = Vec::new();
            if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
                let _ = ws_sender.send(Message::Binary(buf.into())).await;
            }
            let _ = ws_sender.close().await;
            return;
        }
    }

    // Register the agent and get channels
    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone()).await;

    info!("Session created: {} (agent in idle mode)", session_id);

    // If a support code was provided, mark it as connected
    if let Some(ref code) = support_code {
        info!("Linking support code {} to session {}", code, session_id);
        support_codes.mark_connected(code, Some(agent_name.clone()), Some(agent_id.clone())).await;
        support_codes.link_session(code, session_id).await;
    }

    // Use Arc<Mutex> for sender so we can use it from multiple places
    let ws_sender = std::sync::Arc::new(tokio::sync::Mutex::new(ws_sender));
    let ws_sender_input = ws_sender.clone();
    let ws_sender_cancel = ws_sender.clone();

    // Task to forward input events from viewers to agent
    let input_forward = tokio::spawn(async move {
        while let Some(input_data) = input_rx.recv().await {
            let mut sender = ws_sender_input.lock().await;
            if sender.send(Message::Binary(input_data.into())).await.is_err() {
                break;
            }
        }
    });

    let sessions_cleanup = sessions.clone();
    let sessions_status = sessions.clone();
    let support_codes_cleanup = support_codes.clone();
    let support_code_cleanup = support_code.clone();
    let support_code_check = support_code.clone();
    let support_codes_check = support_codes.clone();

    // Task to check for cancellation every 2 seconds
    let cancel_check = tokio::spawn(async move {
        let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));
        loop {
            interval.tick().await;
            if let Some(ref code) = support_code_check {
                if support_codes_check.is_cancelled(code).await {
                    info!("Support code {} was cancelled, disconnecting agent", code);
                    // Send disconnect message
                    let disconnect_msg = proto::Message {
                        payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
                            reason: "Support session was cancelled by technician".to_string(),
                        })),
                    };
                    let mut buf = Vec::new();
                    if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
                        let mut sender = ws_sender_cancel.lock().await;
                        let _ = sender.send(Message::Binary(buf.into())).await;
                        let _ = sender.close().await;
                    }
                    break;
                }
            }
        }
    });

    // Main loop: receive messages from agent
    while let Some(msg) = ws_receiver.next().await {
        match msg {
            Ok(Message::Binary(data)) => {
                // Try to decode as protobuf message
                match proto::Message::decode(data.as_ref()) {
                    Ok(proto_msg) => {
                        match &proto_msg.payload {
                            Some(proto::message::Payload::VideoFrame(_)) => {
                                // Broadcast frame to all viewers (only sent when streaming)
                                let _ = frame_tx.send(data.to_vec());
                            }
                            Some(proto::message::Payload::ChatMessage(chat)) => {
                                // Broadcast chat message to all viewers
                                info!("Chat from client: {}", chat.content);
                                let _ = frame_tx.send(data.to_vec());
                            }
                            Some(proto::message::Payload::AgentStatus(status)) => {
                                // Update session with agent status
                                sessions_status.update_agent_status(
                                    session_id,
                                    Some(status.os_version.clone()),
                                    status.is_elevated,
                                    status.uptime_secs,
                                    status.display_count,
                                    status.is_streaming,
                                ).await;
                                info!("Agent status update: {} - streaming={}, uptime={}s",
                                    status.hostname, status.is_streaming, status.uptime_secs);
                            }
                            Some(proto::message::Payload::Heartbeat(_)) => {
                                // Update heartbeat timestamp
                                sessions_status.update_heartbeat(session_id).await;
                            }
                            Some(proto::message::Payload::HeartbeatAck(_)) => {
                                // Agent acknowledged our heartbeat
                                sessions_status.update_heartbeat(session_id).await;
                            }
                            _ => {}
                        }
                    }
                    Err(e) => {
                        warn!("Failed to decode agent message: {}", e);
                    }
                }
            }
            Ok(Message::Close(_)) => {
                info!("Agent disconnected: {}", agent_id);
                break;
            }
            Ok(Message::Ping(data)) => {
                // Pong is handled automatically by axum
                let _ = data;
            }
            Ok(_) => {}
            Err(e) => {
                error!("WebSocket error from agent {}: {}", agent_id, e);
                break;
            }
        }
    }

    // Cleanup
    input_forward.abort();
    cancel_check.abort();
    sessions_cleanup.remove_session(session_id).await;

    // Mark support code as completed if one was used (unless cancelled)
    if let Some(ref code) = support_code_cleanup {
        if !support_codes_cleanup.is_cancelled(code).await {
            support_codes_cleanup.mark_completed(code).await;
            info!("Support code {} marked as completed", code);
        }
    }

    info!("Session {} ended", session_id);
}

/// Handle a viewer WebSocket connection
async fn handle_viewer_connection(
    socket: WebSocket,
    sessions: SessionManager,
    session_id_str: String,
) {
    // Parse session ID
    let session_id = match uuid::Uuid::parse_str(&session_id_str) {
        Ok(id) => id,
        Err(_) => {
            warn!("Invalid session ID: {}", session_id_str);
            return;
        }
    };

    // Generate unique viewer ID
    let viewer_id = Uuid::new_v4().to_string();

    // Join the session (this sends StartStream to agent if first viewer)
    let (mut frame_rx, input_tx) = match sessions.join_session(session_id, viewer_id.clone()).await {
        Some(channels) => channels,
        None => {
            warn!("Session not found: {}", session_id);
            return;
        }
    };

    info!("Viewer {} joined session: {}", viewer_id, session_id);

    let (mut ws_sender, mut ws_receiver) = socket.split();

    // Task to forward frames from agent to this viewer
    let frame_forward = tokio::spawn(async move {
        while let Ok(frame_data) = frame_rx.recv().await {
            if ws_sender.send(Message::Binary(frame_data.into())).await.is_err() {
                break;
            }
        }
    });

    let sessions_cleanup = sessions.clone();
    let viewer_id_cleanup = viewer_id.clone();

    // Main loop: receive input from viewer and forward to agent
    while let Some(msg) = ws_receiver.next().await {
        match msg {
            Ok(Message::Binary(data)) => {
                // Try to decode as protobuf message
                match proto::Message::decode(data.as_ref()) {
                    Ok(proto_msg) => {
                        match &proto_msg.payload {
                            Some(proto::message::Payload::MouseEvent(_)) |
                            Some(proto::message::Payload::KeyEvent(_)) |
                            Some(proto::message::Payload::SpecialKey(_)) => {
                                // Forward input to agent
                                let _ = input_tx.send(data.to_vec()).await;
                            }
                            Some(proto::message::Payload::ChatMessage(chat)) => {
                                // Forward chat message to agent
                                info!("Chat from technician: {}", chat.content);
                                let _ = input_tx.send(data.to_vec()).await;
                            }
                            _ => {}
                        }
                    }
                    Err(e) => {
                        warn!("Failed to decode viewer message: {}", e);
                    }
                }
            }
            Ok(Message::Close(_)) => {
                info!("Viewer {} disconnected from session: {}", viewer_id, session_id);
                break;
            }
            Ok(_) => {}
            Err(e) => {
                error!("WebSocket error from viewer {}: {}", viewer_id, e);
                break;
            }
        }
    }

    // Cleanup (this sends StopStream to agent if last viewer)
    frame_forward.abort();
    sessions_cleanup.leave_session(session_id, &viewer_id_cleanup).await;
    info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);
}