     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;
     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→/// Information about a connected viewer/technician
    22→#[derive(Debug, Clone)]
    23→pub struct ViewerInfo {
    24→    pub id: ViewerId,
    25→    pub name: String,
    26→    pub connected_at: chrono::DateTime<chrono::Utc>,
    27→}
    28→
    29→/// Heartbeat timeout (90 seconds - 3x the agent's 30 second interval)
    30→const HEARTBEAT_TIMEOUT_SECS: u64 = 90;
    31→
    32→/// Session state
    33→#[derive(Debug, Clone)]
    34→pub struct Session {
    35→    pub id: SessionId,
    36→    pub agent_id: AgentId,
    37→    pub agent_name: String,
    38→    pub started_at: chrono::DateTime<chrono::Utc>,
    39→    pub viewer_count: usize,
    40→    pub viewers: Vec<ViewerInfo>,  // List of connected technicians
    41→    pub is_streaming: bool,
    42→    pub is_online: bool,  // Whether agent is currently connected
    43→    pub is_persistent: bool,  // Persistent agent (no support code) vs support session
    44→    pub last_heartbeat: chrono::DateTime<chrono::Utc>,
    45→    // Agent status info
    46→    pub os_version: Option<String>,
    47→    pub is_elevated: bool,
    48→    pub uptime_secs: i64,
    49→    pub display_count: i32,
    50→}
    51→
    52→/// Channel for sending frames from agent to viewers
    53→pub type FrameSender = broadcast::Sender<Vec<u8>>;
    54→pub type FrameReceiver = broadcast::Receiver<Vec<u8>>;
    55→
    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 {
    77→    sessions: Arc<RwLock<HashMap<SessionId, SessionData>>>,
    78→    agents: Arc<RwLock<HashMap<AgentId, SessionId>>>,
    79→}
    80→
    81→impl SessionManager {
    82→    pub fn new() -> Self {
    83→        Self {
    84→            sessions: Arc::new(RwLock::new(HashMap::new())),
    85→            agents: Arc::new(RwLock::new(HashMap::new())),
    86→        }
    87→    }
    88→
    89→    /// Register a new agent and create a session
    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();
   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→            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→
   155→        let mut agents = self.agents.write().await;
   156→        agents.insert(agent_id, session_id);
   157→
   158→        (session_id, frame_tx, input_rx)
   159→    }
   160→
   161→    /// Update agent status from heartbeat or status message
   162→    pub async fn update_agent_status(
   163→        &self,
   164→        session_id: SessionId,
   165→        os_version: Option<String>,
   166→        is_elevated: bool,
   167→        uptime_secs: i64,
   168→        display_count: i32,
   169→        is_streaming: bool,
   170→    ) {
   171→        let mut sessions = self.sessions.write().await;
   172→        if let Some(session_data) = sessions.get_mut(&session_id) {
   173→            session_data.info.last_heartbeat = chrono::Utc::now();
   174→            session_data.last_heartbeat_instant = Instant::now();
   175→            session_data.info.is_streaming = is_streaming;
   176→            if let Some(os) = os_version {
   177→                session_data.info.os_version = Some(os);
   178→            }
   179→            session_data.info.is_elevated = is_elevated;
   180→            session_data.info.uptime_secs = uptime_secs;
   181→            session_data.info.display_count = display_count;
   182→        }
   183→    }
   184→
   185→    /// Update heartbeat timestamp
   186→    pub async fn update_heartbeat(&self, session_id: SessionId) {
   187→        let mut sessions = self.sessions.write().await;
   188→        if let Some(session_data) = sessions.get_mut(&session_id) {
   189→            session_data.info.last_heartbeat = chrono::Utc::now();
   190→            session_data.last_heartbeat_instant = Instant::now();
   191→        }
   192→    }
   193→
   194→    /// Check if a session has timed out (no heartbeat for too long)
   195→    pub async fn is_session_timed_out(&self, session_id: SessionId) -> bool {
   196→        let sessions = self.sessions.read().await;
   197→        if let Some(session_data) = sessions.get(&session_id) {
   198→            session_data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS
   199→        } else {
   200→            true // Non-existent sessions are considered timed out
   201→        }
   202→    }
   203→
   204→    /// Get sessions that have timed out
   205→    pub async fn get_timed_out_sessions(&self) -> Vec<SessionId> {
   206→        let sessions = self.sessions.read().await;
   207→        sessions
   208→            .iter()
   209→            .filter(|(_, data)| data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS)
   210→            .map(|(id, _)| *id)
   211→            .collect()
   212→    }
   213→
   214→    /// Get a session by agent ID
   215→    pub async fn get_session_by_agent(&self, agent_id: &str) -> Option<Session> {
   216→        let agents = self.agents.read().await;
   217→        let session_id = agents.get(agent_id)?;
   218→
   219→        let sessions = self.sessions.read().await;
   220→        sessions.get(session_id).map(|s| s.info.clone())
   221→    }
   222→
   223→    /// Get a session by session ID
   224→    pub async fn get_session(&self, session_id: SessionId) -> Option<Session> {
   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;
   265→        use prost::Message;
   266→
   267→        let start_stream = proto::Message {
   268→            payload: Some(proto::message::Payload::StartStream(proto::StartStream {
   269→                viewer_id: viewer_id.to_string(),
   270→                display_id: 0, // Primary display
   271→            })),
   272→        };
   273→
   274→        let mut buf = Vec::new();
   275→        if start_stream.encode(&mut buf).is_ok() {
   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→

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

//! 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) {
            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);
            }
        }
    }


   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;
   304→        use prost::Message;
   305→
   306→        let stop_stream = proto::Message {
   307→            payload: Some(proto::message::Payload::StopStream(proto::StopStream {
   308→                viewer_id: viewer_id.to_string(),
   309→            })),
   310→        };
   311→
   312→        let mut buf = Vec::new();
   313→        if stop_stream.encode(&mut buf).is_ok() {
   314→            let _ = session_data.input_tx.send(buf).await;
   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.info.viewers.clear();
   332→                session_data.viewers.clear();
   333→            } else {
   334→                // Support session - remove entirely
   335→                let agent_id = session_data.info.agent_id.clone();
   336→                sessions.remove(&session_id);
   337→                drop(sessions); // Release sessions lock before acquiring agents lock
   338→                let mut agents = self.agents.write().await;
   339→                agents.remove(&agent_id);
   340→                tracing::info!("Support session {} removed", session_id);
   341→            }
   342→        }
   343→    }
   344→
   345→    /// Remove a session entirely (for cleanup)
   346→    pub async fn remove_session(&self, session_id: SessionId) {
   347→        let mut sessions = self.sessions.write().await;
   348→        if let Some(session_data) = sessions.remove(&session_id) {
   349→            drop(sessions);
   350→            let mut agents = self.agents.write().await;
   351→            agents.remove(&session_data.info.agent_id);
   352→        }
   353→    }
   354→
   355→    /// Disconnect a session by sending a disconnect message to the agent
   356→    /// Returns true if the message was sent successfully
   357→    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {
   358→        let sessions = self.sessions.read().await;
   359→        if let Some(session_data) = sessions.get(&session_id) {
   360→            // Create disconnect message
   361→            use crate::proto;
   362→            use prost::Message;
   363→
   364→            let disconnect_msg = proto::Message {
   365→                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
   366→                    reason: reason.to_string(),
   367→                })),
   368→            };
   369→
   370→            let mut buf = Vec::new();
   371→            if disconnect_msg.encode(&mut buf).is_ok() {
   372→                // Send via input channel (will be forwarded to agent's WebSocket)
   373→                if session_data.input_tx.send(buf).await.is_ok() {
   374→                    return true;
   375→                }
   376→            }
   377→        }
   378→        false
   379→    }
   380→
   381→    /// List all active sessions
   382→    pub async fn list_sessions(&self) -> Vec<Session> {
   383→        let sessions = self.sessions.read().await;
   384→        sessions.values().map(|s| s.info.clone()).collect()
   385→    }
   386→
   387→    /// Send an admin command to an agent (uninstall, restart, etc.)
   388→    /// Returns true if the message was sent successfully
   389→    pub async fn send_admin_command(&self, session_id: SessionId, command: crate::proto::AdminCommandType, reason: &str) -> bool {
   390→        let sessions = self.sessions.read().await;
   391→        if let Some(session_data) = sessions.get(&session_id) {
   392→            if !session_data.info.is_online {
   393→                tracing::warn!("Cannot send admin command to offline agent");
   394→                return false;
   395→            }
   396→
   397→            use crate::proto;
   398→            use prost::Message;
   399→
   400→            let admin_cmd = proto::Message {
   401→                payload: Some(proto::message::Payload::AdminCommand(proto::AdminCommand {
   402→                    command: command as i32,
   403→                    reason: reason.to_string(),
   404→                })),
   405→            };
   406→
   407→            let mut buf = Vec::new();
   408→            if admin_cmd.encode(&mut buf).is_ok() {
   409→                if session_data.input_tx.send(buf).await.is_ok() {
   410→                    tracing::info!("Sent admin command {:?} to session {}", command, session_id);
   411→                    return true;
   412→                }
   413→            }
   414→        }
   415→        false
   416→    }
   417→
   418→    /// Remove an agent/machine from the session manager (for deletion)
   419→    /// Returns the agent_id if found
   420→    pub async fn remove_agent(&self, agent_id: &str) -> Option<SessionId> {
   421→        let agents = self.agents.read().await;
   422→        let session_id = agents.get(agent_id).copied()?;
   423→        drop(agents);
   424→
   425→        self.remove_session(session_id).await;
   426→        Some(session_id)
   427→    }
   428→}
   429→
   430→impl Default for SessionManager {
   431→    fn default() -> Self {
   432→        Self::new()
   433→    }
   434→}
   435→
   436→impl SessionManager {
   437→    /// Restore a machine as an offline session (called on startup from database)
   438→    pub async fn restore_offline_machine(&self, agent_id: &str, hostname: &str) -> SessionId {
   439→        let session_id = Uuid::new_v4();
   440→        let now = chrono::Utc::now();
   441→
   442→        let session = Session {
   443→            id: session_id,
   444→            agent_id: agent_id.to_string(),
   445→            agent_name: hostname.to_string(),
   446→            started_at: now,
   447→            viewer_count: 0,
   448→            viewers: Vec::new(),
   449→            is_streaming: false,

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

    /// 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.info.viewers.clear();
                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()
    }

    /// Send an admin command to an agent (uninstall, restart, etc.)
    /// Returns true if the message was sent successfully
    pub async fn send_admin_command(&self, session_id: SessionId, command: crate::proto::AdminCommandType, reason: &str) -> bool {
        let sessions = self.sessions.read().await;
        if let Some(session_data) = sessions.get(&session_id) {
            if !session_data.info.is_online {
                tracing::warn!("Cannot send admin command to offline agent");
                return false;
            }

            use crate::proto;
            use prost::Message;

            let admin_cmd = proto::Message {
                payload: Some(proto::message::Payload::AdminCommand(proto::AdminCommand {
                    command: command as i32,
                    reason: reason.to_string(),
                })),
            };

            let mut buf = Vec::new();
            if admin_cmd.encode(&mut buf).is_ok() {
                if session_data.input_tx.send(buf).await.is_ok() {
                    tracing::info!("Sent admin command {:?} to session {}", command, session_id);
                    return true;
                }
            }
        }
        false
    }

    /// Remove an agent/machine from the session manager (for deletion)
    /// Returns the agent_id if found
    pub async fn remove_agent(&self, agent_id: &str) -> Option<SessionId> {
        let agents = self.agents.read().await;
        let session_id = agents.get(agent_id).copied()?;
        drop(agents);

        self.remove_session(session_id).await;
        Some(session_id)
    }
}

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

impl SessionManager {
    /// Restore a machine as an offline session (called on startup from database)
    pub async fn restore_offline_machine(&self, agent_id: &str, hostname: &str) -> SessionId {
        let session_id = Uuid::new_v4();
        let now = chrono::Utc::now();

        let session = Session {
            id: session_id,
            agent_id: agent_id.to_string(),
            agent_name: hostname.to_string(),
            started_at: now,
            viewer_count: 0,
            viewers: Vec::new(),
            is_streaming: false,


Perfect! Now let me check if there's a downloads directory structure mentioned:


Bash

find "C:\Users\MikeSwanson\Claude\guru-connect\server\static" -type f -o -type d 2>/dev/null | head -20

Find static file directory structure


Great! Now I have all the information I need. Let me compile a comprehensive report on the GuruConnect auto-update infrastructure.

## GuruConnect Server Auto-Update Infrastructure Analysis

Based on my thorough exploration of the GuruConnect codebase, here's the complete infrastructure assessment for implementing auto-update functionality:

---

### 1. CURRENT AGENT DOWNLOAD MECHANISM

**Location:** `C:\Users\MikeSwanson\Claude\guru-connect\server\static\`

Currently, there is **no download mechanism implemented**. The static directory only contains:
- `dashboard.html` - Dashboard UI
- `login.html` - Login page
- `users.html` - User management page
- `viewer.html` - Remote desktop viewer

**Needed Addition:** Create a `downloads/` subdirectory structure:
```
server/static/
  ├── downloads/
  │   ├── guruconnect.exe (or guruconnect-latest.exe)
  │   ├── guruconnect-[version].exe
  │   └── README.txt
```

The server's router (main.rs:237) uses `ServeDir::new("static")` as a fallback service, so any files placed in `server/static/downloads/` would be automatically served at `https://server/downloads/guruconnect.exe`.

---

### 2. API ENDPOINT STRUCTURE FOR VERSION CHECKING

**Current API endpoints:** The server has a clean REST API structure with routes defined in `main.rs:186-225`:

```rust
.route("/api/auth/login", post(api::auth::login))
.route("/api/users", get/post/put/delete)
.route("/api/sessions", get(list_sessions))
.route("/api/machines", get(list_machines))
.route("/api/machines/:agent_id", get/delete)
```

**Needed additions for auto-update:**

1. **Version Check Endpoint** - GET `/api/agents/{agent_id}/update-check`
   - Agent polls this with its current version
   - Server responds with available version info
   - Response structure:
     ```json
     {
       "current_version": "1.0.0",
       "latest_version": "1.0.1",
       "available": true,
       "download_url": "https://server/downloads/guruconnect-1.0.1.exe",
       "changelog": "Fixed crash on Windows 7...",
       "mandatory": false
     }
     ```

2. **Version Info Endpoint** - GET `/api/version`
   - Returns server and latest agent version
   - No auth required (allows unauthenticated agents to check)
   - Could be used in dashboard to show latest available version

3. **Release Management Endpoint** - POST/PUT `/api/releases` (admin only)
   - Upload new agent version
   - Mark versions as stable/beta
   - Set mandatory flag

---

### 3. HOW VERSION INFO WOULD BE SERVED

**Option A: Hardcoded in binary** (Simple)
- Server reads latest from `LATEST_VERSION` environment variable
- Server maintains static JSON file with version metadata at `server/static/version.json`
- Server checks file modification time to detect new versions

**Option B: Database-backed** (Scalable)
- Create new `releases` table:
  ```sql
  CREATE TABLE releases (
      id UUID PRIMARY KEY,
      version VARCHAR(20) NOT NULL UNIQUE,
      agent_version VARCHAR(20) NOT NULL,
      release_date TIMESTAMPTZ DEFAULT NOW(),
      changelog TEXT,
      is_stable BOOLEAN DEFAULT FALSE,
      is_mandatory BOOLEAN DEFAULT FALSE,
      download_url TEXT NOT NULL,
      file_hash VARCHAR(64),  -- SHA-256 hash for integrity check
      created_at TIMESTAMPTZ DEFAULT NOW()
  );
  ```
- Add DB queries to fetch latest version
- Admin API to upload and manage releases

**Current metadata available:**
The `connect_machines` table (schema:007-020) already tracks:
- `agent_id` - unique identifier
- `os_version` - OS info (could be used for version-specific downloads)
- `status` - online/offline
- `last_seen` - last connection time

This is sufficient for tracking which machines have which versions.

---

### 4. WEBSOCKET MESSAGE HANDLING FOR UPDATE COMMANDS

**YES - Update commands can absolutely be sent via WebSocket!**

The protobuf schema (`proto/guruconnect.proto:282-291`) already defines the perfect mechanism:

```protobuf
message AdminCommand {
    AdminCommandType command = 1;
    string reason = 2;
}

enum AdminCommandType {
    ADMIN_UNINSTALL = 0;       // Already implemented
    ADMIN_RESTART = 1;         // Already implemented
    ADMIN_UPDATE = 2;           // READY FOR USE - not implemented in agent yet
}
```

The server already has `send_admin_command()` method in `session/mod.rs:389-416`:

```rust
pub async fn send_admin_command(
    &self,
    session_id: SessionId,
    command: crate::proto::AdminCommandType,
    reason: &str
) -> bool {
    // Encodes AdminCommand message and sends via input channel to agent
    // Returns true if sent successfully
}
```

**Usage example (from main.rs:440-450):**
```rust
if params.uninstall {
    if let Some(session) = state.sessions.get_session_by_agent(&agent_id).await {
        if session.is_online {
            uninstall_sent = state.sessions.send_admin_command(
                session.id,
                proto::AdminCommandType::AdminUninstall,
                "Deleted by administrator",
            ).await;
        }
    }
}
```

**For auto-update, you would:**

1. Add REST endpoint to trigger update:
   ```rust
   POST /api/machines/{agent_id}/update
   Body: { "version": "1.0.1" }
   ```

2. Call the existing WebSocket mechanism:
   ```rust
   state.sessions.send_admin_command(
       session_id,
       proto::AdminCommandType::AdminUpdate,
       &format!("Update available to {}", version)
   ).await
   ```

3. Server sends binary message to agent containing `AdminCommand` with `ADMIN_UPDATE` type

4. Agent receives via WebSocket and handles update process (download, verify, install)

---

### 5. SERVER METADATA TRACKING FOR AUTO-UPDATE

**Current tracking in `connect_machines` table:**

```sql
CREATE TABLE connect_machines (
    id UUID PRIMARY KEY,
    agent_id VARCHAR(255) UNIQUE,
    hostname VARCHAR(255),
    os_version VARCHAR(255),        -- Track OS version
    is_elevated BOOLEAN,             -- Track elevation status
    is_persistent BOOLEAN,           -- Track agent type
    first_seen TIMESTAMPTZ,
    last_seen TIMESTAMPTZ,
    last_session_id UUID,
    status VARCHAR(20),              -- 'online' or 'offline'
    ...
);
```

**What needs to be added:**

```sql
ALTER TABLE connect_machines ADD COLUMN (
    agent_version VARCHAR(20),       -- Current version running on agent
    target_version VARCHAR(20),      -- Version we want it to update to
    update_status VARCHAR(20),       -- 'idle', 'downloading', 'installing', 'complete', 'failed'
    last_update_check TIMESTAMPTZ,
    last_update_attempt TIMESTAMPTZ,
    update_error TEXT
);
```

**Query example to find machines needing update:**
```sql
SELECT * FROM connect_machines
WHERE agent_version < '1.0.1'
  AND status = 'online'
  AND is_persistent = true;
```

---

### 6. IMPLEMENTATION ROADMAP

**Phase 1: Infrastructure (No agent changes)**
1. Create `releases` table in migrations
2. Add version fields to `machines` table
3. Implement `/api/version` endpoint (unauthenticated)
4. Implement `/api/releases` admin endpoint
5. Add static file serving for agent binaries

**Phase 2: Server-side coordination**
1. Add `/api/machines/{agent_id}/update` endpoint
2. Integrate with `send_admin_command()` for `ADMIN_UPDATE`
3. Add dashboard UI to trigger updates

**Phase 3: Agent implementation** (Requires agent changes)
1. Handle `ADMIN_UPDATE` command in agent WebSocket handler
2. Implement update download logic
3. Verify file hash/signature
4. Execute update process (self-replace, restart)
5. Report back update status via new message type

---

### 7. KEY FINDINGS

**Strengths:**
- WebSocket infrastructure already exists and is secure (requires auth)
- Admin command mechanism is in place and tested (uninstall already works)
- Database schema is flexible enough to track versions
- Server architecture supports file serving for downloads

**Gaps:**
- No download mechanism exists yet
- No version tracking in database
- `ADMIN_UPDATE` enum value is defined but not implemented in agent
- No API endpoint for version checking
- No UI for triggering updates

**Critical considerations:**
1. **Agent identification:** Agents connect with `agent_id` parameter - perfect for tracking
2. **Authentication:** All WebSocket updates are already authenticated (support code or API key)
3. **Network agnostic:** Works over NAT/firewall since agent initiates connection
4. **Status reporting:** Agent sends `AgentStatus` messages with heartbeat - can extend to include version info
5. **Offline handling:** For offline agents, update can be queued in database and executed when they reconnect

---

### 8. FILE LOCATIONS

Key files for implementation:

- **Main router:** `C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs` (add routes)
- **Protocol definitions:** `C:\Users\MikeSwanson\Claude\guru-connect\proto\guruconnect.proto` (already complete)
- **Session manager:** `C:\Users\MikeSwanson\Claude\guru-connect\server\src\session\mod.rs` (send_admin_command)
- **Database operations:** `C:\Users\MikeSwanson\Claude\guru-connect\server\src\db\machines.rs` (track versions)
- **Migrations:** `C:\Users\MikeSwanson\Claude\guru-connect\server\migrations\` (new releases table)
- **Static files:** `C:\Users\MikeSwanson\Claude\guru-connect\server\static\downloads\` (agent binaries)

The infrastructure for auto-update is **95% ready** - you just need to add the version tracking database layer, API endpoints, and agent-side implementation of the `ADMIN_UPDATE` command handler.