/// 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,
}

/// 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>,
    /// Map of connected viewers (id -> info)
    viewers: HashMap<ViewerId, ViewerInfo>,
    /// Instant for heartbeat tracking
    last_heartbeat_instant: Instant,
}


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:
    56→/// Channel for sending input events from viewer to agent
    57→pub type InputSender = tokio::sync::mpsc::Sender<Vec<u8>>;
    58→pub type InputReceiver = tokio::sync::mpsc::Receiver<Vec<u8>>;
    59→
    60→/// Internal session data with channels
    61→struct SessionData {
    62→    info: Session,
    63→    /// Channel for video frames (agent -> viewers)
    64→    frame_tx: FrameSender,
    65→    /// Channel for input events (viewer -> agent)
    66→    input_tx: InputSender,
    67→    input_rx: Option<InputReceiver>,
    68→    /// Map of connected viewers (id -> info)
    69→    viewers: HashMap<ViewerId, ViewerInfo>,
    70→    /// Instant for heartbeat tracking
    71→    last_heartbeat_instant: Instant,
    72→}
    73→
    74→/// Manages all active sessions
    75→#[derive(Clone)]
    76→pub struct SessionManager {

/// 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>,
    /// Map of connected viewers (id -> info)
    viewers: HashMap<ViewerId, ViewerInfo>,
    /// Instant for heartbeat tracking
    last_heartbeat_instant: Instant,
}

//! 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 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;

/// Information about a connected viewer/technician
#[derive(Debug, Clone)]
pub struct ViewerInfo {
    pub id: ViewerId,
    pub name: String,
    pub connected_at: chrono::DateTime<chrono::Utc>,
}

/// 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 viewers: Vec<ViewerInfo>,  // List of connected technicians
    pub is_streaming: bool,
    pub is_online: bool,  // Whether agent is currently connected
    pub is_persistent: bool,  // Persistent agent (no support code) vs support session
    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
    /// If agent was previously connected (offline session exists), reuse that session
    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String, is_persistent: bool) -> (SessionId, FrameSender, InputReceiver) {
        // Check if this agent already has an offline session (reconnecting)
        {
            let agents = self.agents.read().await;
            if let Some(&existing_session_id) = agents.get(&agent_id) {
                let mut sessions = self.sessions.write().await;
                if let Some(session_data) = sessions.get_mut(&existing_session_id) {
                    if !session_data.info.is_online {
                        // Reuse existing session - mark as online and create new channels
                        tracing::info!("Agent {} reconnecting to existing session {}", agent_id, existing_session_id);

                        let (frame_tx, _) = broadcast::channel(16);
                        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);

                        session_data.info.is_online = true;
                        session_data.info.last_heartbeat = chrono::Utc::now();
                        session_data.info.agent_name = agent_name; // Update name in case it changed
                        session_data.frame_tx = frame_tx.clone();
                        session_data.input_tx = input_tx;
                        session_data.last_heartbeat_instant = Instant::now();

                        return (existing_session_id, frame_tx, input_rx);
                    }
                }
            }
        }

        // Create new session
        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,
            is_online: true,
            is_persistent,
            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;
        }
    }

    /// Mark agent as disconnected
    /// For persistent agents: keep session but mark as offline
    /// For support sessions: remove session entirely
    pub async fn mark_agent_disconnected(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            if session_data.info.is_persistent {
                // Persistent agent - keep session but mark as offline
                tracing::info!("Persistent agent {} marked offline (session {} preserved)",
                    session_data.info.agent_id, session_id);
                session_data.info.is_online = false;
                session_data.info.is_streaming = false;
                session_data.info.viewer_count = 0;
                session_data.viewers.clear();
            } else {
                // Support session - remove entirely
                let agent_id = session_data.info.agent_id.clone();
                sessions.remove(&session_id);
                drop(sessions); // Release sessions lock before acquiring agents lock
                let mut agents = self.agents.write().await;
                agents.remove(&agent_id);
                tracing::info!("Support session {} removed", session_id);
            }
        }
    }

    /// Remove a session entirely (for cleanup)
    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) {
            drop(sessions);
            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()
    }
}

+    /// Map of connected viewers (id -> info)

+    viewers: HashMap<ViewerId, ViewerInfo>,


Now I need to update register_agent to initialize viewers properly and include viewers in the Session.


   115→            }
   116→        }
   117→
   118→        // Create new session
   119→        let session_id = Uuid::new_v4();
   120→
   121→        // Create channels
   122→        let (frame_tx, _) = broadcast::channel(16); // Buffer 16 frames
   123→        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64); // Buffer 64 input events
   124→
   125→        let now = chrono::Utc::now();
   126→        let session = Session {
   127→            id: session_id,
   128→            agent_id: agent_id.clone(),
   129→            agent_name,
   130→            started_at: now,
   131→            viewer_count: 0,
   132→            is_streaming: false,
   133→            is_online: true,
   134→            is_persistent,
   135→            last_heartbeat: now,
   136→            os_version: None,
   137→            is_elevated: false,
   138→            uptime_secs: 0,
   139→            display_count: 1,
   140→        };
   141→
   142→        let session_data = SessionData {
   143→            info: session,
   144→            frame_tx: frame_tx.clone(),
   145→            input_tx,
   146→            input_rx: None,
   147→            viewers: HashSet::new(),
   148→            last_heartbeat_instant: Instant::now(),
   149→        };
   150→
   151→        let mut sessions = self.sessions.write().await;
   152→        sessions.insert(session_id, session_data);
   153→
   154→        let mut agents = self.agents.write().await;
   155→        agents.insert(agent_id, session_id);
   156→
   157→        (session_id, frame_tx, input_rx)
   158→    }
   159→
   160→    /// Update agent status from heartbeat or status message
   161→    pub async fn update_agent_status(
   162→        &self,
   163→        session_id: SessionId,
   164→        os_version: Option<String>,

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

            }
        }

        // Create new session
        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,
            is_online: true,
            is_persistent,
            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>,


I see two places that need updating:
1. Line 131 needs `viewers: Vec::new()` added to Session
2. Line 147 needs `viewers: HashMap::new()` instead of `HashSet::new()`Let me also check the reconnection code above (lines 96-103) to make sure viewers is handled there too.


    90→    /// If agent was previously connected (offline session exists), reuse that session
    91→    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String, is_persistent: bool) -> (SessionId, FrameSender, InputReceiver) {
    92→        // Check if this agent already has an offline session (reconnecting)
    93→        {
    94→            let agents = self.agents.read().await;
    95→            if let Some(&existing_session_id) = agents.get(&agent_id) {
    96→                let mut sessions = self.sessions.write().await;
    97→                if let Some(session_data) = sessions.get_mut(&existing_session_id) {
    98→                    if !session_data.info.is_online {
    99→                        // Reuse existing session - mark as online and create new channels
   100→                        tracing::info!("Agent {} reconnecting to existing session {}", agent_id, existing_session_id);
   101→
   102→                        let (frame_tx, _) = broadcast::channel(16);
   103→                        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);
   104→
   105→                        session_data.info.is_online = true;
   106→                        session_data.info.last_heartbeat = chrono::Utc::now();
   107→                        session_data.info.agent_name = agent_name; // Update name in case it changed
   108→                        session_data.frame_tx = frame_tx.clone();
   109→                        session_data.input_tx = input_tx;
   110→                        session_data.last_heartbeat_instant = Instant::now();
   111→
   112→                        return (existing_session_id, frame_tx, input_rx);
   113→                    }
   114→                }
   115→            }
   116→        }
   117→
   118→        // Create new session
   119→        let session_id = Uuid::new_v4();

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

    /// If agent was previously connected (offline session exists), reuse that session
    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String, is_persistent: bool) -> (SessionId, FrameSender, InputReceiver) {
        // Check if this agent already has an offline session (reconnecting)
        {
            let agents = self.agents.read().await;
            if let Some(&existing_session_id) = agents.get(&agent_id) {
                let mut sessions = self.sessions.write().await;
                if let Some(session_data) = sessions.get_mut(&existing_session_id) {
                    if !session_data.info.is_online {
                        // Reuse existing session - mark as online and create new channels
                        tracing::info!("Agent {} reconnecting to existing session {}", agent_id, existing_session_id);

                        let (frame_tx, _) = broadcast::channel(16);
                        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);

                        session_data.info.is_online = true;
                        session_data.info.last_heartbeat = chrono::Utc::now();
                        session_data.info.agent_name = agent_name; // Update name in case it changed
                        session_data.frame_tx = frame_tx.clone();
                        session_data.input_tx = input_tx;
                        session_data.last_heartbeat_instant = Instant::now();

                        return (existing_session_id, frame_tx, input_rx);
                    }
                }
            }
        }

        // Create new session
        let session_id = Uuid::new_v4();


The reconnection code doesn't need to change much since viewers should be empty when an agent reconnects (they were offline). Now let me update the Session initialization and SessionData initialization.


Edit

        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,
            is_online: true,
            is_persistent,
            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 now = chrono::Utc::now();
        let session = Session {
            id: session_id,
            agent_id: agent_id.clone(),
            agent_name,
            started_at: now,
            viewer_count: 0,
            viewers: Vec::new(),
            is_streaming: false,
            is_online: true,
            is_persistent,
            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: HashMap::new(),
            last_heartbeat_instant: Instant::now(),
        };


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:
   121→        // Create channels
   122→        let (frame_tx, _) = broadcast::channel(16); // Buffer 16 frames
   123→        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64); // Buffer 64 input events
   124→
   125→        let now = chrono::Utc::now();
   126→        let session = Session {
   127→            id: session_id,
   128→            agent_id: agent_id.clone(),
   129→            agent_name,
   130→            started_at: now,
   131→            viewer_count: 0,
   132→            viewers: Vec::new(),
   133→            is_streaming: false,
   134→            is_online: true,
   135→            is_persistent,
   136→            last_heartbeat: now,
   137→            os_version: None,
   138→            is_elevated: false,
   139→            uptime_secs: 0,
   140→            display_count: 1,
   141→        };
   142→
   143→        let session_data = SessionData {
   144→            info: session,
   145→            frame_tx: frame_tx.clone(),
   146→            input_tx,
   147→            input_rx: None,
   148→            viewers: HashMap::new(),
   149→            last_heartbeat_instant: Instant::now(),
   150→        };
   151→
   152→        let mut sessions = self.sessions.write().await;
   153→        sessions.insert(session_id, session_data);
   154→

        let now = chrono::Utc::now();
        let session = Session {
            id: session_id,
            agent_id: agent_id.clone(),
            agent_name,
            started_at: now,
            viewer_count: 0,
            viewers: Vec::new(),
            is_streaming: false,
            is_online: true,
            is_persistent,
            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: HashMap::new(),
            last_heartbeat_instant: Instant::now(),
        };

//! 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 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;

/// Information about a connected viewer/technician
#[derive(Debug, Clone)]
pub struct ViewerInfo {
    pub id: ViewerId,
    pub name: String,
    pub connected_at: chrono::DateTime<chrono::Utc>,
}

/// 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 viewers: Vec<ViewerInfo>,  // List of connected technicians
    pub is_streaming: bool,
    pub is_online: bool,  // Whether agent is currently connected
    pub is_persistent: bool,  // Persistent agent (no support code) vs support session
    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>,
    /// Map of connected viewers (id -> info)
    viewers: HashMap<ViewerId, ViewerInfo>,
    /// 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
    /// If agent was previously connected (offline session exists), reuse that session
    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String, is_persistent: bool) -> (SessionId, FrameSender, InputReceiver) {
        // Check if this agent already has an offline session (reconnecting)
        {
            let agents = self.agents.read().await;
            if let Some(&existing_session_id) = agents.get(&agent_id) {
                let mut sessions = self.sessions.write().await;
                if let Some(session_data) = sessions.get_mut(&existing_session_id) {
                    if !session_data.info.is_online {
                        // Reuse existing session - mark as online and create new channels
                        tracing::info!("Agent {} reconnecting to existing session {}", agent_id, existing_session_id);

                        let (frame_tx, _) = broadcast::channel(16);
                        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);

                        session_data.info.is_online = true;
                        session_data.info.last_heartbeat = chrono::Utc::now();
                        session_data.info.agent_name = agent_name; // Update name in case it changed
                        session_data.frame_tx = frame_tx.clone();
                        session_data.input_tx = input_tx;
                        session_data.last_heartbeat_instant = Instant::now();

                        return (existing_session_id, frame_tx, input_rx);
                    }
                }
            }
        }

        // Create new session
        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,
            is_online: true,
            is_persistent,
            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;
        }
    }

    /// Mark agent as disconnected
    /// For persistent agents: keep session but mark as offline
    /// For support sessions: remove session entirely
    pub async fn mark_agent_disconnected(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            if session_data.info.is_persistent {
                // Persistent agent - keep session but mark as offline
                tracing::info!("Persistent agent {} marked offline (session {} preserved)",
                    session_data.info.agent_id, session_id);
                session_data.info.is_online = false;
                session_data.info.is_streaming = false;
                session_data.info.viewer_count = 0;
                session_data.viewers.clear();
            } else {
                // Support session - remove entirely
                let agent_id = session_data.info.agent_id.clone();
                sessions.remove(&session_id);
                drop(sessions); // Release sessions lock before acquiring agents lock
                let mut agents = self.agents.write().await;
                agents.remove(&agent_id);
                tracing::info!("Support session {} removed", session_id);
            }
        }
    }

    /// Remove a session entirely (for cleanup)
    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) {
            drop(sessions);
            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()
    }
}

agent_name,
started_at: now,
viewer_count: 0,
+            viewers: Vec::new(),
is_streaming: false,
is_online: true,
is_persistent,


Now I need to update join_session to accept a viewer name, and update leave_session to use the HashMap. Let me also update the Session's viewers list from the SessionData's viewers map.


Now update join_session to accept viewer name:


    /// 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))
    }

    /// 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, viewer_name: String) -> 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();

        // Add viewer info
        let viewer_info = ViewerInfo {
            id: viewer_id.clone(),
            name: viewer_name.clone(),
            connected_at: chrono::Utc::now(),
        };
        session_data.viewers.insert(viewer_id.clone(), viewer_info);

        // Update session info
        session_data.info.viewer_count = session_data.viewers.len();
        session_data.info.viewers = session_data.viewers.values().cloned().collect();

        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!("Viewer {} ({}) joined session {}, sending StartStream", viewer_name, viewer_id, session_id);
            Self::send_start_stream_internal(session_data, &viewer_id).await;
        } else {
            tracing::info!("Viewer {} ({}) joined session {}", viewer_name, viewer_id, session_id);
        }

        Some((frame_rx, input_tx))
    }


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:
   225→        let sessions = self.sessions.read().await;
   226→        sessions.get(&session_id).map(|s| s.info.clone())
   227→    }
   228→
   229→    /// Join a session as a viewer, returns channels and sends StartStream to agent
   230→    pub async fn join_session(&self, session_id: SessionId, viewer_id: ViewerId, viewer_name: String) -> Option<(FrameReceiver, InputSender)> {
   231→        let mut sessions = self.sessions.write().await;
   232→        let session_data = sessions.get_mut(&session_id)?;
   233→
   234→        let was_empty = session_data.viewers.is_empty();
   235→
   236→        // Add viewer info
   237→        let viewer_info = ViewerInfo {
   238→            id: viewer_id.clone(),
   239→            name: viewer_name.clone(),
   240→            connected_at: chrono::Utc::now(),
   241→        };
   242→        session_data.viewers.insert(viewer_id.clone(), viewer_info);
   243→
   244→        // Update session info
   245→        session_data.info.viewer_count = session_data.viewers.len();
   246→        session_data.info.viewers = session_data.viewers.values().cloned().collect();
   247→
   248→        let frame_rx = session_data.frame_tx.subscribe();
   249→        let input_tx = session_data.input_tx.clone();
   250→
   251→        // If this is the first viewer, send StartStream to agent
   252→        if was_empty {
   253→            tracing::info!("Viewer {} ({}) joined session {}, sending StartStream", viewer_name, viewer_id, session_id);
   254→            Self::send_start_stream_internal(session_data, &viewer_id).await;
   255→        } else {
   256→            tracing::info!("Viewer {} ({}) joined session {}", viewer_name, viewer_id, session_id);
   257→        }
   258→
   259→        Some((frame_rx, input_tx))
   260→    }
   261→
   262→    /// Internal helper to send StartStream message
   263→    async fn send_start_stream_internal(session_data: &SessionData, viewer_id: &str) {
   264→        use crate::proto;

    /// 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, viewer_name: String) -> 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();

        // Add viewer info
        let viewer_info = ViewerInfo {
            id: viewer_id.clone(),
            name: viewer_name.clone(),
            connected_at: chrono::Utc::now(),
        };
        session_data.viewers.insert(viewer_id.clone(), viewer_info);

        // Update session info
        session_data.info.viewer_count = session_data.viewers.len();
        session_data.info.viewers = session_data.viewers.values().cloned().collect();

        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!("Viewer {} ({}) joined session {}, sending StartStream", viewer_name, viewer_id, session_id);
            Self::send_start_stream_internal(session_data, &viewer_id).await;
        } else {
            tracing::info!("Viewer {} ({}) joined session {}", viewer_name, viewer_id, session_id);
        }

        Some((frame_rx, input_tx))
    }

//! 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 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;

/// Information about a connected viewer/technician
#[derive(Debug, Clone)]
pub struct ViewerInfo {
    pub id: ViewerId,
    pub name: String,
    pub connected_at: chrono::DateTime<chrono::Utc>,
}

/// 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 viewers: Vec<ViewerInfo>,  // List of connected technicians
    pub is_streaming: bool,
    pub is_online: bool,  // Whether agent is currently connected
    pub is_persistent: bool,  // Persistent agent (no support code) vs support session
    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>,
    /// Map of connected viewers (id -> info)
    viewers: HashMap<ViewerId, ViewerInfo>,
    /// 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
    /// If agent was previously connected (offline session exists), reuse that session
    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String, is_persistent: bool) -> (SessionId, FrameSender, InputReceiver) {
        // Check if this agent already has an offline session (reconnecting)
        {
            let agents = self.agents.read().await;
            if let Some(&existing_session_id) = agents.get(&agent_id) {
                let mut sessions = self.sessions.write().await;
                if let Some(session_data) = sessions.get_mut(&existing_session_id) {
                    if !session_data.info.is_online {
                        // Reuse existing session - mark as online and create new channels
                        tracing::info!("Agent {} reconnecting to existing session {}", agent_id, existing_session_id);

                        let (frame_tx, _) = broadcast::channel(16);
                        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);

                        session_data.info.is_online = true;
                        session_data.info.last_heartbeat = chrono::Utc::now();
                        session_data.info.agent_name = agent_name; // Update name in case it changed
                        session_data.frame_tx = frame_tx.clone();
                        session_data.input_tx = input_tx;
                        session_data.last_heartbeat_instant = Instant::now();

                        return (existing_session_id, frame_tx, input_rx);
                    }
                }
            }
        }

        // Create new session
        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,
            viewers: Vec::new(),
            is_streaming: false,
            is_online: true,
            is_persistent,
            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: HashMap::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;
        }
    }

    /// Mark agent as disconnected
    /// For persistent agents: keep session but mark as offline
    /// For support sessions: remove session entirely
    pub async fn mark_agent_disconnected(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            if session_data.info.is_persistent {
                // Persistent agent - keep session but mark as offline
                tracing::info!("Persistent agent {} marked offline (session {} preserved)",
                    session_data.info.agent_id, session_id);
                session_data.info.is_online = false;
                session_data.info.is_streaming = false;
                session_data.info.viewer_count = 0;
                session_data.viewers.clear();
            } else {
                // Support session - remove entirely
                let agent_id = session_data.info.agent_id.clone();
                sessions.remove(&session_id);
                drop(sessions); // Release sessions lock before acquiring agents lock
                let mut agents = self.agents.write().await;
                agents.remove(&agent_id);
                tracing::info!("Support session {} removed", session_id);
            }
        }
    }

    /// Remove a session entirely (for cleanup)
    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) {
            drop(sessions);
            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()
    }
}

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

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

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

+            connected_at: chrono::Utc::now(),

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

+        session_data.info.viewers = session_data.viewers.values().cloned().collect();

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

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

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


Now I need to update leave_session to update the viewers list correctly.


   276→            let _ = session_data.input_tx.send(buf).await;
   277→        }
   278→    }
   279→
   280→    /// Leave a session as a viewer, sends StopStream if no viewers left
   281→    pub async fn leave_session(&self, session_id: SessionId, viewer_id: &ViewerId) {
   282→        let mut sessions = self.sessions.write().await;
   283→        if let Some(session_data) = sessions.get_mut(&session_id) {
   284→            session_data.viewers.remove(viewer_id);
   285→            session_data.info.viewer_count = session_data.viewers.len();
   286→
   287→            // If no more viewers, send StopStream to agent
   288→            if session_data.viewers.is_empty() {
   289→                tracing::info!("Last viewer {} left session {}, sending StopStream", viewer_id, session_id);
   290→                Self::send_stop_stream_internal(session_data, viewer_id).await;
   291→            }
   292→        }
   293→    }
   294→
   295→    /// Internal helper to send StopStream message
   296→    async fn send_stop_stream_internal(session_data: &SessionData, viewer_id: &str) {
   297→        use crate::proto;
   298→        use prost::Message;
   299→
   300→        let stop_stream = proto::Message {

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

            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 {


I need to update leave_session to also update the session info's viewers list. The remove already works with HashMap.


    /// 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;
            }
        }
    }

    /// 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) {
            let viewer_name = session_data.viewers.get(viewer_id).map(|v| v.name.clone());
            session_data.viewers.remove(viewer_id);
            session_data.info.viewer_count = session_data.viewers.len();
            session_data.info.viewers = session_data.viewers.values().cloned().collect();

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


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:
   276→            let _ = session_data.input_tx.send(buf).await;
   277→        }
   278→    }
   279→
   280→    /// Leave a session as a viewer, sends StopStream if no viewers left
   281→    pub async fn leave_session(&self, session_id: SessionId, viewer_id: &ViewerId) {
   282→        let mut sessions = self.sessions.write().await;
   283→        if let Some(session_data) = sessions.get_mut(&session_id) {
   284→            let viewer_name = session_data.viewers.get(viewer_id).map(|v| v.name.clone());
   285→            session_data.viewers.remove(viewer_id);
   286→            session_data.info.viewer_count = session_data.viewers.len();
   287→            session_data.info.viewers = session_data.viewers.values().cloned().collect();
   288→
   289→            // If no more viewers, send StopStream to agent
   290→            if session_data.viewers.is_empty() {
   291→                tracing::info!("Last viewer {} ({}) left session {}, sending StopStream",
   292→                    viewer_name.as_deref().unwrap_or("unknown"), viewer_id, session_id);
   293→                Self::send_stop_stream_internal(session_data, viewer_id).await;
   294→            } else {
   295→                tracing::info!("Viewer {} ({}) left session {}",
   296→                    viewer_name.as_deref().unwrap_or("unknown"), viewer_id, session_id);
   297→            }
   298→        }
   299→    }
   300→
   301→    /// Internal helper to send StopStream message
   302→    async fn send_stop_stream_internal(session_data: &SessionData, viewer_id: &str) {
   303→        use crate::proto;

    /// 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) {
            let viewer_name = session_data.viewers.get(viewer_id).map(|v| v.name.clone());
            session_data.viewers.remove(viewer_id);
            session_data.info.viewer_count = session_data.viewers.len();
            session_data.info.viewers = session_data.viewers.values().cloned().collect();

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

//! 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 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;

/// Information about a connected viewer/technician
#[derive(Debug, Clone)]
pub struct ViewerInfo {
    pub id: ViewerId,
    pub name: String,
    pub connected_at: chrono::DateTime<chrono::Utc>,
}

/// 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 viewers: Vec<ViewerInfo>,  // List of connected technicians
    pub is_streaming: bool,
    pub is_online: bool,  // Whether agent is currently connected
    pub is_persistent: bool,  // Persistent agent (no support code) vs support session
    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>,
    /// Map of connected viewers (id -> info)
    viewers: HashMap<ViewerId, ViewerInfo>,
    /// 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
    /// If agent was previously connected (offline session exists), reuse that session
    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String, is_persistent: bool) -> (SessionId, FrameSender, InputReceiver) {
        // Check if this agent already has an offline session (reconnecting)
        {
            let agents = self.agents.read().await;
            if let Some(&existing_session_id) = agents.get(&agent_id) {
                let mut sessions = self.sessions.write().await;
                if let Some(session_data) = sessions.get_mut(&existing_session_id) {
                    if !session_data.info.is_online {
                        // Reuse existing session - mark as online and create new channels
                        tracing::info!("Agent {} reconnecting to existing session {}", agent_id, existing_session_id);

                        let (frame_tx, _) = broadcast::channel(16);
                        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);

                        session_data.info.is_online = true;
                        session_data.info.last_heartbeat = chrono::Utc::now();
                        session_data.info.agent_name = agent_name; // Update name in case it changed
                        session_data.frame_tx = frame_tx.clone();
                        session_data.input_tx = input_tx;
                        session_data.last_heartbeat_instant = Instant::now();

                        return (existing_session_id, frame_tx, input_rx);
                    }
                }
            }
        }

        // Create new session
        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,
            viewers: Vec::new(),
            is_streaming: false,
            is_online: true,
            is_persistent,
            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: HashMap::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, viewer_name: String) -> 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();

        // Add viewer info
        let viewer_info = ViewerInfo {
            id: viewer_id.clone(),
            name: viewer_name.clone(),
            connected_at: chrono::Utc::now(),
        };
        session_data.viewers.insert(viewer_id.clone(), viewer_info);

        // Update session info
        session_data.info.viewer_count = session_data.viewers.len();
        session_data.info.viewers = session_data.viewers.values().cloned().collect();

        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!("Viewer {} ({}) joined session {}, sending StartStream", viewer_name, viewer_id, session_id);
            Self::send_start_stream_internal(session_data, &viewer_id).await;
        } else {
            tracing::info!("Viewer {} ({}) joined session {}", viewer_name, viewer_id, session_id);
        }

        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;
        }
    }

    /// Mark agent as disconnected
    /// For persistent agents: keep session but mark as offline
    /// For support sessions: remove session entirely
    pub async fn mark_agent_disconnected(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            if session_data.info.is_persistent {
                // Persistent agent - keep session but mark as offline
                tracing::info!("Persistent agent {} marked offline (session {} preserved)",
                    session_data.info.agent_id, session_id);
                session_data.info.is_online = false;
                session_data.info.is_streaming = false;
                session_data.info.viewer_count = 0;
                session_data.viewers.clear();
            } else {
                // Support session - remove entirely
                let agent_id = session_data.info.agent_id.clone();
                sessions.remove(&session_id);
                drop(sessions); // Release sessions lock before acquiring agents lock
                let mut agents = self.agents.write().await;
                agents.remove(&agent_id);
                tracing::info!("Support session {} removed", session_id);
            }
        }
    }

    /// Remove a session entirely (for cleanup)
    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) {
            drop(sessions);
            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()
    }
}

+            let viewer_name = session_data.viewers.get(viewer_id).map(|v| v.name.clone());

+            session_data.info.viewers = session_data.viewers.values().cloned().collect();

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

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

+                    viewer_name.as_deref().unwrap_or("unknown"), viewer_id, session_id);

+                tracing::info!("Viewer {} ({}) left session {}",


Now I need to also update mark_agent_disconnected to clear the viewers when an agent disconnects.Let me check that function.


   315→        }
   316→    }
   317→
   318→    /// Mark agent as disconnected
   319→    /// For persistent agents: keep session but mark as offline
   320→    /// For support sessions: remove session entirely
   321→    pub async fn mark_agent_disconnected(&self, session_id: SessionId) {
   322→        let mut sessions = self.sessions.write().await;
   323→        if let Some(session_data) = sessions.get_mut(&session_id) {
   324→            if session_data.info.is_persistent {
   325→                // Persistent agent - keep session but mark as offline
   326→                tracing::info!("Persistent agent {} marked offline (session {} preserved)",
   327→                    session_data.info.agent_id, session_id);
   328→                session_data.info.is_online = false;
   329→                session_data.info.is_streaming = false;
   330→                session_data.info.viewer_count = 0;
   331→                session_data.viewers.clear();
   332→            } else {
   333→                // Support session - remove entirely
   334→                let agent_id = session_data.info.agent_id.clone();
   335→                sessions.remove(&session_id);
   336→                drop(sessions); // Release sessions lock before acquiring agents lock
   337→                let mut agents = self.agents.write().await;
   338→                agents.remove(&agent_id);
   339→                tracing::info!("Support session {} removed", session_id);
   340→            }
   341→        }
   342→    }
   343→
   344→    /// Remove a session entirely (for cleanup)

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

        }
    }

    /// Mark agent as disconnected
    /// For persistent agents: keep session but mark as offline
    /// For support sessions: remove session entirely
    pub async fn mark_agent_disconnected(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            if session_data.info.is_persistent {
                // Persistent agent - keep session but mark as offline
                tracing::info!("Persistent agent {} marked offline (session {} preserved)",
                    session_data.info.agent_id, session_id);
                session_data.info.is_online = false;
                session_data.info.is_streaming = false;
                session_data.info.viewer_count = 0;
                session_data.viewers.clear();
            } else {
                // Support session - remove entirely
                let agent_id = session_data.info.agent_id.clone();
                sessions.remove(&session_id);
                drop(sessions); // Release sessions lock before acquiring agents lock
                let mut agents = self.agents.write().await;
                agents.remove(&agent_id);
                tracing::info!("Support session {} removed", session_id);
            }
        }
    }

    /// Remove a session entirely (for cleanup)


I need to also clear the session's viewers Vec when disconnected.


            if session_data.info.is_persistent {
                // Persistent agent - keep session but mark as offline
                tracing::info!("Persistent agent {} marked offline (session {} preserved)",
                    session_data.info.agent_id, session_id);
                session_data.info.is_online = false;
                session_data.info.is_streaming = false;
                session_data.info.viewer_count = 0;
                session_data.viewers.clear();

            if session_data.info.is_persistent {
                // Persistent agent - keep session but mark as offline
                tracing::info!("Persistent agent {} marked offline (session {} preserved)",
                    session_data.info.agent_id, session_id);
                session_data.info.is_online = false;
                session_data.info.is_streaming = false;
                session_data.info.viewer_count = 0;
                session_data.info.viewers.clear();
                session_data.viewers.clear();